From 335e3387616329b5f3c873950e192d8e9ccfe192 Mon Sep 17 00:00:00 2001
From: Samuel Williams <samuel.williams@oriontransfer.co.nz>
Date: Thu, 13 Mar 2025 13:30:51 +1300
Subject: [PATCH] Add maximum_header_count limit.

---
 lib/protocol/http1/connection.rb | 8 +++++++-
 lib/protocol/http1/error.rb      | 3 +++
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/lib/protocol/http1/connection.rb b/lib/protocol/http1/connection.rb
index ca81618..aa0ac95 100644
--- a/lib/protocol/http1/connection.rb
+++ b/lib/protocol/http1/connection.rb
@@ -46,13 +46,14 @@ module HTTP1
 		VALID_FIELD_VALUE = /\A#{FIELD_VALUE}\z/.freeze
 		
 		DEFAULT_MAXIMUM_LINE_LENGTH = 8192
+		MAXIMUM_HEADER_COUNT = 128
 		
 		class Connection
 			CRLF = "\r\n"
 			HTTP10 = "HTTP/1.0"
 			HTTP11 = "HTTP/1.1"
 			
-			def initialize(stream, persistent: true, state: :idle, maximum_line_length: DEFAULT_MAXIMUM_LINE_LENGTH)
+			def initialize(stream, persistent: true, state: :idle, maximum_line_length: DEFAULT_MAXIMUM_LINE_LENGTH, maximum_header_count: MAXIMUM_HEADER_COUNT)
 				@stream = stream
 				
 				@persistent = persistent
@@ -61,6 +62,7 @@ def initialize(stream, persistent: true, state: :idle, maximum_line_length: DEFA
 				@count = 0
 				
 				@maximum_line_length = maximum_line_length
+				@maximum_header_count = maximum_header_count
 			end
 			
 			attr :stream
@@ -381,6 +383,10 @@ def read_headers
 				fields = []
 				
 				while line = read_line
+					if @maximum_header_count and fields.size > @maximum_header_count
+						raise HeaderCountError, "Too many headers: #{fields.size} > #{@maximum_header_count}!"
+					end
+					
 					# Empty line indicates end of headers:
 					break if line.empty?
 					
diff --git a/lib/protocol/http1/error.rb b/lib/protocol/http1/error.rb
index dd2c015..12595f1 100644
--- a/lib/protocol/http1/error.rb
+++ b/lib/protocol/http1/error.rb
@@ -17,6 +17,9 @@ class ProtocolError < Error
 		class LineLengthError < Error
 		end
 		
+		class HeaderCountError < Error
+		end
+		
 		# The request was not able to be parsed correctly, or failed some kind of validation.
 		class BadRequest < Error
 		end