Add middleware for header count and header list item limits in limits plugin (#8285)

Author: smyrick

apollo-router/src/plugins/limits/fixtures/header_count_limit.router.yaml

@@ -0,0 +1,2 @@
+limits:
+  http_max_request_headers: 5

apollo-router/src/plugins/limits/fixtures/header_list_items_limit.router.yaml

@@ -0,0 +1,2 @@
+limits:
+  http_max_header_list_items: 3

apollo-router/src/plugins/limits/layer.rs

@@ -18,6 +18,14 @@ pub(super) enum BodyLimitError {
     PayloadTooLarge,
 }
 
+#[derive(thiserror::Error, Debug, Display)]
+pub(super) enum HeaderLimitError {
+    /// Request header fields too many
+    TooManyHeaders,
+    /// Request header list too many items
+    TooManyHeaderListItems,
+}
+
 struct BodyLimitControlInner {
     limit: AtomicUsize,
     current: AtomicUsize,
@@ -416,3 +424,113 @@ mod test {
         }
     }
 }
+
+/// Layer that limits the number of headers in an HTTP request
+pub(crate) struct RequestHeaderCountLimitLayer {
+    max_headers: Option<usize>,
+}
+
+impl RequestHeaderCountLimitLayer {
+    pub(crate) fn new(max_headers: Option<usize>) -> Self {
+        Self { max_headers }
+    }
+}
+
+impl<S> Layer<S> for RequestHeaderCountLimitLayer {
+    type Service = RequestHeaderCountLimit<S>;
+
+    fn layer(&self, inner: S) -> Self::Service {
+        RequestHeaderCountLimit::new(inner, self.max_headers)
+    }
+}
+
+pub(crate) struct RequestHeaderCountLimit<S> {
+    inner: S,
+    max_headers: Option<usize>,
+}
+
+impl<S> RequestHeaderCountLimit<S> {
+    fn new(inner: S, max_headers: Option<usize>) -> Self {
+        Self { inner, max_headers }
+    }
+}
+
+impl<ReqBody, RespBody, S> Service<http::Request<ReqBody>> for RequestHeaderCountLimit<S>
+where
+    S: Service<http::Request<ReqBody>, Response = http::Response<RespBody>>,
+    S::Error: From<HeaderLimitError>,
+{
+    type Response = S::Response;
+    type Error = S::Error;
+    type Future = futures::future::Either<futures::future::Ready<Result<S::Response, S::Error>>, S::Future>;
+
+    fn poll_ready(&mut self, cx: &mut std::task::Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.inner.poll_ready(cx)
+    }
+
+    fn call(&mut self, req: http::Request<ReqBody>) -> Self::Future {
+        if let Some(max_headers) = self.max_headers {
+            if req.headers().len() > max_headers {
+                // This will be converted to the appropriate response by the error handling
+                return futures::future::Either::Left(futures::future::ready(Err(HeaderLimitError::TooManyHeaders.into())));
+            }
+        }
+        futures::future::Either::Right(self.inner.call(req))
+    }
+}
+
+/// Layer that limits the number of items in header lists (for headers with multiple values)
+pub(crate) struct RequestHeaderListItemsLimitLayer {
+    max_items: Option<usize>,
+}
+
+impl RequestHeaderListItemsLimitLayer {
+    pub(crate) fn new(max_items: Option<usize>) -> Self {
+        Self { max_items }
+    }
+}
+
+impl<S> Layer<S> for RequestHeaderListItemsLimitLayer {
+    type Service = RequestHeaderListItemsLimit<S>;
+
+    fn layer(&self, inner: S) -> Self::Service {
+        RequestHeaderListItemsLimit::new(inner, self.max_items)
+    }
+}
+
+pub(crate) struct RequestHeaderListItemsLimit<S> {
+    inner: S,
+    max_items: Option<usize>,
+}
+
+impl<S> RequestHeaderListItemsLimit<S> {
+    fn new(inner: S, max_items: Option<usize>) -> Self {
+        Self { inner, max_items }
+    }
+}
+
+impl<ReqBody, RespBody, S> Service<http::Request<ReqBody>> for RequestHeaderListItemsLimit<S>
+where
+    S: Service<http::Request<ReqBody>, Response = http::Response<RespBody>>,
+    S::Error: From<HeaderLimitError>,
+{
+    type Response = S::Response;
+    type Error = S::Error;
+    type Future = futures::future::Either<futures::future::Ready<Result<S::Response, S::Error>>, S::Future>;
+
+    fn poll_ready(&mut self, cx: &mut std::task::Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.inner.poll_ready(cx)
+    }
+
+    fn call(&mut self, req: http::Request<ReqBody>) -> Self::Future {
+        if let Some(max_items) = self.max_items {
+            for (header_name, _) in req.headers().iter() {
+                let header_values_count = req.headers().get_all(header_name).iter().count();
+                if header_values_count > max_items {
+                    return futures::future::Either::Left(futures::future::ready(Err(HeaderLimitError::TooManyHeaderListItems.into())));
+                }
+            }
+        }
+        futures::future::Either::Right(self.inner.call(req))
+    }
+}

