From b2283d8f2505f67de2918c792b0beaf0448e697a Mon Sep 17 00:00:00 2001
From: Blake Williams <blakewilliams@github.com>
Date: Sun, 17 Apr 2022 12:21:36 -0400
Subject: [PATCH] Fix rack.after_reply exceptions breaking connections (#2861)

* Fix rack.after_reply exceptions breaking connections

Currently, when a `rack.after_reply` callable raises an exception we
attempt to handle it like other client errors by responding with an HTTP
500 response. This however doesn't work because `rack.after_reply`
callbacks are called after the response body has already been written to
the client.

This can cause issues with re-used connections. This is because 2 HTTP
responses are being returned for a single request. If a second HTTP
request is made before the error handling logic completes the timing can
line up causing the second HTTP response to be served a 500 from the
first HTTP requests `rack.after_reply` callbacks raising.

That may look roughly like:

1. Request 1 starts, opening a reusable TCP connection
2. Request 1 is written to and "completed"
3. Request 1 `rack.after_reply` callables are called
4. Request 2 starts, reusing the same TCP connection as request 1
5. `rack.after_reply` raises, calls `client_error` and serves a 500
   response
6. Request 2 receives the 500 response.

This is somewhat difficult to reproduce using HTTP clients since it's a
race condition whether or not the 500 is written at the "correct" time
or not.

To prevent this issue the `rack.after_reply` callables are now wrapped
in a begin/rescue/end block that rescues from `StandardError` and logs
instead of attempting to serve a 500 response.

* Assert against less specific exception
---
 lib/puma/request.rb      |  6 ++++-
 test/test_rack_server.rb | 50 ++++++++++++++++++++++++++++++++++++++--
 2 files changed, 53 insertions(+), 3 deletions(-)

diff --git a/lib/puma/request.rb b/lib/puma/request.rb
index 03f531bb..b4bdbb89 100644
--- a/lib/puma/request.rb
+++ b/lib/puma/request.rb
@@ -182,7 +182,11 @@ module Puma
           res_body.close if res_body.respond_to? :close
         end
 
-        after_reply.each { |o| o.call }
+        begin
+          after_reply.each { |o| o.call }
+        rescue StandardError => e
+          @log_writer.debug_error e
+        end
       end
 
       res_info[:keep_alive]
diff --git a/test/test_rack_server.rb b/test/test_rack_server.rb
index 0d2e5a37..377ac24c 100644
--- a/test/test_rack_server.rb
+++ b/test/test_rack_server.rb
@@ -36,8 +36,8 @@ class TestRackServer < Minitest::Test
   def setup
     @simple = lambda { |env| [200, { "X-Header" => "Works" }, ["Hello"]] }
     @server = Puma::Server.new @simple
-    port = (@server.add_tcp_listener "127.0.0.1", 0).addr[1]
-    @tcp = "http://127.0.0.1:#{port}"
+    @port = (@server.add_tcp_listener "127.0.0.1", 0).addr[1]
+    @tcp = "http://127.0.0.1:#{@port}"
     @stopped = false
   end
 
@@ -108,6 +108,52 @@ class TestRackServer < Minitest::Test
     assert_equal true, closed
   end
 
+  def test_after_reply_exception
+    @server.app = lambda do |env|
+      env['rack.after_reply'] << lambda { raise ArgumentError, "oops" }
+      @simple.call(env)
+    end
+
+    @server.run
+
+    socket = TCPSocket.open "127.0.0.1", @port
+    socket.puts "GET /test HTTP/1.1\r\n"
+    socket.puts "Connection: Keep-Alive\r\n"
+    socket.puts "\r\n"
+
+    headers = socket.readline("\r\n\r\n")
+      .split("\r\n")
+      .drop(1)
+      .map { |line| line.split(/:\s?/) }
+      .to_h
+
+    content_length = headers["Content-Length"].to_i
+    real_response_body = socket.read(content_length)
+
+    assert_equal "Hello", real_response_body
+
+    # When after_reply breaks the connection it will write the expected HTTP
+    # response followed by a second HTTP response: HTTP/1.1 500
+    #
+    # This sleeps to give the server time to write the invalid/extra HTTP
+    # response.
+    #
+    # * If we can read from the socket, we know that extra content has been
+    #   written to the connection and assert that it's our erroneous 500
+    #   response.
+    # * If we would block trying to read from the socket, we can assume that
+    #   the erroneous 500 response wasn't/won't be written.
+    sleep 0.1
+    assert_raises IO::WaitReadable do
+      content = socket.read_nonblock(12)
+      refute_includes content, "500"
+    end
+
+    socket.close
+
+    stop
+  end
+
   def test_common_logger
     log = StringIO.new