summary refs log tree commit homepage
diff options
context:
space:
mode:
authorEric Wong <normalperson@yhbt.net>2010-10-05 22:09:20 +0000
committerEric Wong <normalperson@yhbt.net>2010-10-05 22:09:20 +0000
commitfd6b47cf1690cb45f2144cd92e0fe1f301c7c37b (patch)
tree37c7df590221318380a086269a8aa4c4a20670d4
parent350e8fa3a94838bcc936782315b3472615fe6517 (diff)
TeeInput methods may be invoked deep in the stack, so
avoid giving them more work to do if a client disconnects
due to a bad upload.
-rw-r--r--lib/unicorn/tee_input.rb35
-rw-r--r--test/unit/test_request.rb4
-rw-r--r--test/unit/test_tee_input.rb2
3 files changed, 17 insertions, 24 deletions
diff --git a/lib/unicorn/tee_input.rb b/lib/unicorn/tee_input.rb
index 0dbfff6..32ee4f2 100644
--- a/lib/unicorn/tee_input.rb
+++ b/lib/unicorn/tee_input.rb
@@ -171,44 +171,25 @@ class Unicorn::TeeInput < Struct.new(:socket, :req, :parser,
 
 private
 
-  def client_error(e)
-    case e
-    when 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 Unicorn::ClientShutdown, "bytes_read=#{tmp.size}", []
-    when Unicorn::HttpParserError
-      e.set_backtrace([])
-    end
-    raise e
-  end
-
   # tees off a +length+ chunk of data from the input into the IO
   # backing store as well as returning it.  +dst+ must be specified.
   # returns nil if reading from the input returns nil
   def tee(length, dst)
     unless parser.body_eof?
-      if parser.filter_body(dst, socket.readpartial(length, buf)).nil?
+      r = socket.kgio_read(length, buf) or eof!
+      unless parser.filter_body(dst, r)
         tmp.write(dst)
         tmp.seek(0, IO::SEEK_END) # workaround FreeBSD/OSX + MRI 1.8.x bug
         return dst
       end
     end
     finalize_input
-    rescue => e
-      client_error(e)
   end
 
   def finalize_input
     while parser.trailers(req, buf).nil?
-      # Don't worry about raising ClientShutdown 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(@@io_chunk_size)
+      r = socket.kgio_read(@@io_chunk_size) or eof!
+      buf << r
     end
     self.socket = nil
   end
@@ -232,4 +213,12 @@ private
     dst
   end
 
+  def eof!
+    # 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 Unicorn::ClientShutdown, "bytes_read=#{tmp.size}", []
+  end
 end
diff --git a/test/unit/test_request.rb b/test/unit/test_request.rb
index 8aae344..5ac511f 100644
--- a/test/unit/test_request.rb
+++ b/test/unit/test_request.rb
@@ -163,6 +163,10 @@ class RequestTest < Test::Unit::TestCase
     length = bs * count
     client = Tempfile.new('big_put')
     def client.kgio_addr; '127.0.0.1'; end
+    def client.kgio_read(*args)
+      readpartial(*args)
+    rescue EOFError
+    end
     client.syswrite(
       "PUT / HTTP/1.1\r\n" \
       "Host: foo\r\n" \
diff --git a/test/unit/test_tee_input.rb b/test/unit/test_tee_input.rb
index 56b8952..263aa8a 100644
--- a/test/unit/test_tee_input.rb
+++ b/test/unit/test_tee_input.rb
@@ -10,7 +10,7 @@ class TestTeeInput < Test::Unit::TestCase
   def setup
     @rs = $/
     @env = {}
-    @rd, @wr = UNIXSocket.pair
+    @rd, @wr = Kgio::UNIXSocket.pair
     @rd.sync = @wr.sync = true
     @start_pid = $$
   end