apollo-router/src/plugins/limits/mod.rs

@@ -19,7 +19,10 @@ use crate::layers::ServiceBuilderExt;
 use crate::plugin::Plugin;
 use crate::plugin::PluginInit;
 use crate::plugins::limits::layer::BodyLimitError;
+use crate::plugins::limits::layer::HeaderLimitError;
 use crate::plugins::limits::layer::RequestBodyLimitLayer;
+use crate::plugins::limits::layer::RequestHeaderCountLimitLayer;
+use crate::plugins::limits::layer::RequestHeaderListItemsLimitLayer;
 use crate::services::router;
 use crate::services::router::BoxService;
 
@@ -116,6 +119,20 @@ pub(crate) struct Config {
     #[schemars(with = "Option<String>", default)]
     pub(crate) http1_max_request_buf_size: Option<ByteSize>,
 
+    /// Limit the maximum number of headers in an HTTP request.
+    ///
+    /// If router receives more headers than this limit, it responds to the client with
+    /// "431 Request Header Fields Too Large".
+    /// When not specified, no limit is enforced at the middleware level.
+    pub(crate) http_max_request_headers: Option<usize>,
+
+    /// Limit the maximum number of items in a header list (for headers with multiple values).
+    ///
+    /// If a single header has more values than this limit, it responds to the client with
+    /// "431 Request Header Fields Too Large".
+    /// When not specified, no limit is enforced at the middleware level.
+    pub(crate) http_max_header_list_items: Option<usize>,
+
     /// Limit the depth of nested list fields in introspection queries
     /// to protect avoid generating huge responses. Returns a GraphQL
     /// error with `{ message: "Maximum introspection depth exceeded" }`
@@ -136,6 +153,8 @@ impl Default for Config {
             http_max_request_bytes: 2_000_000,
             http1_max_request_headers: None,
             http1_max_request_buf_size: None,
+            http_max_request_headers: None,
+            http_max_header_list_items: None,
             parser_max_tokens: 15_000,
 
             // This is `apollo-parser`’s default, which protects against stack overflow
@@ -174,6 +193,12 @@ impl Plugin for LimitsPlugin {
             // Here we need to convert to and from the underlying http request types so that we can use existing middleware.
             .map_request(Into::into)
             .map_response(Into::into)
+            .layer(RequestHeaderCountLimitLayer::new(
+                self.config.http_max_request_headers,
+            ))
+            .layer(RequestHeaderListItemsLimitLayer::new(
+                self.config.http_max_header_list_items,
+            ))
             .layer(RequestBodyLimitLayer::new(
                 self.config.http_max_request_bytes,
             ))
@@ -209,7 +234,10 @@ impl LimitsPlugin {
                 }
 
                 match root_cause.downcast_ref::<BodyLimitError>() {
-                    None => Err(e),
+                    None => match root_cause.downcast_ref::<HeaderLimitError>() {
+                        None => Err(e),
+                        Some(header_error) => Ok(header_error.into_response(ctx)),
+                    },
                     Some(_) => Ok(BodyLimitError::PayloadTooLarge.into_response(ctx)),
                 }
             }
@@ -236,6 +264,37 @@ impl BodyLimitError {
     }
 }
 
+impl HeaderLimitError {
+    fn into_response(&self, ctx: Context) -> router::Response {
+        match self {
+            HeaderLimitError::TooManyHeaders => router::Response::error_builder()
+                .error(
+                    graphql::Error::builder()
+                        .message("Request header fields too many")
+                        .extension_code("INVALID_GRAPHQL_REQUEST")
+                        .extension("details", "Request header fields too many")
+                        .build(),
+                )
+                .status_code(StatusCode::REQUEST_HEADER_FIELDS_TOO_LARGE)
+                .context(ctx)
+                .build()
+                .unwrap(),
+            HeaderLimitError::TooManyHeaderListItems => router::Response::error_builder()
+                .error(
+                    graphql::Error::builder()
+                        .message("Request header list too many items")
+                        .extension_code("INVALID_GRAPHQL_REQUEST")
+                        .extension("details", "Request header list too many items")
+                        .build(),
+                )
+                .status_code(StatusCode::REQUEST_HEADER_FIELDS_TOO_LARGE)
+                .context(ctx)
+                .build()
+                .unwrap(),
+        }
+    }
+}
+
 register_plugin!("apollo", "limits", LimitsPlugin);
 
 #[cfg(test)]
@@ -430,6 +489,121 @@ mod test {
         );
     }
 
