diff --git a/bottlecap/src/http.rs b/bottlecap/src/http.rs index ec9b821a6..f6df95871 100644 --- a/bottlecap/src/http.rs +++ b/bottlecap/src/http.rs @@ -13,7 +13,7 @@ use tracing::{debug, error, warn}; #[must_use] pub fn get_client(config: &Arc) -> reqwest::Client { - build_client(config).unwrap_or_else(|e| { + build_client(config, true).unwrap_or_else(|e| { error!( "Unable to parse proxy configuration: {}, no proxy will be used", e @@ -23,7 +23,28 @@ pub fn get_client(config: &Arc) -> reqwest::Client { }) } -fn build_client(config: &Arc) -> Result> { +/// Same as [`get_client`], but never follows redirects. reqwest keeps custom +/// headers such as `DD-API-KEY` when it follows a redirect to another host. +/// Returns `None` rather than fall back to a client that follows redirects. +#[must_use] +pub fn get_client_without_redirects(config: &Arc) -> Option { + build_client(config, false) + .or_else(|e| { + error!("Unable to parse proxy configuration: {e}, no proxy will be used"); + Ok::<_, Box>( + create_reqwest_client_builder()? + .redirect(reqwest::redirect::Policy::none()) + .build()?, + ) + }) + .inspect_err(|e| error!("Unable to build an HTTP client without redirects: {e}")) + .ok() +} + +fn build_client( + config: &Arc, + follow_redirects: bool, +) -> Result> { let mut client = create_reqwest_client_builder()? .timeout(Duration::from_secs(config.flush_timeout)) // Disable connection pooling to avoid stale connections after Lambda freeze/resume. @@ -33,6 +54,9 @@ fn build_client(config: &Arc) -> Result, + /// Must not follow redirects, so the API key only goes to `flag_config_base_url`. + client: reqwest::Client, + flag_config_base_url: String, + env: Option, + timeout: Duration, +} + +impl Passthrough { + #[must_use] + pub fn new( + config: &Config, + api_key_factory: Arc, + client: reqwest::Client, + ) -> Self { + Self { + api_key_factory, + client, + flag_config_base_url: format!("https://{FLAG_CONFIG_SUBDOMAIN}.{}", config.site), + env: config.env.clone(), + timeout: UPSTREAM_TIMEOUT, + } + } + + pub fn router(self) -> Router { + Router::new() + .route(FLAG_CONFIG_PATH, get(flag_config)) + .with_state(self) + } +} + +async fn flag_config( + State(passthrough): State, + RawQuery(query): RawQuery, + headers: HeaderMap, +) -> Response { + let deadline = Instant::now() + passthrough.timeout; + // Resolve the key in its own task: if the tracer disconnects, this handler + // is dropped, and a key load cancelled midway would leave every other + // caller of the factory waiting on it forever. + let factory = Arc::clone(&passthrough.api_key_factory); + let api_key = match timeout_at( + deadline, + tokio::spawn(async move { factory.get_api_key().await }), + ) + .await + { + Ok(Ok(Some(api_key))) => api_key, + // 403 makes the flag provider log the failure without retrying. + Ok(Ok(None)) => { + error!( + "TRACE_AGENT | Failed to resolve API key, can't forward the flag config request" + ); + return (StatusCode::FORBIDDEN, "Datadog API key unavailable").into_response(); + } + // A 5xx is retried with backoff, and the key load keeps running. + Ok(Err(e)) => { + error!("TRACE_AGENT | API key resolution failed: {e}"); + return StatusCode::BAD_GATEWAY.into_response(); + } + Err(_) => { + error!("TRACE_AGENT | Timed out resolving the API key for the flag config request"); + return StatusCode::BAD_GATEWAY.into_response(); + } + }; + + let mut url = format!("{}{FLAG_CONFIG_PATH}", passthrough.flag_config_base_url); + if let Some(query) = &query { + url.push('?'); + url.push_str(query); + } + + let mut request = passthrough + .client + .get(url) + .timeout(deadline.saturating_duration_since(Instant::now())) + .header("DD-API-KEY", api_key); + // The tracer drops `dd_env` when it uses a custom base URL. + let has_env = query.as_deref().is_some_and(|q| { + q.split('&') + .any(|pair| pair == "dd_env" || pair.starts_with("dd_env=")) + }); + if let (Some(env), false) = (&passthrough.env, has_env) { + request = request.query(&[("dd_env", env)]); + } + for name in FORWARDED_REQUEST_HEADERS { + if let Some(value) = headers.get(&name) { + request = request.header(name, value); + } + } + + let upstream = match request.send().await { + Ok(upstream) => upstream, + Err(e) => { + error!("TRACE_AGENT | Flag config request to Datadog failed: {e}"); + return StatusCode::BAD_GATEWAY.into_response(); + } + }; + + let status = upstream.status(); + let mut response_headers = HeaderMap::new(); + for name in FORWARDED_RESPONSE_HEADERS { + if let Some(value) = upstream.headers().get(&name) { + response_headers.insert(name, value.clone()); + } + } + match upstream.bytes().await { + Ok(body) => { + debug!("TRACE_AGENT | Forwarded flag config request: {status}"); + (status, response_headers, body).into_response() + } + Err(e) => { + error!("TRACE_AGENT | Failed to read the flag config response: {e}"); + StatusCode::BAD_GATEWAY.into_response() + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use axum::body::{Body, to_bytes}; + use axum::http::Request; + use httpmock::prelude::*; + use tower::ServiceExt; + + fn passthrough(base_url: String, api_key_factory: ApiKeyFactory) -> Passthrough { + let config = Arc::new(Config { + http_protocol: Some("http1".to_string()), + ..Config::default() + }); + Passthrough { + api_key_factory: Arc::new(api_key_factory), + client: crate::http::get_client_without_redirects(&config).expect("client"), + flag_config_base_url: base_url, + env: Some("prod".to_string()), + timeout: UPSTREAM_TIMEOUT, + } + } + + async fn call(passthrough: Passthrough, uri: &str, if_none_match: Option<&str>) -> Response { + let mut request = Request::get(uri); + if let Some(etag) = if_none_match { + request = request.header(header::IF_NONE_MATCH, etag); + } + passthrough + .router() + .oneshot(request.body(Body::empty()).expect("request")) + .await + .expect("response") + } + + #[tokio::test] + async fn forwards_with_api_key_and_env_and_returns_the_reply() { + let server = MockServer::start_async().await; + let mock = server + .mock_async(|when, then| { + when.method(GET) + .path(FLAG_CONFIG_PATH) + .header("DD-API-KEY", "test-key") + .query_param("dd_env", "prod"); + then.status(200) + .header("ETag", "\"abc\"") + .header("Content-Encoding", "gzip") + .body("rules"); + }) + .await; + + let response = call( + passthrough(server.base_url(), ApiKeyFactory::new("test-key")), + FLAG_CONFIG_PATH, + None, + ) + .await; + + mock.assert_async().await; + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(response.headers()[header::ETAG], "\"abc\""); + assert_eq!(response.headers()[header::CONTENT_ENCODING], "gzip"); + let body = to_bytes(response.into_body(), usize::MAX) + .await + .expect("body"); + assert_eq!(&body[..], b"rules"); + } + + #[tokio::test] + async fn passes_not_modified_through() { + let server = MockServer::start_async().await; + let mock = server + .mock_async(|when, then| { + when.method(GET) + .path(FLAG_CONFIG_PATH) + .header("If-None-Match", "\"abc\""); + then.status(304); + }) + .await; + + let response = call( + passthrough(server.base_url(), ApiKeyFactory::new("test-key")), + FLAG_CONFIG_PATH, + Some("\"abc\""), + ) + .await; + + mock.assert_async().await; + assert_eq!(response.status(), StatusCode::NOT_MODIFIED); + } + + #[tokio::test] + async fn keeps_an_existing_env_query() { + let server = MockServer::start_async().await; + let mock = server + .mock_async(|when, then| { + when.method(GET) + .path(FLAG_CONFIG_PATH) + .query_param("dd_env", "staging"); + then.status(200); + }) + .await; + + let response = call( + passthrough(server.base_url(), ApiKeyFactory::new("test-key")), + &format!("{FLAG_CONFIG_PATH}?dd_env=staging"), + None, + ) + .await; + + mock.assert_async().await; + assert_eq!(response.status(), StatusCode::OK); + } + + #[tokio::test] + async fn replies_forbidden_without_an_api_key() { + let no_key = ApiKeyFactory::new_from_resolver(Arc::new(|| Box::pin(async { None })), None); + + let response = call( + passthrough("http://127.0.0.1:1".to_string(), no_key), + FLAG_CONFIG_PATH, + None, + ) + .await; + + assert_eq!(response.status(), StatusCode::FORBIDDEN); + } + + #[tokio::test] + async fn replies_bad_gateway_when_datadog_is_unreachable() { + let response = call( + passthrough( + "http://127.0.0.1:1".to_string(), + ApiKeyFactory::new("test-key"), + ), + FLAG_CONFIG_PATH, + None, + ) + .await; + + assert_eq!(response.status(), StatusCode::BAD_GATEWAY); + } + + #[tokio::test] + async fn does_not_follow_redirects() { + let elsewhere = MockServer::start_async().await; + let leaked = elsewhere + .mock_async(|when, then| { + when.any_request(); + then.status(200); + }) + .await; + let server = MockServer::start_async().await; + server + .mock_async(|when, then| { + when.method(GET).path(FLAG_CONFIG_PATH); + then.status(302) + .header("Location", elsewhere.url(FLAG_CONFIG_PATH)); + }) + .await; + + let response = call( + passthrough(server.base_url(), ApiKeyFactory::new("test-key")), + FLAG_CONFIG_PATH, + None, + ) + .await; + + assert_eq!(response.status(), StatusCode::FOUND); + assert_eq!(leaked.hits_async().await, 0); + } + + #[tokio::test] + async fn slow_api_key_times_out_and_keeps_loading() { + let server = MockServer::start_async().await; + let mock = server + .mock_async(|when, then| { + when.method(GET) + .path(FLAG_CONFIG_PATH) + .header("DD-API-KEY", "slow-key"); + then.status(200); + }) + .await; + let slow_key = ApiKeyFactory::new_from_resolver( + Arc::new(|| { + Box::pin(async { + tokio::time::sleep(Duration::from_millis(300)).await; + Some("slow-key".to_string()) + }) + }), + None, + ); + let passthrough = Passthrough { + timeout: Duration::from_millis(100), + ..passthrough(server.base_url(), slow_key) + }; + + let first = call(passthrough.clone(), FLAG_CONFIG_PATH, None).await; + assert_eq!(first.status(), StatusCode::BAD_GATEWAY); + + // The load finishes in the background; a load cancelled midway would + // leave the factory waiting forever and the next request would hang. + tokio::time::sleep(Duration::from_millis(400)).await; + let second = tokio::time::timeout( + Duration::from_secs(5), + call(passthrough, FLAG_CONFIG_PATH, None), + ) + .await + .expect("the key load never finished"); + assert_eq!(second.status(), StatusCode::OK); + mock.assert_async().await; + } + + #[tokio::test] + async fn does_not_forward_other_paths() { + let response = call( + passthrough( + "http://127.0.0.1:1".to_string(), + ApiKeyFactory::new("test-key"), + ), + "/api/v2/other", + None, + ) + .await; + + assert_eq!(response.status(), StatusCode::NOT_FOUND); + } +} diff --git a/bottlecap/src/traces/trace_agent.rs b/bottlecap/src/traces/trace_agent.rs index bb7bdf6cf..33c0fbf1b 100644 --- a/bottlecap/src/traces/trace_agent.rs +++ b/bottlecap/src/traces/trace_agent.rs @@ -34,6 +34,7 @@ use crate::{ tags::{self, provider}, traces::{ INVOCATION_SPAN_RESOURCE, + passthrough::Passthrough, proxy_aggregator::{self, ProxyRequest}, span_dedup::DedupKey, span_dedup_service::DedupHandle, @@ -247,6 +248,9 @@ pub struct TraceAgent { /// `None` when the caller wants the default port. See /// [`TraceAgent::with_receiver_port`]. receiver_port: Option, + /// Routes forwarded to Datadog with the API key. See + /// [`TraceAgent::with_passthrough`]. + passthrough: Option, } #[derive(Clone, Copy)] @@ -345,9 +349,18 @@ impl TraceAgent { span_deduper, router_extension: None, receiver_port: None, + passthrough: None, } } + /// Adds the pass-through routes (for example feature flag config), which + /// forward to Datadog right away with the resolved API key. + #[must_use] + pub fn with_passthrough(mut self, passthrough: Passthrough) -> Self { + self.passthrough = Some(passthrough); + self + } + /// Attaches a [`RouterExtension`] that will be applied to the /// fully-assembled production router before the outer fallback and /// body-limit layers. Without this call, the agent exposes only its @@ -473,6 +486,10 @@ impl TraceAgent { .merge(proxy_router) .merge(info_router); + if let Some(passthrough) = &self.passthrough { + router = router.merge(passthrough.clone().router()); + } + if let Some(extension) = &self.router_extension { router = extension.extend(router)?; }