Ignore/close response bodies for status codes that don't allow them. (#23)

Author: ioquatix

lib/protocol/rack/body.rb

@@ -5,13 +5,18 @@
 
 require_relative "body/streaming"
 require_relative "body/enumerable"
+require_relative "constants"
 require "protocol/http/body/completable"
 
 module Protocol
 	module Rack
 		module Body
 			CONTENT_LENGTH = "content-length"
 
+			def self.no_content?(status)
+				status == 204 or status == 205 or status == 304
+			end
+			
 			def self.wrap(env, status, headers, body, input = nil)
 				# In no circumstance do we want this header propagating out:
 				if length = headers.delete(CONTENT_LENGTH)
@@ -33,12 +38,29 @@ def self.wrap(env, status, headers, body, input = nil)
 					end
 				elsif body.respond_to?(:each)
 					body = Body::Enumerable.wrap(body, length)
-				else
+				elsif body
 					body = Body::Streaming.new(body, input)
+				else
+					Console.warn(self, "Rack response body was nil, ignoring!")
 				end
 
-				if response_finished = env[RACK_RESPONSE_FINISHED] and response_finished.any?
-					body = ::Protocol::HTTP::Body::Completable.new(body, completion_callback(response_finished, env, status, headers))
+				if body and no_content?(status)
+					unless body.empty?
+						Console.warn(self, "Rack response body was not empty, and status code indicates no content!", body: body, status: status)
+					end
+					
+					body.close
+					body = nil
+				end
+				
+				response_finished = env[RACK_RESPONSE_FINISHED]
+				
+				if response_finished&.any?
+					if body
+						body = ::Protocol::HTTP::Body::Completable.new(body, completion_callback(response_finished, env, status, headers))
+					else
+						completion_callback(response_finished, env, status, headers).call(nil)
+					end
 				end
 
 				return body

test/protocol/rack/body.rb

@@ -0,0 +1,81 @@
+# frozen_string_literal: true
+
+# Released under the MIT License.
+# Copyright, 2024, by Samuel Williams.
+
+require "protocol/rack/body"
+require "protocol/http/body/readable"
+require "console"
+
+describe Protocol::Rack::Body do
+	with "#no_content?" do
+		it "returns true for status codes that indicate no content" do
+			expect(subject.no_content?(204)).to be == true
+			expect(subject.no_content?(205)).to be == true
+			expect(subject.no_content?(304)).to be == true
+			expect(subject.no_content?(200)).to be == false
+		end
+	end
+	
+	with "#wrap" do
+		let(:env) { {} }
+		let(:headers) { {} }
+		
+		it "handles nil body" do
+			expect(Console).to receive(:warn).and_return(nil)
+			
+			result = subject.wrap(env, 200, headers, nil)
+			expect(result).to be_nil
+		end
+		
+		with "non-empty body and no-content status" do
+			let(:mock_body) do
+				Protocol::HTTP::Body::Buffered.new(["content"])
+			end
+			
+			[204, 205, 304].each do |status|
+				it "closes body and returns nil for status #{status}", unique: status do
+					expect(Console).to receive(:warn).and_return(nil)
+					expect(mock_body).to receive(:close)
+					
+					result = subject.wrap(env, status, headers, mock_body)
+					
+					expect(result).to be_nil
+				end
+			end
+		end
+		
+		with "empty body and no-content status" do
+			let(:mock_body) do
+				Protocol::HTTP::Body::Buffered.new
+			end
+			
+			it "closes body and returns nil for no-content status" do
+				expect(Console).not.to receive(:warn)
+				expect(mock_body).to receive(:close)
+				
+				result = subject.wrap(env, 204, headers, mock_body)
+				
+				expect(result).to be_nil
+			end
+		end
+		
+		with "body and normal status" do
+			let(:mock_body) do
+				body = Object.new
+				
+				def body.each
+					yield "content"
+				end
+				
+				body
+			end
+			
+			it "wraps body properly with status 200" do
+				result = subject.wrap(env, 200, headers, mock_body)
+				expect(result).to be_a(Protocol::Rack::Body::Enumerable)
+				expect(result).not.to be_nil
+			end
+		end
+	end
+end