Validate trailers on add.

Author: ioquatix

lib/protocol/http/header/generic.rb

@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+# Released under the MIT License.
+# Copyright, 2025, by Samuel Williams.
+
+require_relative "split"
+
+module Protocol
+	module HTTP
+		module Header
+			# Represents generic or custom headers that can be used in trailers.
+			#
+			# This class is used as the default policy for headers not explicitly defined in the POLICY hash.
+			#
+			# It allows generic headers to be used in HTTP trailers, which is important for:
+			# - Custom application headers.
+			# - gRPC status headers (grpc-status, grpc-message).
+			# - Headers used by proxies and middleware.
+			# - Future HTTP extensions.
+			class Generic < Split
+				# Whether this header is acceptable in HTTP trailers.
+				# Generic headers are allowed in trailers by default to support extensibility.
+				# @returns [Boolean] `true`, generic headers are allowed in trailers.
+				def self.trailer?
+					true
+				end
+			end
+		end
+	end
+end

lib/protocol/http/headers.rb

@@ -20,6 +20,7 @@
 require_relative "header/trailer"
 require_relative "header/server_timing"
 require_relative "header/digest"
+require_relative "header/generic"
 
 require_relative "header/accept"
 require_relative "header/accept_charset"
@@ -158,7 +159,26 @@ def trailer!(&block)
 				return trailer(&block)
 			end
 
+			# Enumerate all the headers in the header, if there are any.
+			# 
+			# @yields {|key, value| ...} The header key and value.
+			# 	@parameter key [String] The header key.
+			# 	@parameter value [String] The raw header value.
+			def header(&block)
+				return to_enum(:header) unless block_given?
+				
+				if @tail and @tail < @fields.size
+					@fields.first(@tail).each(&block)
+				else
+					@fields.each(&block)
+				end
+			end
+			
 			# Enumerate all headers in the trailer, if there are any.
+			# 
+			# @yields {|key, value| ...} The header key and value.
+			# 	@parameter key [String] The header key.
+			# 	@parameter value [String] The raw header value.
 			def trailer(&block)
 				return to_enum(:trailer) unless block_given?
 
@@ -191,7 +211,7 @@ def empty?
 			# 	@parameter key [String] The header key.
 			# 	@parameter value [String] The raw header value.
 			def each(&block)
-				self.to_h.each(&block)
+				@fields.each(&block)
 			end
 
 			# @returns [Boolean] Whether the headers include the specified key.
@@ -227,9 +247,18 @@ def extract(keys)
 			#
 			# @parameter key [String] the header key.
 			# @parameter value [String] the header value to assign.
-			def add(key, value)
+			# @parameter trailer [Boolean] whether this header is being added as a trailer.
+			def add(key, value, trailer: (@tail != nil))
 				value = value.to_s
 
+				if trailer
+					policy = @policy[key.downcase]
+					
+					if !policy or !policy.trailer?
+						raise InvalidTrailerError, key
+					end
+				end
+				
 				if @indexed
 					merge_into(@indexed, key.downcase, value)
 				end
@@ -297,6 +326,10 @@ def merge(headers)
 			end
 
 			# The policy for various headers, including how they are merged and normalized.
+			#
+			# A policy may be `false` to indicate that the header may only be specified once and is a simple string.
+			#
+			# Otherwise, the policy is a class which implements the header normalization logic, including `parse` and `coerce` class methods.
 			POLICY = {
 				# Headers which may only be specified once:
 				"content-disposition" => false,
@@ -315,8 +348,11 @@ def merge(headers)
 				"user-agent" => false,
 				"trailer" => Header::Trailer,
 
-				# Custom headers:
+				# Connection handling:
 				"connection" => Header::Connection,
+				"upgrade" => false,
+				
+				# Cache handling:
 				"cache-control" => Header::CacheControl,
 				"te" => Header::TE,
 				"vary" => Header::Vary,
@@ -354,16 +390,21 @@ def merge(headers)
 
 				# Accept headers:
 				"accept" => Header::Accept,
+				"accept-ranges" => Header::Split,
 				"accept-charset" => Header::AcceptCharset,
 				"accept-encoding" => Header::AcceptEncoding,
 				"accept-language" => Header::AcceptLanguage,
 
+				# Content negotiation headers:
+				"content-encoding" => Header::Split,
+				"content-range" => false,
+				
 				# Performance headers:
 				"server-timing" => Header::ServerTiming,
 
 				# Content integrity headers:
 				"digest" => Header::Digest,
-			}.tap{|hash| hash.default = Split}
+			}.tap{|hash| hash.default = Header::Generic}
 
 			# Delete all header values for the given key, and return the merged value.
 			#
