add Kcar::Response#rack to give proper Rack responses

We can't unchunk bodies and still say
"Transfer-Encoding: chunked", nor can we say we have Trailers
without chunked encoding.  So we'll unconditionally disable
unchunking and pass the response through as-is

2 files changed, 97 insertions, 7 deletions

diff --git a/lib/kcar/response.rb b/lib/kcar/response.rb
index ef93f12..38d35ac 100644
--- a/lib/kcar/response.rb
+++ b/lib/kcar/response.rb
@@ -4,7 +4,7 @@ module Kcar
 
 # This may be used to generate a Rack response
 #
-class Response < Struct.new(:sock, :unchunk, :hdr, :buf, :parser)
+class Response < Struct.new(:sock, :hdr, :unchunk, :buf, :parser)
 
   # :stopdoc:
   LAST_CHUNK = "0\r\n"
@@ -17,15 +17,21 @@ class Response < Struct.new(:sock, :unchunk, :hdr, :buf, :parser)
   # initializes a socket, +sock+ must respond to the "readpartial"
   # method.  +unchunk+ may be set to disable transparent unchunking
   # +hdr+ may be a Hash, Array, or Rack::Utils::HeaderHash
-  def initialize(sock, unchunk = true, hdr = {})
-    super(sock, unchunk, hdr, "", Parser.new)
+  def initialize(sock, hdr = {}, unchunk = true)
+    super(sock, hdr, unchunk, "", Parser.new)
   end
 
-  # returns a 3-element array suitable for use as a Rack response:
+  # returns a 3-element array that resembles a Rack response, but is
+  # more useful for additional processing by other code.
+  #
   #    [ status, headers, body ]
   #
+  # Use Kcar::Response#rack if you want to guaranteee a proper Rack response.
+  #
   # this method will not return until the response headers are fully parsed,
   # but the body returned will be this Kcar::Response handler itself.
+  # +unchunk+ must be true to guarantee trailers will be stored in the
+  # returned +header+ object
   def read
     buf << sock.readpartial(READ_SIZE) if buf.empty?
     while (response = parser.headers(hdr, buf)).nil?
@@ -34,6 +40,17 @@ class Response < Struct.new(:sock, :unchunk, :hdr, :buf, :parser)
     response << self
   end
 
+  # returns a 3-element array suitable for use as a Rack response:
+  #    [ status, headers, body ]
+  #
+  # this method will not return until the response headers are fully parsed,
+  # but the body returned will be this Kcar::Response handler itself.
+  # It is not guaranteed that trailers will be stored in the returned +header+
+  def rack
+    self.unchunk = false
+    read
+  end
+
   # this is expected to be called by our Rack server, it will close
   # our given +sock+ object if keepalive is not used otherwise it
   # will just reset the parser and clear the header object
diff --git a/test/test_response.rb b/test/test_response.rb
index 54a3bf3..ab1c30c 100644
--- a/test/test_response.rb
+++ b/test/test_response.rb
@@ -163,7 +163,7 @@ class TestSession < Test::Unit::TestCase
       @s << "HTTP/1.1 200 OK\r\nTrailer: Foo\r\n"
       @s << "Transfer-Encoding: chunked\r\n\r\n"
     end
-    @response = Kcar::Response.new(@c, false)
+    @response = Kcar::Response.new(@c, {}, false)
     status, headers, body = @response.read
     assert_equal status, "200 OK"
     expect = {
@@ -191,7 +191,7 @@ class TestSession < Test::Unit::TestCase
       @s << "HTTP/1.0 200 OK\r\n"
       @s << "Content-Type: text/plain\r\n\r\n"
     end
-    @response = Kcar::Response.new(@c, false)
+    @response = Kcar::Response.new(@c, {}, false)
     status, headers, body = @response.read
     assert_equal status, "200 OK"
     expect = { "Content-Type" => "text/plain", }
@@ -261,7 +261,7 @@ class TestSession < Test::Unit::TestCase
       @s << "HTTP/1.1 200 OK\r\nTrailer: Foo\r\n"
       @s << "Transfer-Encoding: chunked\r\n\r\n"
     end
-    @response = Kcar::Response.new(@c, false)
+    @response = Kcar::Response.new(@c, {}, false)
     status, headers, body = @response.read
     assert_equal status, "200 OK"
     expect = {
@@ -339,4 +339,77 @@ class TestSession < Test::Unit::TestCase
     assert ! @c.closed?
   end
 
+  def test_rack_preserve_chunk_hash
+    pid = fork do
+      @s << "HTTP/1.1 200 OK\r\n"
+      @s << "Trailer: Foo\r\n"
+      @s << "Transfer-Encoding: chunked\r\n\r\n"
+      @s << "5\r\nhello\r\n0\r\nFoo: bar\r\n\r\n"
+    end
+    @response = Kcar::Response.new(@c)
+    status, headers, body = @response.rack
+    assert_equal status, "200 OK"
+    expect = {
+      "Trailer" => "Foo",
+      "Transfer-Encoding" => "chunked",
+    }
+    assert_equal expect, headers
+    tmp = []
+    assert body.parser.keepalive?
+    assert_nothing_raised { body.each { |chunk| tmp << chunk.dup } }
+    assert_equal ["5\r\n", "hello\r\n", "0\r\n", "Foo: bar\r\n\r\n"], tmp
+    expect["Foo"] = "bar"
+    assert_equal expect, headers
+    _, status = Process.waitpid2(pid)
+    assert status.success?
+    body.close
+    assert ! @c.closed?
+  end
+
+  def test_rack_preserve_chunk_ary
+    pid = fork do
+      @s << "HTTP/1.1 200 OK\r\n"
+      @s << "Trailer: Foo\r\n"
+      @s << "Transfer-Encoding: chunked\r\n\r\n"
+      @s << "5\r\nhello\r\n0\r\nFoo: bar\r\n\r\n"
+    end
+    @response = Kcar::Response.new(@c, [])
+    status, headers, body = @response.rack
+    assert_equal status, "200 OK"
+    expect = [ %w(Trailer Foo), %w(Transfer-Encoding chunked) ]
+    assert_equal expect, headers
+    tmp = []
+    assert body.parser.keepalive?
+    assert_nothing_raised { body.each { |chunk| tmp << chunk.dup } }
+    assert_equal ["5\r\n", "hello\r\n", "0\r\n", "Foo: bar\r\n\r\n"], tmp
+    expect << %w(Foo bar)
+    assert_equal expect, headers
+    _, status = Process.waitpid2(pid)
+    assert status.success?
+    body.close
+    assert ! @c.closed?
+  end
+
+  def test_rack_preserve_chunk_ary_no_keepalive
+    pid = fork do
+      @s << "HTTP/1.1 200 OK\r\n"
+      @s << "Connection: close\r\n"
+      @s << "Trailer: Foo\r\n"
+      @s << "Transfer-Encoding: chunked\r\n\r\n"
+      @s << "5\r\nhello\r\n0\r\nFoo: bar\r\n\r\n"
+    end
+    @s.close
+    @response = Kcar::Response.new(@c, [])
+    status, headers, body = @response.rack
+    assert_equal status, "200 OK"
+    tmp = []
+    assert ! body.parser.keepalive?
+    assert_nothing_raised { body.each { |chunk| tmp << chunk.dup } }
+    assert_equal ["5\r\nhello\r\n0\r\nFoo: bar\r\n\r\n"], tmp
+    _, status = Process.waitpid2(pid)
+    assert status.success?
+    body.close
+    assert @c.closed?
+  end
+
 end