client/legacy_client_session: set transfer-encoding: chunked if content-length is not specified

2 files changed, 84 insertions, 11 deletions

diff --git a/lib/plum/client/legacy_client_session.rb b/lib/plum/client/legacy_client_session.rb
index 2e52be2..087c643 100644
--- a/lib/plum/client/legacy_client_session.rb
+++ b/lib/plum/client/legacy_client_session.rb
@@ -30,8 +30,18 @@ module Plum
     end
 
     def request(headers, body, options, &headers_cb)
+      headers["host"] = headers[":authority"] || headers["host"] || @config[:hostname]
+      if body
+        if headers["content-length"] || headers["transfer-encoding"]
+          chunked = false
+        else
+          chunked = true
+          headers["transfer-encoding"] = "chunked"
+        end
+      end
+
       response = Response.new
-      @requests << [response, headers, body, headers_cb]
+      @requests << [response, headers, body, chunked, headers_cb]
       consume_queue
       response
     end
@@ -45,20 +55,40 @@ module Plum
     def consume_queue
       return if @response || @requests.empty?
 
-      response, headers, body, cb = @requests.shift
-      headers["host"] = headers[":authority"] || headers["host"] || @config[:hostname]
+      response, headers, body, chunked, cb = @requests.shift
       @response = response
       @headers_callback = cb
 
-      @socket << "%s %s HTTP/1.1\r\n" % [headers[":method"], headers[":path"]]
+      @socket << construct_request(headers)
+
+      if body
+        if chunked
+          read_object(body) { |chunk|
+            @socket << chunk.bytesize.to_s(16) << "\r\n" << chunk << "\r\n"
+          }
+        else
+          read_object(body) { |chunk| @socket << chunk }
+        end
+      end
+    end
+
+    def construct_request(headers)
+      out = String.new
+      out << "%s %s HTTP/1.1\r\n" % [headers[":method"], headers[":path"]]
       headers.each { |key, value|
         next if key.start_with?(":") # HTTP/2 psuedo headers
-        @socket << "%s: %s\r\n" % [key, value]
+        out << "%s: %s\r\n" % [key, value]
       }
-      @socket << "\r\n"
+      out << "\r\n"
+    end
 
-      if body
-        @socket << body
+    def read_object(body)
+      if body.is_a?(String)
+        yield body
+      else # IO
+        until body.eof?
+          yield body.readpartial(1024)
+        end
       end
     end
 
diff --git a/test/plum/client/test_legacy_client_session.rb b/test/plum/client/test_legacy_client_session.rb
index 2c8f146..421a261 100644
--- a/test/plum/client/test_legacy_client_session.rb
+++ b/test/plum/client/test_legacy_client_session.rb
@@ -37,11 +37,54 @@ class LegacyClientSessionTest < Minitest::Test
     io = StringIO.new
     session = LegacyClientSession.new(io, Client::DEFAULT_CONFIG.merge(hostname: "aa"))
     res = session.request({ ":method" => "GET", ":path" => "/aa" }, "aa", {})
-    assert("GET /aa HTTP/1.1\r\nhost: aa\r\n\r\naa")
+    assert_equal("GET /aa HTTP/1.1\r\nhost: aa\r\ntransfer-encoding: chunked\r\n\r\n2\r\naa\r\n", io.string)
     io.string << "HTTP/1.1 200 OK\r\nContent-Length: 3\r\n\r\naaa"
     session.succ until res.finished?
     assert(res.finished?)
-    assert("aaa", res.body)
-    assert({ ":status" => "200", "content-length" => "3" }, res.headers)
+    assert_equal("aaa", res.body)
+    assert_equal({ ":status" => "200", "content-length" => "3" }, res.headers)
+  end
+
+  def test_chunked_chunked_string
+    io = StringIO.new
+    session = LegacyClientSession.new(io, Client::DEFAULT_CONFIG.merge(hostname: "hostname"))
+    res = session.request({ ":method" => "GET", ":path" => "/aa" }, "a" * 1025, {})
+    assert_equal(<<-EOR, io.string)
+GET /aa HTTP/1.1\r
+host: hostname\r
+transfer-encoding: chunked\r
+\r
+401\r
+#{"a"*1025}\r
+    EOR
+  end
+
+  def test_chunked_chunked_io
+    io = StringIO.new
+    session = LegacyClientSession.new(io, Client::DEFAULT_CONFIG.merge(hostname: "hostname"))
+    res = session.request({ ":method" => "GET", ":path" => "/aa" }, StringIO.new("a" * 1025), {})
+    assert_equal(<<-EOR, io.string)
+GET /aa HTTP/1.1\r
+host: hostname\r
+transfer-encoding: chunked\r
+\r
+400\r
+#{"a"*1024}\r
+1\r
+a\r
+    EOR
+  end
+
+  def test_chunked_sized
+    io = StringIO.new
+    session = LegacyClientSession.new(io, Client::DEFAULT_CONFIG.merge(hostname: "hostname"))
+    res = session.request({ ":method" => "GET", ":path" => "/aa", "content-length" => 1025 }, StringIO.new("a" * 1025), {})
+    assert_equal((<<-EOR).chomp, io.string)
+GET /aa HTTP/1.1\r
+content-length: 1025\r
+host: hostname\r
+\r
+#{"a"*1025}
+    EOR
   end
 end