Validate streaming responses from sent chunks

lib/committee/middleware/response_validation.rb

@@ -18,13 +18,25 @@ def handle(request)
         streaming_content_parser = retrieve_streaming_content_parser(headers)
 
         if streaming_content_parser
+          original_response = response
+          streamed_response = []
+          response = Enumerator.new do |yielder|
+            original_response.each do |chunk|
+              streamed_response << chunk
+              yielder << chunk
+            end
+          end
           response = Rack::BodyProxy.new(response) do
             begin
-              validate(request, status, headers, response, streaming_content_parser)
-            rescue => e
-              handle_exception(e, request.env)
+              original_response.close if original_response.respond_to?(:close)
+            ensure
+              begin
+                validate(request, status, headers, streamed_response, streaming_content_parser)
+              rescue => e
+                handle_exception(e, request.env)
 
-              raise e if @raise
+                raise e if @raise
+              end
             end
           end
         else

test/middleware/response_validation_test.rb

@@ -212,6 +212,31 @@ def app
         get "/events/stream"
         assert_equal 200, last_response.status
       end
+
+      it 'validates the chunks sent from a body that cannot be read after close' do
+        closed = false
+        body = Object.new
+        body.define_singleton_method(:each) do |&block|
+          raise IOError, "body is closed" if closed
+
+          ["hello"].each(&block)
+        end
+        body.define_singleton_method(:close) { closed = true }
+        validated_body = nil
+        options = { schema: open_api_3_streaming_response_schema, streaming_content_parsers: { 'text/event-stream' => ->(body) { validated_body = body } }, raise: true, }
+        @app = Rack::Builder.new {
+          use Committee::Middleware::ResponseValidation, options
+          run ->(_) { [200, { 'content-type' => 'text/event-stream' }, body] }
+        }
+
+        status, _headers, response_body = @app.call(Rack::MockRequest.env_for("/events/stream"))
+
+        assert_equal 200, status
+        assert_equal ["hello"], response_body.each.to_a
+        response_body.close
+        assert closed
+        assert_equal "hello", validated_body
+      end
     end
 
     describe 'application/x-json-stream; customized streaming event' do