From e4256da292f9626d7dfca60e08f65651a0a9139a Mon Sep 17 00:00:00 2001
From: Eric Wong <normalperson@yhbt.net>
Date: Sat, 14 Nov 2009 00:23:19 +0000
Subject: raise Unicorn::ClientShutdown if client aborts in TeeInput

Leaving the EOFError exception as-is bad because most
applications/frameworks run an application-wide exception
handler to pretty-print and/or log the exception with a huge
backtrace.

Since there's absolutely nothing we can do in the server-side
app to deal with clients prematurely shutting down, having a
backtrace does not make sense.  Having a backtrace can even be
harmful since it creates unnecessary noise for application
engineers monitoring or tracking down real bugs.
---
 lib/unicorn.rb           |  6 +++++
 lib/unicorn/tee_input.rb | 11 +++++++++
 test/unit/test_server.rb | 61 ++++++++++++++++++++++++++++++++++++++++++++++--
 3 files changed, 76 insertions(+), 2 deletions(-)

diff --git a/lib/unicorn.rb b/lib/unicorn.rb
index a696402..c6c311e 100644
--- a/lib/unicorn.rb
+++ b/lib/unicorn.rb
@@ -8,6 +8,12 @@ autoload :Rack, 'rack'
 # a Unicorn web server.  It contains a minimalist HTTP server with just enough
 # functionality to service web application requests fast as possible.
 module Unicorn
+
+  # raise this inside TeeInput when a client disconnects inside the
+  # application dispatch
+  class ClientShutdown < EOFError
+  end
+
   autoload :Const, 'unicorn/const'
   autoload :HttpRequest, 'unicorn/http_request'
   autoload :HttpResponse, 'unicorn/http_response'
diff --git a/lib/unicorn/tee_input.rb b/lib/unicorn/tee_input.rb
index 69397c0..50ddb5b 100644
--- a/lib/unicorn/tee_input.rb
+++ b/lib/unicorn/tee_input.rb
@@ -135,10 +135,21 @@ module Unicorn
         end
       end
       finalize_input
+      rescue EOFError
+        # in case client only did a premature shutdown(SHUT_WR)
+        # we do support clients that shutdown(SHUT_WR) after the
+        # _entire_ request has been sent, and those will not have
+        # raised EOFError on us.
+        socket.close if socket
+        raise ClientShutdown, "bytes_read=#{@tmp.size}", []
     end
 
     def finalize_input
       while parser.trailers(req, buf).nil?
+        # Don't worry about throw-ing :http_499 here on EOFError, tee()
+        # will catch EOFError when app is processing it, otherwise in
+        # initialize we never get any chance to enter the app so the
+        # EOFError will just get trapped by Unicorn and not the Rack app
         buf << socket.readpartial(Const::CHUNK_SIZE)
       end
       self.socket = nil
diff --git a/test/unit/test_server.rb b/test/unit/test_server.rb
index bbb06da..a7f6a35 100644
--- a/test/unit/test_server.rb
+++ b/test/unit/test_server.rb
@@ -12,11 +12,12 @@ include Unicorn
 
 class TestHandler 
 
-  def call(env) 
-  #   response.socket.write("HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\n\r\nhello!\n")
+  def call(env)
     while env['rack.input'].read(4096)
     end
     [200, { 'Content-Type' => 'text/plain' }, ['hello!\n']]
+    rescue Unicorn::ClientShutdown => e
+      $stderr.syswrite("#{e.class}: #{e.message} #{e.backtrace.empty?}\n")
    end
 end
 
@@ -103,6 +104,62 @@ class WebServerTest < Test::Unit::TestCase
     assert_equal 'hello!\n', results[0], "Handler didn't really run"
   end
 
+  def test_client_shutdown_writes
+    sock = nil
+    buf = nil
+    bs = 15609315 * rand
+    assert_nothing_raised do
+      sock = TCPSocket.new('127.0.0.1', @port)
+      sock.syswrite("PUT /hello HTTP/1.1\r\n")
+      sock.syswrite("Host: example.com\r\n")
+      sock.syswrite("Transfer-Encoding: chunked\r\n")
+      sock.syswrite("Trailer: X-Foo\r\n")
+      sock.syswrite("\r\n")
+      sock.syswrite("%x\r\n" % [ bs ])
+      sock.syswrite("F" * bs)
+      sock.syswrite("\r\n0\r\nX-")
+      "Foo: bar\r\n\r\n".each_byte do |x|
+        sock.syswrite x.chr
+        sleep 0.05
+      end
+      # we wrote the entire request before shutting down, server should
+      # continue to process our request and never hit EOFError on our sock
+      sock.shutdown(Socket::SHUT_WR)
+      buf = sock.read
+    end
+    assert_equal 'hello!\n', buf.split(/\r\n\r\n/).last
+    lines = File.readlines("test_stderr.#$$.log")
+    assert lines.grep(/^Unicorn::ClientShutdown: /).empty?
+    assert_nothing_raised { sock.close }
+  end
+
+  def test_client_shutdown_write_truncates
+    sock = nil
+    buf = nil
+    bs = 15609315 * rand
+    assert_nothing_raised do
+      sock = TCPSocket.new('127.0.0.1', @port)
+      sock.syswrite("PUT /hello HTTP/1.1\r\n")
+      sock.syswrite("Host: example.com\r\n")
+      sock.syswrite("Transfer-Encoding: chunked\r\n")
+      sock.syswrite("Trailer: X-Foo\r\n")
+      sock.syswrite("\r\n")
+      sock.syswrite("%x\r\n" % [ bs ])
+      sock.syswrite("F" * (bs / 2.0))
+
+      # shutdown prematurely, this will force the server to abort
+      # processing on us even during app dispatch
+      sock.shutdown(Socket::SHUT_WR)
+      IO.select([sock], nil, nil, 60) or raise "Timed out"
+      buf = sock.read
+    end
+    assert_equal "", buf
+    lines = File.readlines("test_stderr.#$$.log")
+    lines = lines.grep(/^Unicorn::ClientShutdown: bytes_read=\d+/)
+    assert_equal 1, lines.size
+    assert_match %r{\AUnicorn::ClientShutdown: bytes_read=\d+ true$}, lines[0]
+    assert_nothing_raised { sock.close }
+  end
 
   def do_test(string, chunk, close_after=nil, shutdown_delay=0)
     # Do not use instance variables here, because it needs to be thread safe
-- 
cgit v1.2.3-24-ge0c7