@@ -403,25 +444,14 @@ def delete(key)
 			# @parameter hash [Hash] The hash to merge into.
 			# @parameter key [String] The header key.
 			# @parameter value [String] The raw header value.
-			# @parameter trailer [Boolean] Whether this header is in the trailer section.
-			protected def merge_into(hash, key, value, trailer = @tail)
+			protected def merge_into(hash, key, value)
 				if policy = @policy[key]
-					# Check if we're adding to trailers and this header is allowed:
-					if trailer && !policy.trailer?
-						raise InvalidTrailerError, key
-					end
-					
 					if current_value = hash[key]
 						current_value << value
 					else
 						hash[key] = policy.parse(value)
 					end
 				else
-					# By default, headers are not allowed in trailers:
-					if trailer
-						raise InvalidTrailerError, key
-					end
-					
 					if hash.key?(key)
 						raise DuplicateHeaderError, key
 					end
@@ -437,13 +467,14 @@ def delete(key)
 			# @returns [Hash] A hash table of `{key, value}` pairs.
 			def to_h
 				unless @indexed
-					@indexed = {}
+					indexed = {}
 
-					@fields.each_with_index do |(key, value), index|
-						trailer = (@tail && index >= @tail)
-						
-						merge_into(@indexed, key.downcase, value, trailer)
+					@fields.each do |key, value|
+						merge_into(indexed, key.downcase, value)
 					end
+					
+					# Deferred assignment so that exceptions in `merge_into` don't leave us in an inconsistent state:
+					@indexed = indexed
 				end
 
 				return @indexed

releases.md

@@ -1,5 +1,10 @@
 # Releases
 
+## Unreleased
+
+  - Move trailer validation to `Headers#add` method to ensure all additions are checked at the time of addition as this is a hard requirement.
+  - Introduce `Headers#header` method to enumerate only the main headers, excluding trailers. This can be used after invoking `Headers#trailer!` to avoid race conditions.
+
 ## v0.57.0
 
   - Always use `#parse` when parsing header values from strings to ensure proper normalization and validation.

test/protocol/http/headers.rb

@@ -14,7 +14,7 @@
 			["Set-Cookie", "hello=world"],
 			["Accept", "*/*"],
 			["set-cookie", "foo=bar"],
-		]
+		].freeze
 	end
 
 	let(:headers) {subject[fields]}
@@ -293,15 +293,15 @@
 
 			headers.add("etag", "abcd")
 
-			expect(trailer.to_h).to be == {"etag" => "abcd"}
+			expect(trailer.to_a).to be == [["etag", "abcd"]]
 		end
 
 		it "can add trailer without explicit header" do
 			trailer = headers.trailer!
 
 			headers.add("etag", "abcd")
 
-			expect(trailer.to_h).to be == {"etag" => "abcd"}
+			expect(trailer.to_a).to be == [["etag", "abcd"]]
 		end
 
 		with "forbidden trailers" do
@@ -335,15 +335,42 @@
 
 				cookie
 				set-cookie
-				
-				x-foo-bar
 			]
 
 			forbidden_trailers.each do |key|
-				it "can't add a #{key.inspect} header in the trailer", unique: key do
-					trailer = headers.trailer!
-					headers.add(key, "example")
-					expect{headers.to_h}.to raise_exception(Protocol::HTTP::InvalidTrailerError)
+				with "forbidden trailer #{key.inspect}", unique: key do
+					it "can't add a #{key.inspect} header in the trailer" do
+						trailer = headers.trailer!
+						expect do
+							headers.add(key, "example", trailer: true)
+						end.to raise_exception(Protocol::HTTP::InvalidTrailerError)
+					end
+					
+					it "can't add a #{key.inspect} header with trailer: true" do
+						expect do
+							headers.add(key, "example", trailer: true)
+						end.to raise_exception(Protocol::HTTP::InvalidTrailerError)
+					end
+				end
+			end
+		end
+		
+		with "unknown trailers", unique: "unknown" do
+			let(:headers) {subject.new}
+			
+			unknown_trailers = %w[
+				x-foo-bar
+				grpc-status
+				grpc-message
+				x-custom-header
+			]
+			
+			unknown_trailers.each do |key|
+				with "unknown trailer #{key.inspect}", unique: key do
+					it "can add unknown header #{key.inspect} as trailer" do
+						headers.add(key, "example", trailer: true)
+						expect(headers).to be(:include?, key)
+					end
 				end
 			end
 		end
