diff --git a/bottlecap/Cargo.lock b/bottlecap/Cargo.lock index 340e0afb7..bd91694b3 100644 --- a/bottlecap/Cargo.lock +++ b/bottlecap/Cargo.lock @@ -568,6 +568,7 @@ dependencies = [ "tikv-jemallocator", "tokio", "tokio-util", + "tower-http", "tracing", "tracing-core", "tracing-subscriber", @@ -3720,6 +3721,7 @@ dependencies = [ "futures-util", "http 1.3.1", "http-body 1.0.1", + "http-body-util", "iri-string", "pin-project-lite", "tower 0.5.2", diff --git a/bottlecap/Cargo.toml b/bottlecap/Cargo.toml index 9a162d5f7..bc295522b 100644 --- a/bottlecap/Cargo.toml +++ b/bottlecap/Cargo.toml @@ -49,6 +49,7 @@ opentelemetry-semantic-conventions = { version = "0.30", features = ["semconv_ex rustls-native-certs = { version = "0.8.1", optional = true } axum = { version = "0.8.4", default-features = false, features = ["default"] } ustr = { version = "1.0.0", default-features = false } +tower-http = { version = "0.6.6", default-features = false, features = ["limit"] } # If you are adding or updating a datadog-owned code dependency, please ensure # that it has a clippy.toml rule for disallowing the reqwest::Client::builder # method in favor of our diff --git a/bottlecap/src/traces/trace_agent.rs b/bottlecap/src/traces/trace_agent.rs index 03500894a..27b0132fd 100644 --- a/bottlecap/src/traces/trace_agent.rs +++ b/bottlecap/src/traces/trace_agent.rs @@ -3,7 +3,7 @@ use axum::{ Router, - extract::{Request, State}, + extract::{DefaultBodyLimit, Request, State}, http::StatusCode, response::{IntoResponse, Response}, routing::{any, post}, @@ -18,6 +18,7 @@ use tokio::sync::{ mpsc::{self, Receiver, Sender}, }; use tokio_util::sync::CancellationToken; +use tower_http::limit::RequestBodyLimitLayer; use tracing::{debug, error}; use crate::traces::trace_processor::SendingTraceProcessor; @@ -72,7 +73,9 @@ const INSTRUMENTATION_INTAKE_PATH: &str = "/api/v2/apmtelemetry"; const TRACER_PAYLOAD_CHANNEL_BUFFER_SIZE: usize = 10; const STATS_PAYLOAD_CHANNEL_BUFFER_SIZE: usize = 10; -pub const MAX_CONTENT_LENGTH: usize = 10 * 1024 * 1024; +pub const TRACE_REQUEST_BODY_LIMIT: usize = 50 * 1024 * 1024; +pub const DEFAULT_REQUEST_BODY_LIMIT: usize = 2 * 1024 * 1024; +pub const MAX_CONTENT_LENGTH: usize = 50 * 1024 * 1024; const LAMBDA_LOAD_SPAN: &str = "aws.lambda.load"; #[derive(Clone)] @@ -231,10 +234,12 @@ impl TraceAgent { V5_TRACE_ENDPOINT_PATH, post(Self::v05_traces).put(Self::v05_traces), ) + .layer(RequestBodyLimitLayer::new(TRACE_REQUEST_BODY_LIMIT)) .with_state(trace_state); let stats_router = Router::new() .route(STATS_ENDPOINT_PATH, post(Self::stats).put(Self::stats)) + .layer(RequestBodyLimitLayer::new(DEFAULT_REQUEST_BODY_LIMIT)) .with_state(stats_state); let proxy_router = Router::new() @@ -254,6 +259,7 @@ impl TraceAgent { INSTRUMENTATION_ENDPOINT_PATH, post(Self::instrumentation_proxy), ) + .layer(RequestBodyLimitLayer::new(DEFAULT_REQUEST_BODY_LIMIT)) .with_state(proxy_state); let info_router = Router::new().route(INFO_ENDPOINT_PATH, any(Self::info)); @@ -264,6 +270,8 @@ impl TraceAgent { .merge(proxy_router) .merge(info_router) .fallback(handler_not_found) + // Disable the default body limit so we can use our own limit + .layer(DefaultBodyLimit::disable()) } async fn graceful_shutdown(shutdown_token: CancellationToken) {