+    #[tokio::test]
+    async fn test_header_count_limit_exceeded() {
+        let plugin = header_count_plugin().await;
+        let resp = plugin
+            .router_service(|_| async { panic!("should have rejected request") })
+            .call(
+                router::Request::fake_builder()
+                    .header("header1", "value1")
+                    .header("header2", "value2")
+                    .header("header3", "value3")
+                    .header("header4", "value4")
+                    .header("header5", "value5")
+                    .header("header6", "value6") // This should exceed the limit of 5
+                    .body(router::body::empty())
+                    .build()
+                    .unwrap(),
+            )
+            .await;
+        assert!(resp.is_ok());
+        let resp = resp.unwrap();
+        assert_eq!(resp.response.status(), StatusCode::REQUEST_HEADER_FIELDS_TOO_LARGE);
+        let body_str = String::from_utf8(
+            router::body::into_bytes(resp.response.into_body())
+                .await
+                .unwrap()
+                .to_vec()
+        ).unwrap();
+        assert!(body_str.contains("Request header fields too many"));
+    }
+
+    #[tokio::test]
+    async fn test_header_count_limit_ok() {
+        let plugin = header_count_plugin().await;
+        let resp = plugin
+            .router_service(|_| async { Ok(router::Response::fake_builder().build().unwrap()) })
+            .call(
+                router::Request::fake_builder()
+                    .header("header1", "value1")
+                    .header("header2", "value2")
+                    .header("header3", "value3")
+                    .body(router::body::empty())
+                    .build()
+                    .unwrap(),
+            )
+            .await;
+        assert!(resp.is_ok());
+        let resp = resp.unwrap();
+        assert_eq!(resp.response.status(), StatusCode::OK);
+    }
+
+    #[tokio::test]
+    async fn test_header_list_items_limit_exceeded() {
+        let plugin = header_list_items_plugin().await;
+        let mut request = router::Request::fake_builder()
+            .body(router::body::empty());
+        
+        // Create a request with a header that has 4 values (exceeds limit of 3)
+        for i in 1..=4 {
+            request = request.header("test-header", format!("value{}", i));
+        }
+        
+        let resp = plugin
+            .router_service(|_| async { panic!("should have rejected request") })
+            .call(request.build().unwrap())
+            .await;
+        assert!(resp.is_ok());
+        let resp = resp.unwrap();
+        assert_eq!(resp.response.status(), StatusCode::REQUEST_HEADER_FIELDS_TOO_LARGE);
+        let body_str = String::from_utf8(
+            router::body::into_bytes(resp.response.into_body())
+                .await
+                .unwrap()
+                .to_vec()
+        ).unwrap();
+        assert!(body_str.contains("Request header list too many items"));
+    }
+
+    #[tokio::test]
+    async fn test_header_list_items_limit_ok() {
+        let plugin = header_list_items_plugin().await;
+        let mut request = router::Request::fake_builder()
+            .body(router::body::empty());
+        
+        // Create a request with a header that has 2 values (within limit of 3)
+        for i in 1..=2 {
+            request = request.header("test-header", format!("value{}", i));
+        }
+        
+        let resp = plugin
+            .router_service(|_| async { Ok(router::Response::fake_builder().build().unwrap()) })
+            .call(request.build().unwrap())
+            .await;
+        assert!(resp.is_ok());
+        let resp = resp.unwrap();
+        assert_eq!(resp.response.status(), StatusCode::OK);
+    }
+
+    async fn header_count_plugin() -> PluginTestHarness<LimitsPlugin> {
+        let plugin: PluginTestHarness<LimitsPlugin> = PluginTestHarness::builder()
+            .config(include_str!("fixtures/header_count_limit.router.yaml"))
+            .build()
+            .await
+            .expect("test harness");
+        plugin
+    }
+
+    async fn header_list_items_plugin() -> PluginTestHarness<LimitsPlugin> {
+        let plugin: PluginTestHarness<LimitsPlugin> = PluginTestHarness::builder()
+            .config(include_str!("fixtures/header_list_items_limit.router.yaml"))
+            .build()
+            .await
+            .expect("test harness");
+        plugin
+    }
+
     async fn plugin() -> PluginTestHarness<LimitsPlugin> {
         let plugin: PluginTestHarness<LimitsPlugin> = PluginTestHarness::builder()
             .config(include_str!("fixtures/content_length_limit.router.yaml"))