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.

4 files changed, 69 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
diff --git a/t/sha1-random-size.ru b/t/sha1-random-size.ru
new file mode 100644
index 0000000..f86d017
--- /dev/null
+++ b/t/sha1-random-size.ru
@@ -0,0 +1,19 @@
+# SHA1 checksum generator
+require 'digest/sha1'
+use Rack::ContentLength
+cap = 16384
+app = lambda do |env|
+  /\A100-continue\z/i =~ env['HTTP_EXPECT'] and
+    return [ 100, {}, [] ]
+  digest = Digest::SHA1.new
+  input = env['rack.input']
+  if buf = input.read(rand(cap))
+    begin
+      raise "#{buf.size} > #{cap}" if buf.size > cap
+      digest.update(buf)
+    end while input.read(rand(cap), buf)
+  end
+
+  [ 200, {'Content-Type' => 'text/plain'}, [ digest.hexdigest << "\n" ] ]
+end
+run app
diff --git a/t/t0102-rack-input-short.sh b/t/t0102-rack-input-short.sh
new file mode 100644
index 0000000..97c079d
--- /dev/null
+++ b/t/t0102-rack-input-short.sh
@@ -0,0 +1,32 @@
+#!/bin/sh
+. ./test-lib.sh
+test -r random_blob || die "random_blob required, run with 'make $0'"
+
+t_plan 4 "rack.input short read tests"
+
+t_begin "setup and startup" && {
+        rtmpfiles curl_out curl_err
+        rainbows_setup $model
+        rainbows -D sha1-random-size.ru -c $unicorn_config
+        blob_sha1=$(rsha1 random_blob)
+        t_info "blob_sha1=$blob_sha1"
+        rainbows_wait_start
+}
+
+t_begin "regular request" && {
+        curl -sSf -T random_blob http://$listen/ > $curl_out 2> $curl_err
+        test x$blob_sha1 = x$(cat $curl_out)
+        test ! -s $curl_err
+}
+
+t_begin "chunked request" && {
+        curl -sSf -T- < random_blob http://$listen/ > $curl_out 2> $curl_err
+        test x$blob_sha1 = x$(cat $curl_out)
+        test ! -s $curl_err
+}
+
+t_begin "shutdown" && {
+        kill $rainbows_pid
+}
+
+t_done