From f6015e3aefb58bdca42b4fc9f63e1be13f49ce56 Mon Sep 17 00:00:00 2001
From: Yu Li <liyu.yukiteru@bytedance.com>
Date: Fri, 12 Jan 2024 19:33:50 +0800
Subject: [PATCH] chore(http1): add test cases for `max_headers`

Signed-off-by: Yu Li <liyu.yukiteru@bytedance.com>
---
 src/proto/h1/role.rs | 119 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 119 insertions(+)

diff --git a/src/proto/h1/role.rs b/src/proto/h1/role.rs
index 4aae355401..a51169e885 100644
--- a/src/proto/h1/role.rs
+++ b/src/proto/h1/role.rs
@@ -2698,6 +2698,125 @@ mod tests {
         assert_eq!(parsed.head.headers["server"], "hello\tworld");
     }
 
+    #[test]
+    fn parse_too_large_headers() {
+        fn gen_req_with_headers(num: usize) -> String {
+            let mut req = String::from("GET / HTTP/1.1\r\n");
+            for i in 0..num {
+                req.push_str(&format!("key{i}: val{i}\r\n"));
+            }
+            req.push_str("\r\n");
+            req
+        }
+        fn gen_resp_with_headers(num: usize) -> String {
+            let mut req = String::from("HTTP/1.1 200 OK\r\n");
+            for i in 0..num {
+                req.push_str(&format!("key{i}: val{i}\r\n"));
+            }
+            req.push_str("\r\n");
+            req
+        }
+        fn parse(gen_size: usize, max_headers: Option<usize>, should_success: bool) {
+            {
+                // server side
+                let mut bytes = BytesMut::from(gen_req_with_headers(gen_size).as_str());
+                let result = Server::parse(
+                    &mut bytes,
+                    ParseContext {
+                        cached_headers: &mut None,
+                        req_method: &mut None,
+                        h1_parser_config: Default::default(),
+                        h1_max_headers: max_headers.unwrap_or(DEFAULT_MAX_HEADERS),
+                        h1_header_read_timeout: None,
+                        h1_header_read_timeout_fut: &mut None,
+                        h1_header_read_timeout_running: &mut false,
+                        timer: Time::Empty,
+                        preserve_header_case: false,
+                        #[cfg(feature = "ffi")]
+                        preserve_header_order: false,
+                        h09_responses: false,
+                        #[cfg(feature = "ffi")]
+                        on_informational: &mut None,
+                    },
+                );
+                if should_success {
+                    result.expect("parse ok").expect("parse complete");
+                } else {
+                    result.expect_err("parse should err");
+                }
+            }
+            {
+                // client side
+                let mut bytes = BytesMut::from(gen_resp_with_headers(gen_size).as_str());
+                let result = Client::parse(
+                    &mut bytes,
+                    ParseContext {
+                        cached_headers: &mut None,
+                        req_method: &mut None,
+                        h1_parser_config: Default::default(),
+                        h1_max_headers: max_headers.unwrap_or(DEFAULT_MAX_HEADERS),
+                        h1_header_read_timeout: None,
+                        h1_header_read_timeout_fut: &mut None,
+                        h1_header_read_timeout_running: &mut false,
+                        timer: Time::Empty,
+                        preserve_header_case: false,
+                        #[cfg(feature = "ffi")]
+                        preserve_header_order: false,
+                        h09_responses: false,
+                        #[cfg(feature = "ffi")]
+                        on_informational: &mut None,
+                    },
+                );
+                if should_success {
+                    result.expect("parse ok").expect("parse complete");
+                } else {
+                    result.expect_err("parse should err");
+                }
+            }
+        }
+
+        // check generator
+        assert_eq!(
+            gen_req_with_headers(0),
+            String::from("GET / HTTP/1.1\r\n\r\n")
+        );
+        assert_eq!(
+            gen_req_with_headers(1),
+            String::from("GET / HTTP/1.1\r\nkey0: val0\r\n\r\n")
+        );
+        assert_eq!(
+            gen_req_with_headers(2),
+            String::from("GET / HTTP/1.1\r\nkey0: val0\r\nkey1: val1\r\n\r\n")
+        );
+        assert_eq!(
+            gen_req_with_headers(3),
+            String::from("GET / HTTP/1.1\r\nkey0: val0\r\nkey1: val1\r\nkey2: val2\r\n\r\n")
+        );
+
+        // default max_headers is 100
+        parse(0, None, true);
+        parse(1, None, true);
+        parse(50, None, true);
+        parse(99, None, true);
+        parse(100, None, true);
+        parse(101, None, false);
+        parse(102, None, false);
+        parse(200, None, false);
+
+        // max_headers is 0, parser does not accept headers
+        parse(0, Some(0), true);
+        parse(1, Some(0), false);
+        parse(100, Some(0), false);
+
+        // max_headers is 200
+        parse(0, Some(200), true);
+        parse(1, Some(200), true);
+        parse(100, Some(200), true);
+        parse(200, Some(200), true);
+        parse(201, Some(200), false);
+        parse(210, Some(200), false);
+    }
+
     #[test]
     fn test_write_headers_orig_case_empty_value() {
         let mut headers = HeaderMap::new();