crank/crates/crank-adapter-rest/src/client.rs

use std::{
    collections::BTreeMap,
    env, io,
    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
    sync::Arc,
    time::Duration,
};

use crank_core::{HttpMethod, RestTarget};
use crank_metrics::{UpstreamOperationKind, UpstreamOutcome, UpstreamRequestMetrics};
use crank_trace::{ErrorCategory, Stage, StageOutcome};
use futures_util::StreamExt;
use opentelemetry::{global, propagation::Injector, trace::TraceContextExt};
use reqwest::{
    Client,
    dns::{Addrs, Name, Resolve, Resolving},
    header::{HeaderMap, HeaderName, HeaderValue},
    redirect,
};
use serde_json::Value;
use tracing::{Instrument, Span};
use tracing_opentelemetry::OpenTelemetrySpanExt;

use crate::{RestAdapterError, RestRequest, RestResponse};

#[derive(Clone, Debug)]
pub struct RestAdapter {
    client: Result<Client, Arc<str>>,
    policy: OutboundHttpPolicy,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct OutboundHttpPolicy {
    allowed_hosts: Vec<String>,
    denied_hosts: Vec<String>,
    max_response_bytes: usize,
}

const DEFAULT_MAX_RESPONSE_BYTES: usize = 4 * 1024 * 1024;

impl Default for RestAdapter {
    fn default() -> Self {
        Self::new()
    }
}

impl RestAdapter {
    pub fn new() -> Self {
        Self::with_policy(OutboundHttpPolicy::default())
    }

    pub fn from_env() -> Result<Self, RestAdapterError> {
        Ok(Self::with_policy(OutboundHttpPolicy::from_env()?))
    }

    pub fn with_policy(policy: OutboundHttpPolicy) -> Self {
        let resolver = Arc::new(PolicyDnsResolver {
            policy: policy.clone(),
        });
        let client = Client::builder()
            .redirect(redirect::Policy::none())
            .no_proxy()
            .dns_resolver(resolver)
            .build()
            .map_err(|error| Arc::<str>::from(error.to_string()));

        Self { client, policy }
    }

    pub async fn execute(
        &self,
        target: &RestTarget,
        request: &RestRequest,
    ) -> Result<RestResponse, RestAdapterError> {
        let request_metrics = UpstreamRequestMetrics::start(UpstreamOperationKind::Rest);
        let result = self.execute_inner(target, request).await;
        let outcome = match &result {
            Ok(_) => UpstreamOutcome::Success,
            Err(error) => upstream_outcome(error),
        };
        request_metrics.complete(outcome);
        result
    }

