rev_thread_spawn/revactor: fix TeeInput for short reads

Explicitly requested short reads may cause too much data to be
returned, which would be bad and potentially break the
application.  We need to ensure proper IO#readpartial-like
semantics in both of these models.

2 files changed, 18 insertions, 3 deletions

diff --git a/lib/rainbows/ev_thread_core.rb b/lib/rainbows/ev_thread_core.rb
index 784d30a..287b726 100644
--- a/lib/rainbows/ev_thread_core.rb
+++ b/lib/rainbows/ev_thread_core.rb
@@ -17,9 +17,19 @@ module Rainbows
     # we pass ourselves off as a Socket to Unicorn::TeeInput and this
     # is the only method Unicorn::TeeInput requires from the socket
     def readpartial(length, buf = "")
-      buf.replace(@state.pop)
+      length == 0 and return buf.replace("")
+      # try bufferred reads first
+      @tbuf && @tbuf.size > 0 and return buf.replace(@tbuf.read(length))
+
+      tmp = @state.pop
+      diff = tmp.size - length
+      if diff > 0
+        @tbuf ||= ::IO::Buffer.new
+        @tbuf.write(tmp[length, tmp.size])
+        tmp = tmp[0, length]
+      end
       resume
-      buf
+      buf.replace(tmp)
     end
 
     def app_spawn(input)
@@ -59,7 +69,7 @@ module Rainbows
         if 0 == @hp.content_length
           app_spawn(HttpRequest::NULL_IO) # common case
         else # nil or len > 0
-          @state = Queue.new
+          @state, @tbuf = Queue.new, nil
           app_spawn(nil)
         end
       when Queue
diff --git a/lib/rainbows/revactor/tee_input.rb b/lib/rainbows/revactor/tee_input.rb
index 92effb4..b8042ad 100644
--- a/lib/rainbows/revactor/tee_input.rb
+++ b/lib/rainbows/revactor/tee_input.rb
@@ -24,6 +24,11 @@ module Rainbows
           begin
             if parser.filter_body(dst, buf << socket.read).nil?
               @tmp.write(dst)
+              diff = dst.size - length
+              if diff > 0
+                dst.replace(dst[0,length])
+                @tmp.seek(-diff, IO::SEEK_CUR)
+              end
               return dst
             end
           rescue EOFError