@@ -359,22 +386,110 @@
 			]
 
 			permitted_trailers.each do |key|
-				it "can add a #{key.inspect} header in the trailer", unique: key do
-					trailer = headers.trailer!
-					headers.add(key, "example")
-					expect(headers).to be(:include?, key)
+				with "permitted trailer #{key.inspect}", unique: key do
+					it "can add a #{key.inspect} header in the trailer" do
+						trailer = headers.trailer!
+						headers.add(key, "example")
+						expect(headers).to be(:include?, key)
+					end
+					
+					it "can add a #{key.inspect} header with trailer: true" do
+						headers.add(key, "example", trailer: true)
+						expect(headers).to be(:include?, key)
+					end
 				end
 			end
 		end
 	end
 
+	with "#header" do
+		it "can enumerate all headers when there are no trailers" do
+			result = headers.header.to_a
+			
+			expect(result).to be == fields
+		end
+		
+		it "enumerates headers but not trailers" do
+			headers.trailer!
+			headers.add("etag", "abcd")
+			headers.add("digest", "sha-256=xyz")
+			
+			header = headers.header.to_a
+			
+			# Should only include the original 5 fields, not the 2 trailers
+			expect(header.size).to be == 5
+			expect(header).to be == fields
+		end
+		
+		it "returns an enumerator when no block is given" do
+			enumerator = headers.header
+			
+			expect(enumerator).to be_a(Enumerator)
+			expect(enumerator.to_a).to be == fields
+		end
+		
+		it "returns an enumerator that excludes trailers" do
+			headers.trailer!
+			headers.add("etag", "abcd")
+			
+			enumerator = headers.header
+			
+			expect(enumerator).to be_a(Enumerator)
+			expect(enumerator.to_a.size).to be == 5
+			expect(enumerator.to_a).to be == [
+				["Content-Type", "text/html"],
+				["connection", "Keep-Alive"],
+				["Set-Cookie", "hello=world"],
+				["Accept", "*/*"],
+				["set-cookie", "foo=bar"]
+			]
+		end
+	end
+	
 	with "#trailer" do
 		it "can enumerate trailer" do
 			headers.add("trailer", "etag")
 			headers.trailer!
 			headers.add("etag", "abcd")
 
-			expect(headers.trailer.to_h).to be == {"etag" => "abcd"}
+			expect(headers.trailer.to_a).to be == [["etag", "abcd"]]
+		end
+	end
+	
+	with "#add with trailer: keyword" do
+		let(:headers) {subject.new}
+		
+		it "allows adding regular headers without trailer: true" do
+			headers.add("content-type", "text/plain")
+			expect(headers["content-type"]).to be == "text/plain"
+		end
+		
+		it "validates trailers immediately when trailer: true" do
+			expect do
+				headers.add("content-type", "text/plain", trailer: true)
+			end.to raise_exception(Protocol::HTTP::InvalidTrailerError)
+		end
+		
+		it "allows permitted trailers with trailer: true" do
+			headers.add("etag", "abcd", trailer: true)
+			expect(headers["etag"]).to be == "abcd"
+		end
+		
+		it "validates trailers without calling trailer! first" do
+			# This should fail immediately, without needing trailer! to be called
+			expect do
+				headers.add("authorization", "Bearer token", trailer: true)
+			end.to raise_exception(Protocol::HTTP::InvalidTrailerError)
+		end
+		
+		it "validates trailers even when headers are not indexed" do
+			# Add without triggering indexing
+			expect do
+				headers.add("host", "example.com", trailer: true)
+			end.to raise_exception(Protocol::HTTP::InvalidTrailerError)
+			
+			# Ensure we haven't triggered indexing yet
+			expect(headers.instance_variable_get(:@indexed)).to be_nil
 		end
 	end