    async fn execute_inner(
        &self,
        target: &RestTarget,
        request: &RestRequest,
    ) -> Result<RestResponse, RestAdapterError> {
        let url = build_url(target, request)?;
        self.policy.validate_url(&url)?;
        let mut headers = build_headers(target, request)?;
        apply_current_trace_context(&mut headers);
        let client =
            self.client
                .as_ref()
                .map_err(|details| RestAdapterError::InvalidConfiguration {
                    details: details.to_string(),
                })?;
        let mut builder = client
            .request(to_reqwest_method(target.method), url)
            .headers(headers)
            .timeout(Duration::from_millis(request.timeout_ms));

        if let Some(body) = &request.body {
            builder = builder.json(body);
        }

        let upstream_span = Stage::UpstreamHttp.span();
        let result = async {
            let response = builder.send().await?;
            let status = response.status();
            let headers = normalize_headers(response.headers());
            let body = decode_body(response, self.policy.max_response_bytes).await?;

            if !status.is_success() {
                return Err(RestAdapterError::UnexpectedStatus {
                    status: status.as_u16(),
                    body,
                });
            }

            Ok(RestResponse {
                status_code: status.as_u16(),
                headers,
                body,
            })
        }
        .instrument(upstream_span.clone())
        .await;
        match &result {
            Ok(_) => StageOutcome::Success.record(&upstream_span),
            Err(_) => {
                StageOutcome::Error.record(&upstream_span);
                ErrorCategory::Upstream.record(&upstream_span);
            }
        }
        result
    }
}

fn upstream_outcome(error: &RestAdapterError) -> UpstreamOutcome {
    match error {
        RestAdapterError::UnexpectedStatus { status, .. } if (400..500).contains(status) => {
            UpstreamOutcome::ClientError
        }
        RestAdapterError::UnexpectedStatus { status, .. } if (500..600).contains(status) => {
            UpstreamOutcome::ServerError
        }
        RestAdapterError::UnexpectedStatus { .. } => UpstreamOutcome::UnexpectedStatus,
        RestAdapterError::Transport(error) if error.is_timeout() => UpstreamOutcome::Timeout,
        RestAdapterError::Transport(_) => UpstreamOutcome::TransportError,
        RestAdapterError::ResponseTooLarge { .. } => UpstreamOutcome::ResponseTooLarge,
        RestAdapterError::TargetNotAllowed { .. } => UpstreamOutcome::Rejected,
        RestAdapterError::WindowExpired => UpstreamOutcome::WindowExpired,
        RestAdapterError::InvalidSseEvent => UpstreamOutcome::InvalidResponse,
        RestAdapterError::InvalidBaseUrl { .. }
        | RestAdapterError::InvalidPathParameter { .. }
        | RestAdapterError::InvalidQueryParameter { .. }
        | RestAdapterError::InvalidHeaderName { .. }
        | RestAdapterError::InvalidHeaderValue { .. } => UpstreamOutcome::InvalidRequest,
        RestAdapterError::InvalidConfiguration { .. } => UpstreamOutcome::Configuration,
    }
}

impl Default for OutboundHttpPolicy {
    fn default() -> Self {
        Self {
            allowed_hosts: Vec::new(),
            denied_hosts: Vec::new(),
            max_response_bytes: DEFAULT_MAX_RESPONSE_BYTES,
        }
    }
}

impl OutboundHttpPolicy {
    pub fn from_env() -> Result<Self, RestAdapterError> {
        let max_response_bytes = match env::var("CRANK_OUTBOUND_MAX_RESPONSE_BYTES") {
            Ok(value) => {
                value
                    .parse::<usize>()
                    .map_err(|_| RestAdapterError::InvalidConfiguration {
                        details: "CRANK_OUTBOUND_MAX_RESPONSE_BYTES must be a positive integer"
                            .to_owned(),
                    })?
            }
            Err(env::VarError::NotPresent) => DEFAULT_MAX_RESPONSE_BYTES,
            Err(error) => {
                return Err(RestAdapterError::InvalidConfiguration {
                    details: error.to_string(),
                });
            }
        };
        if max_response_bytes == 0 {
            return Err(RestAdapterError::InvalidConfiguration {
                details: "CRANK_OUTBOUND_MAX_RESPONSE_BYTES must be greater than zero".to_owned(),
            });
        }

        Ok(Self {
            allowed_hosts: host_patterns_from_env("CRANK_OUTBOUND_ALLOWED_HOSTS")?,
            denied_hosts: host_patterns_from_env("CRANK_OUTBOUND_DENIED_HOSTS")?,
            max_response_bytes,
        })
    }

    pub fn allowing_hosts(hosts: impl IntoIterator<Item = impl Into<String>>) -> Self {
        Self {
            allowed_hosts: hosts.into_iter().map(Into::into).collect(),
            ..Self::default()
        }
    }

    pub fn with_max_response_bytes(mut self, max_response_bytes: usize) -> Self {
        self.max_response_bytes = max_response_bytes;
        self
    }

    pub fn validate_base_url(&self, base_url: &str) -> Result<(), RestAdapterError> {
        let url = reqwest::Url::parse(base_url).map_err(|_| RestAdapterError::InvalidBaseUrl {
            url: base_url.to_owned(),
        })?;
        self.validate_url(&url)
    }

    fn validate_url(&self, url: &reqwest::Url) -> Result<(), RestAdapterError> {
        if !matches!(url.scheme(), "http" | "https")
            || !url.username().is_empty()
            || url.password().is_some()
        {
            return Err(RestAdapterError::TargetNotAllowed {
                target: url.to_string(),
            });
        }
        let host = url
            .host_str()
            .ok_or_else(|| RestAdapterError::TargetNotAllowed {
                target: url.to_string(),
            })?;
        self.validate_host(host)?;
        if !self.is_explicitly_allowed(host) && is_local_hostname(host) {
            return Err(RestAdapterError::TargetNotAllowed {
                target: host.to_owned(),
            });
        }
        if let Ok(address) = host.parse::<IpAddr>()
            && !self.is_explicitly_allowed(host)
            && !is_public_ip(address)
        {
            return Err(RestAdapterError::TargetNotAllowed {
                target: host.to_owned(),
            });
        }
        Ok(())
    }

    fn validate_host(&self, host: &str) -> Result<(), RestAdapterError> {
        let host = normalize_host(host);
        let denied = self
            .denied_hosts
            .iter()
            .any(|pattern| host_matches(pattern, &host));
        if denied {
            return Err(RestAdapterError::TargetNotAllowed { target: host });
        }
        Ok(())
    }

