Close a partially read request body once the response has been sent

lib/protocol/rack/adapter/generic.rb

@@ -10,6 +10,8 @@
 require_relative "../response"
 require_relative "../rewindable"
 
+require "protocol/http/body/completable"
+
 module Protocol
 	module Rack
 		module Adapter
@@ -177,6 +179,9 @@ def handle_error(env, status, headers, body, error)
 				def call(request)
 					env = self.make_environment(request)
 
+					# Keep the input we created, in case the application replaces it:
+					input = env[RACK_INPUT]
+
 					status, headers, body = @app.call(env)
 
 					# The status must always be an integer:
@@ -191,11 +196,32 @@ def call(request)
 
 					headers, meta = self.wrap_headers(headers)
 
-					return Response.wrap(env, status, headers, meta, body, request)
+					response = Response.wrap(env, status, headers, meta, body, request)
+
+					return close_partially_read_input(input, response)
 				rescue => error
+					close_partially_read_input(input)
+
 					return self.handle_error(env, status, headers, body, error)
 				end
 
+				# If the application read part of the request body, close the rest once the response has been sent, otherwise the server will wait for it. Input the application never read is left for the server to discard.
+				#
+				# @parameter input [Input | Nil] The input given to the application.
+				# @parameter response [Protocol::HTTP::Response | Nil] The response being sent, if any.
+				# @returns [Protocol::HTTP::Response | Nil] The response.
+				def close_partially_read_input(input, response = nil)
+					if input.is_a?(Input) and input.partially_read?
+						if response&.body
+							response.body = ::Protocol::HTTP::Body::Completable.new(response.body, proc{input.close})
+						else
+							input.close
+						end
+					end
+
+					return response
+				end
+
 				# Generate a suitable response for the given exception.
 				# 
 				# @parameter exception [Exception] The exception that occurred.

lib/protocol/rack/adapter/rack2.rb

@@ -87,6 +87,9 @@ def make_environment(request)
 				def call(request)
 					env = self.make_environment(request)
 
+					# Keep the input we created, in case the application replaces it:
+					input = env[RACK_INPUT]
+
 					status, headers, body = @app.call(env)
 
 					if status
@@ -110,8 +113,12 @@ def call(request)
 					# 	body = hijack_body
 					# end
 
-					return Response.wrap(env, status, headers, meta, body, request)
+					response = Response.wrap(env, status, headers, meta, body, request)
+
+					return close_partially_read_input(input, response)
 				rescue => error
+					close_partially_read_input(input)
+
 					return self.handle_error(env, status, headers, body, error)
 				end
 

lib/protocol/rack/input.rb

@@ -21,6 +21,7 @@ class Input
 			def initialize(body, ...)
 				@body = body
 				@closed = false
+				@reading = false
 
 				super(...)
 			end
@@ -70,6 +71,11 @@ def rewind
 				return false
 			end
 
+			# Whether the application has read part of the input, but not all of it.
+			def partially_read?
+				@reading and !closed?
+			end
+
 			# Whether the stream has been closed.
 			def closed?
 				@closed or @body.nil?
@@ -88,6 +94,8 @@ def flush
 
 			def sysread(size, buffer)
 				if @body
+					@reading = true
+
 					# User's may forget to call #close...
 					if chunk = @body.read
 						# If the user reads exactly the content length, we close the stream automatically:

test/protocol/rack/adapter/rack2.rb

@@ -116,6 +116,63 @@
 				end
 			end
 		end
