From 5bc239fd154a7eaebeb024394f8e0b507bbf4c5a Mon Sep 17 00:00:00 2001
From: Eric Wong <normalperson@yhbt.net>
Date: Fri, 19 Nov 2010 20:51:57 +0000
Subject: stream_input: small cleanups and fixes

No need to accept any number of args, that could hide bugs in
applications that could give three or more arguments.  We also
raise ArgumentError when given a negative length argument to
read.
---
 lib/unicorn/stream_input.rb    | 37 +++++++++++++------------------------
 test/unit/test_stream_input.rb |  7 +++++++
 2 files changed, 20 insertions(+), 24 deletions(-)

diff --git a/lib/unicorn/stream_input.rb b/lib/unicorn/stream_input.rb
index ef8997e..45d1336 100644
--- a/lib/unicorn/stream_input.rb
+++ b/lib/unicorn/stream_input.rb
@@ -39,28 +39,25 @@ class Unicorn::StreamInput
   # ios.read(length [, buffer]) will return immediately if there is
   # any data and only block when nothing is available (providing
   # IO#readpartial semantics).
-  def read(*args)
-    length = args.shift
-    rv = args.shift || ''
-    if length.nil?
-      read_all(rv)
-    else
+  def read(length = nil, rv = '')
+    if length
       if length <= @rbuf.size
-        rv.replace(@rbuf.slice(0, length))
-        @rbuf.replace(@rbuf.slice(length, @rbuf.size) || '')
+        length < 0 and raise ArgumentError, "negative length #{length} given"
+        rv.replace(@rbuf.slice!(0, length))
       else
-        rv.replace(@rbuf)
-        length -= @rbuf.size
-        @rbuf.replace('')
-        until length == 0 || eof? || (rv.size > 0 && @chunked)
-          @socket.kgio_read(length, @buf) or eof!
+        to_read = length - @rbuf.size
+        rv.replace(@rbuf.slice!(0, @rbuf.size))
+        until to_read == 0 || eof? || (rv.size > 0 && @chunked)
+          @socket.kgio_read(to_read, @buf) or eof!
           filter_body(@rbuf, @buf)
           rv << @rbuf
-          length -= @rbuf.size
-          @rbuf.replace('')
+          to_read -= @rbuf.size
         end
+        @rbuf.replace('')
       end
       rv = nil if rv.empty? && length != 0
+    else
+      read_all(rv)
     end
     rv
   end
@@ -84,15 +81,7 @@ class Unicorn::StreamInput
 
     begin
       @rbuf.gsub!(re, '') and return $1
-      if eof?
-        if @rbuf.empty?
-          return nil
-        else
-          rv = @rbuf.dup
-          @rbuf.replace('')
-          return rv
-        end
-      end
+      return @rbuf.empty? ? nil : @rbuf.slice!(0, @rbuf.size) if eof?
       @socket.kgio_read(@@io_chunk_size, @buf) or eof!
       filter_body(once = '', @buf)
       @rbuf << once
diff --git a/test/unit/test_stream_input.rb b/test/unit/test_stream_input.rb
index adf4571..0c4ec98 100644
--- a/test/unit/test_stream_input.rb
+++ b/test/unit/test_stream_input.rb
@@ -21,6 +21,13 @@ class TestStreamInput < Test::Unit::TestCase
     Process.waitall
   end
 
+  def test_read_negative
+    r = init_request('hello')
+    si = Unicorn::StreamInput.new(@rd, r)
+    assert_raises(ArgumentError) { si.read(-1) }
+    assert_equal 'hello', si.read
+  end
+
   def test_read_small
     r = init_request('hello')
     si = Unicorn::StreamInput.new(@rd, r)
-- 
cgit v1.2.3-24-ge0c7