    fn is_explicitly_allowed(&self, host: &str) -> bool {
        let host = normalize_host(host);
        self.allowed_hosts
            .iter()
            .any(|pattern| host_matches(pattern, &host))
    }
}

#[derive(Clone, Debug)]
struct PolicyDnsResolver {
    policy: OutboundHttpPolicy,
}

impl Resolve for PolicyDnsResolver {
    fn resolve(&self, name: Name) -> Resolving {
        let host = normalize_host(name.as_str());
        let policy = self.policy.clone();
        Box::pin(async move {
            policy
                .validate_host(&host)
                .map_err(|error| boxed_io_error(error.to_string()))?;
            let explicitly_allowed = policy.is_explicitly_allowed(&host);
            let resolved = tokio::net::lookup_host((host.as_str(), 0))
                .await
                .map_err(|error| Box::new(error) as Box<dyn std::error::Error + Send + Sync>)?;
            let addresses = resolved
                .filter(|address| explicitly_allowed || is_public_ip(address.ip()))
                .collect::<Vec<SocketAddr>>();
            if addresses.is_empty() {
                return Err(boxed_io_error(format!(
                    "outbound target {host} did not resolve to an allowed address"
                )));
            }
            Ok(Box::new(addresses.into_iter()) as Addrs)
        })
    }
}

fn boxed_io_error(message: String) -> Box<dyn std::error::Error + Send + Sync> {
    Box::new(io::Error::new(io::ErrorKind::PermissionDenied, message))
}

fn host_patterns_from_env(name: &str) -> Result<Vec<String>, RestAdapterError> {
    let value = match env::var(name) {
        Ok(value) => value,
        Err(env::VarError::NotPresent) => return Ok(Vec::new()),
        Err(error) => {
            return Err(RestAdapterError::InvalidConfiguration {
                details: error.to_string(),
            });
        }
    };
    value
        .split(',')
        .map(str::trim)
        .filter(|value| !value.is_empty())
        .map(|value| {
            let wildcard = value.starts_with("*.");
            let normalized = normalize_host(value.trim_start_matches("*."));
            let valid_ip = !wildcard && normalized.parse::<IpAddr>().is_ok();
            if normalized.is_empty()
                || normalized.contains('/')
                || (!valid_ip && normalized.contains(':'))
                || (wildcard && normalized.parse::<IpAddr>().is_ok())
            {
                return Err(RestAdapterError::InvalidConfiguration {
                    details: format!("{name} contains an invalid host pattern: {value}"),
                });
            }
            Ok(if wildcard {
                format!("*.{normalized}")
            } else {
                normalized
            })
        })
        .collect()
}

fn normalize_host(host: &str) -> String {
    host.trim()
        .trim_start_matches('[')
        .trim_end_matches(']')
        .trim_end_matches('.')
        .to_ascii_lowercase()
}

fn is_local_hostname(host: &str) -> bool {
    let host = normalize_host(host);
    host == "localhost" || host.ends_with(".localhost")
}

fn host_matches(pattern: &str, host: &str) -> bool {
    pattern.strip_prefix("*.").map_or_else(
        || pattern == host,
        |suffix| host != suffix && host.ends_with(&format!(".{suffix}")),
    )
}

fn is_public_ip(address: IpAddr) -> bool {
    match address {
        IpAddr::V4(address) => is_public_ipv4(address),
        IpAddr::V6(address) => is_public_ipv6(address),
    }
}

fn is_public_ipv4(address: Ipv4Addr) -> bool {
    let octets = address.octets();
    !(address.is_private()
        || address.is_loopback()
        || address.is_link_local()
        || address.is_broadcast()
        || address.is_documentation()
        || address.is_unspecified()
        || address.is_multicast()
        || octets[0] == 0
        || (octets[0] == 100 && (64..=127).contains(&octets[1]))
        || (octets[0] == 192 && octets[1] == 0 && octets[2] == 0)
        || (octets[0] == 198 && (18..=19).contains(&octets[1]))
        || octets[0] >= 240)
}

fn is_public_ipv6(address: Ipv6Addr) -> bool {
    let segments = address.segments();
    if let Some(address) = address.to_ipv4_mapped() {
        return is_public_ipv4(address);
    }
    if segments[..6].iter().all(|segment| *segment == 0) {
        let [a, b] = segments[6].to_be_bytes();
        let [c, d] = segments[7].to_be_bytes();
        return is_public_ipv4(Ipv4Addr::new(a, b, c, d));
    }
    !(address.is_unspecified()
        || address.is_loopback()
        || address.is_multicast()
        || (segments[0] & 0xfe00) == 0xfc00
        || (segments[0] & 0xffc0) == 0xfe80
        || (segments[0] & 0xffc0) == 0xfec0
        || (segments[0] == 0x0064
            && segments[1] == 0xff9b
            && segments[2..6].iter().all(|segment| *segment == 0))
        || (segments[0] == 0x0064 && segments[1] == 0xff9b && segments[2] == 1)
        || segments[0] == 0x2002
        || (segments[0] == 0x2001 && matches!(segments[1], 0 | 0x0db8)))
}

fn build_url(target: &RestTarget, request: &RestRequest) -> Result<reqwest::Url, RestAdapterError> {
    let base_url =
        reqwest::Url::parse(&target.base_url).map_err(|_| RestAdapterError::InvalidBaseUrl {
            url: target.base_url.clone(),
        })?;
    let path = substitute_path_params(&target.path_template, &request.path_params);
    let mut url = base_url.join(path.trim_start_matches('/')).map_err(|_| {
        RestAdapterError::InvalidBaseUrl {
            url: target.base_url.clone(),
        }
    })?;

    {
        let mut query = url.query_pairs_mut();
        for (key, value) in &request.query_params {
            query.append_pair(key, value);
        }
    }

    Ok(url)
}

fn substitute_path_params(path_template: &str, path_params: &BTreeMap<String, String>) -> String {
    let mut rendered = path_template.to_owned();

    for (key, value) in path_params {
        rendered = rendered.replace(&format!("{{{key}}}"), value);
    }

    rendered
}

fn build_headers(
    target: &RestTarget,
    request: &RestRequest,
) -> Result<HeaderMap, RestAdapterError> {
    let mut headers = HeaderMap::new();

    for (name, value) in &target.static_headers {
        insert_header(&mut headers, name, value)?;
    }

    for (name, value) in &request.headers {
        insert_header(&mut headers, name, value)?;
    }

    Ok(headers)
}

fn insert_header(headers: &mut HeaderMap, name: &str, value: &str) -> Result<(), RestAdapterError> {
    let header_name =
        HeaderName::try_from(name).map_err(|_| RestAdapterError::InvalidHeaderName {
            header: name.to_owned(),
        })?;
    if is_trace_propagation_header(&header_name) {
        return Ok(());
    }
    let header_value =
        HeaderValue::try_from(value).map_err(|_| RestAdapterError::InvalidHeaderValue {
            header: name.to_owned(),
        })?;

    headers.insert(header_name, header_value);
    Ok(())
}

fn is_trace_propagation_header(name: &HeaderName) -> bool {
    matches!(name.as_str(), "traceparent" | "tracestate" | "baggage")
}

fn apply_current_trace_context(headers: &mut HeaderMap) {
    for header in ["traceparent", "tracestate", "baggage"] {
        headers.remove(header);
    }

    let context = Span::current().context();
    if !context.span().span_context().is_valid() {
        return;
    }
    global::get_text_map_propagator(|propagator| {
        propagator.inject_context(&context, &mut ReqwestHeaderInjector(headers));
    });
}

struct ReqwestHeaderInjector<'a>(&'a mut HeaderMap);