+
+		with "a request body" do
+			let(:body) {Protocol::HTTP::Body::Buffered.new(["Hello", " ", "World"])}
+			let(:request) {Protocol::HTTP::Request.new("https", "example.com", "POST", "/", "http/1.1", Protocol::HTTP::Headers.new, body)}
+
+			with "that is partially read" do
+				let(:app) do
+					->(env) do
+						@input = env["rack.input"]
+						@input.read(5)
+
+						[400, {}, ["Rejected"]]
+					end
+				end
+
+				it "closes the input once the response has been sent" do
+					expect(response.status).to be == 400
+					expect(@input).not.to be(:closed?)
+
+					expect(response.read).to be == "Rejected"
+
+					expect(@input).to be(:closed?)
+				end
+			end
+
+			with "that is partially read when the application fails" do
+				let(:app) do
+					->(env) do
+						@input = env["rack.input"]
+						@input.read(5)
+
+						raise StandardError, "Invalid upload"
+					end
+				end
+
+				it "closes the input" do
+					expect(response.status).to be == 500
+					expect(@input).to be(:closed?)
+				end
+			end
+
+			with "that is not read" do
+				let(:app) do
+					->(env) do
+						@input = env["rack.input"]
+
+						[200, {}, ["Ignored"]]
+					end
+				end
+
+				it "leaves the input for the server" do
+					response.close
+
+					expect(@input).not.to be(:closed?)
+				end
+			end
+		end
 	end
 
 	with "#make_response" do

test/protocol/rack/adapter/rack3.rb

@@ -73,6 +73,81 @@
 				expect(response.read).to be == "Hello World!"
 			end
 		end
+
+		with "a request body" do
+			let(:body) {Protocol::HTTP::Body::Buffered.new(["Hello", " ", "World"])}
+			let(:request) {Protocol::HTTP::Request.new("https", "example.com", "POST", "/", "http/1.1", Protocol::HTTP::Headers.new, body)}
+
+			with "that is partially read" do
+				let(:app) do
+					->(env) do
+						@input = env["rack.input"]
+						@input.read(5)
+
+						[400, {}, ["Rejected"]]
+					end
+				end
+
+				it "closes the input once the response has been sent" do
+					expect(response.status).to be == 400
+					expect(@input).not.to be(:closed?)
+
+					expect(response.read).to be == "Rejected"
+
+					expect(@input).to be(:closed?)
+				end
+			end
+
+			with "that is partially read and no response body" do
+				let(:app) do
+					->(env) do
+						@input = env["rack.input"]
+						@input.read(5)
+
+						[204, {}, []]
+					end
+				end
+
+				it "closes the input immediately" do
+					expect(response.body).to be_nil
+					expect(@input).to be(:closed?)
+				end
+			end
+
+			with "that is partially read when the application fails" do
+				let(:app) do
+					->(env) do
+						@input = env["rack.input"]
+						@input.read(5)
+
+						raise StandardError, "Invalid upload"
+					end
+				end
+
+				it "closes the input" do
+					expect(response.status).to be == 500
+					expect(@input).to be(:closed?)
+				end
+			end
+
+			with "that is not read" do
+				let(:app) do
+					->(env) do
+						@input = env["rack.input"]
+
+						[200, {}, ["Ignored"]]
+					end
+				end
+
+				it "leaves the input for the server" do
+					expect(response.body).to be_a(Protocol::HTTP::Body::Buffered)
+
+					response.close
+
+					expect(@input).not.to be(:closed?)
+				end
+			end
+		end
 	end
 
 	with "#make_response" do

test/protocol/rack/input.rb

@@ -146,6 +146,31 @@
 			end
 		end
 
+		with "#partially_read?" do
+			it "is not partially read before it is read" do
+				expect(input).not.to be(:partially_read?)
+			end
+
+			it "is partially read after reading part of it" do
+				input.read(3)
+
+				expect(input).to be(:partially_read?)
+			end
+
+			it "is not partially read after reading all of it" do
+				input.read
+
+				expect(input).not.to be(:partially_read?)
+			end
+
+			it "is not partially read once closed" do
+				input.read(3)
+				input.close
+
+				expect(input).not.to be(:partially_read?)
+			end
+		end
+
 		with "#rewind" do
 			it "can rewind after reading all input" do
 				expect(input.read).to be == sample_data.join