Fix head request handling. (#29)

Author: samuel-williams-shopify

lib/protocol/rack/body.rb

@@ -6,7 +6,9 @@
 require_relative "body/streaming"
 require_relative "body/enumerable"
 require_relative "constants"
+
 require "protocol/http/body/completable"
+require "protocol/http/body/head"
 
 module Protocol
 	module Rack
@@ -37,8 +39,9 @@ def self.no_content?(status)
 			# @parameter headers [Hash] The response headers.
 			# @parameter body [Object] The response body to wrap.
 			# @parameter input [Object] Optional input for streaming bodies.
+			# @parameter head [Boolean] Indicates if this is a HEAD request, which should not have a body.
 			# @returns [Protocol::HTTP::Body] The wrapped response body.
-			def self.wrap(env, status, headers, body, input = nil)
+			def self.wrap(env, status, headers, body, input = nil, head = false)
 				# In no circumstance do we want this header propagating out:
 				if length = headers.delete(CONTENT_LENGTH)
 					# We don't really trust the user to provide the right length to the transport.
@@ -84,6 +87,19 @@ def self.wrap(env, status, headers, body, input = nil)
 					end
 				end
 
+				# There are two main situations we need to handle:
+				# 1. The application has the `Rack::Head` middleware in the stack, which means we should not return a body, and the application is also responsible for setting the content-length header. `Rack::Head` will result in an empty enumerable body.
+				# 2. The application does not have `Rack::Head`, in which case it will return a body and we need to extract the length.
+				# In both cases, we need to ensure that the body is wrapped correctly. If there is no body and we don't know the length, we also just return `nil`.
+				if head
+					if body
+						body = ::Protocol::HTTP::Body::Head.for(body)
+					elsif length
+						body = ::Protocol::HTTP::Body::Head.new(length)
+					end
+					# Otherwise, body is `nil` and we don't know the length either.
+				end
+				
 				return body
 			end
 

lib/protocol/rack/response.rb

@@ -50,12 +50,7 @@ def self.wrap(env, status, headers, meta, body, request = nil)
 					body = hijack_body
 				end
 
-				body = Body.wrap(env, status, headers, body, request&.body)
-
-				if request&.head?
-					# I thought about doing this in Output.wrap, but decided the semantics are too tricky. Specifically, the various ways a rack response body can be wrapped, and the need to invoke #close at the right point.
-					body = ::Protocol::HTTP::Body::Head.for(body)
-				end
+				body = Body.wrap(env, status, headers, body, request&.body, request&.head?)
 
 				protocol = meta[RACK_PROTOCOL]
 

test/protocol/rack/body.rb

@@ -18,8 +18,8 @@
 	end
 
 	with "#wrap" do
-		let(:env) { {} }
-		let(:headers) { {} }
+		let(:env) {Hash.new}
+		let(:headers) {Hash.new}
 
 		it "handles nil body" do
 			expect(Console).to receive(:warn).and_return(nil)
@@ -28,6 +28,42 @@
 			expect(result).to be_nil
 		end
 
+		with "head request" do
+			it "handles head request with content-length and empty body" do
+				headers["content-length"] = "123"
+				
+				result = subject.wrap(env, 200, headers, [], nil, true)
+				
+				expect(result).to be_a(Protocol::HTTP::Body::Head)
+				expect(result.length).to be == 123
+			end
+			
+			it "handles head request with no content-length and empty body" do
+				result = subject.wrap(env, 200, headers, [], nil, true)
+				
+				expect(result).to be_a(Protocol::HTTP::Body::Head)
+				expect(result.length).to be == 0
+			end
+			
+			it "handles head request with content-length and nil body" do
+				headers["content-length"] = "123"
+				
+				expect(Console).to receive(:warn).and_return(nil)
+				result = subject.wrap(env, 200, headers, nil, nil, true)
+				
+				expect(result).to be_a(Protocol::HTTP::Body::Head)
+				expect(result.length).to be == 123
+			end
+			
+			it "handles head request with no content-length and nil body" do
+				expect(Console).to receive(:warn).and_return(nil)
+				
+				result = subject.wrap(env, 200, headers, nil, nil, true)
+				
+				expect(result).to be_nil
+			end
+		end
+		
 		with "non-empty body and no-content status" do
 			let(:mock_body) do
 				Protocol::HTTP::Body::Buffered.new(["content"])