impl Injector for ReqwestHeaderInjector<'_> {
    fn set(&mut self, key: &str, value: String) {
        let Ok(name) = HeaderName::try_from(key) else {
            return;
        };
        let Ok(value) = HeaderValue::try_from(value) else {
            return;
        };
        self.0.insert(name, value);
    }
}

async fn decode_body(
    response: reqwest::Response,
    max_response_bytes: usize,
) -> Result<Value, RestAdapterError> {
    if response
        .content_length()
        .is_some_and(|length| length > max_response_bytes as u64)
    {
        return Err(RestAdapterError::ResponseTooLarge {
            limit_bytes: max_response_bytes,
        });
    }
    let mut stream = response.bytes_stream();
    let mut bytes = Vec::new();
    while let Some(chunk) = stream.next().await {
        let chunk = chunk?;
        if bytes.len().saturating_add(chunk.len()) > max_response_bytes {
            return Err(RestAdapterError::ResponseTooLarge {
                limit_bytes: max_response_bytes,
            });
        }
        bytes.extend_from_slice(&chunk);
    }

    if bytes.is_empty() {
        return Ok(Value::Null);
    }

    match serde_json::from_slice::<Value>(&bytes) {
        Ok(value) => Ok(value),
        Err(_) => Ok(Value::String(
            String::from_utf8_lossy(&bytes).trim().to_owned(),
        )),
    }
}

fn normalize_headers(headers: &HeaderMap) -> BTreeMap<String, String> {
    headers
        .iter()
        .filter_map(|(name, value)| {
            value
                .to_str()
                .ok()
                .map(|value| (name.as_str().to_owned(), value.to_owned()))
        })
        .collect()
}

fn to_reqwest_method(method: HttpMethod) -> reqwest::Method {
    match method {
        HttpMethod::Get => reqwest::Method::GET,
        HttpMethod::Post => reqwest::Method::POST,
        HttpMethod::Put => reqwest::Method::PUT,
        HttpMethod::Patch => reqwest::Method::PATCH,
        HttpMethod::Delete => reqwest::Method::DELETE,
    }
}