Merge pull request #2287 from eugeneius/chunked_content_length

Set CONTENT_LENGTH for chunked requests
puma · May 31, 2020 · f9ddd58 · f9ddd58
2 parents 22002ac + bdcfb60
commit f9ddd58
Show file tree

Hide file tree

Showing 3 changed files with 60 additions and 7 deletions.
diff --git a/History.md b/History.md
@@ -46,6 +46,7 @@
   * Improvements to `out_of_band` hook (#2234)
   * Prefer the rackup file specified by the CLI (#2225)
   * Fix for spawning subprocesses with fork_worker option (#2267)
+  * Set `CONTENT_LENGTH` for chunked requests (#2287)
 
 * Refactor
   * Remove unused loader argument from Plugin initializer (#2095)

diff --git a/lib/puma/client.rb b/lib/puma/client.rb
@@ -420,7 +420,10 @@ def read_chunked_body
           raise EOFError
         end
 
-        return true if decode_chunk(chunk)
+        if decode_chunk(chunk)
+          @env[CONTENT_LENGTH] = @chunked_content_length
+          return true
+        end
       end
     end
 
@@ -432,20 +435,28 @@ def setup_chunked_body(body)
       @body = Tempfile.new(Const::PUMA_TMP_BASE)
       @body.binmode
       @tempfile = @body
+      @chunked_content_length = 0
+
+      if decode_chunk(body)
+        @env[CONTENT_LENGTH] = @chunked_content_length
+        return true
+      end
+    end
 
-      return decode_chunk(body)
+    def write_chunk(str)
+      @chunked_content_length += @body.write(str)
     end
 
     def decode_chunk(chunk)
       if @partial_part_left > 0
         if @partial_part_left <= chunk.size
           if @partial_part_left > 2
-            @body << chunk[0..(@partial_part_left-3)] # skip the \r\n
+            write_chunk(chunk[0..(@partial_part_left-3)]) # skip the \r\n
           end
           chunk = chunk[@partial_part_left..-1]
           @partial_part_left = 0
         else
-          @body << chunk if @partial_part_left > 2 # don't include the last \r\n
+          write_chunk(chunk) if @partial_part_left > 2 # don't include the last \r\n
           @partial_part_left -= chunk.size
           return false
         end
@@ -492,12 +503,12 @@ def decode_chunk(chunk)
 
           case
           when got == len
-            @body << part[0..-3] # to skip the ending \r\n
+            write_chunk(part[0..-3]) # to skip the ending \r\n
           when got <= len - 2
-            @body << part
+            write_chunk(part)
             @partial_part_left = len - part.size
           when got == len - 1 # edge where we get just \r but not \n
-            @body << part[0..-2]
+            write_chunk(part[0..-2])
             @partial_part_left = len - part.size
           end
         else

diff --git a/test/test_puma_server.rb b/test/test_puma_server.rb
@@ -498,21 +498,26 @@ def test_Expect_100
 
   def test_chunked_request
     body = nil
+    content_length = nil
     server_run app: ->(env) {
       body = env['rack.input'].read
+      content_length = env['CONTENT_LENGTH']
       [200, {}, [""]]
     }
 
     data = send_http_and_read "GET / HTTP/1.1\r\nConnection: close\r\nTransfer-Encoding: chunked\r\n\r\n1\r\nh\r\n4\r\nello\r\n0\r\n\r\n"
 
     assert_equal "HTTP/1.1 200 OK\r\nConnection: close\r\nContent-Length: 0\r\n\r\n", data
     assert_equal "hello", body
+    assert_equal 5, content_length
   end
 
   def test_chunked_request_pause_before_value
     body = nil
+    content_length = nil
     server_run app: ->(env) {
       body = env['rack.input'].read
+      content_length = env['CONTENT_LENGTH']
       [200, {}, [""]]
     }
 
@@ -525,12 +530,15 @@ def test_chunked_request_pause_before_value
 
     assert_equal "HTTP/1.1 200 OK\r\nConnection: close\r\nContent-Length: 0\r\n\r\n", data
     assert_equal "hello", body
+    assert_equal 5, content_length
   end
 
   def test_chunked_request_pause_between_chunks
     body = nil
+    content_length = nil
     server_run app: ->(env) {
       body = env['rack.input'].read
+      content_length = env['CONTENT_LENGTH']
       [200, {}, [""]]
     }
 
@@ -543,12 +551,15 @@ def test_chunked_request_pause_between_chunks
 
     assert_equal "HTTP/1.1 200 OK\r\nConnection: close\r\nContent-Length: 0\r\n\r\n", data
     assert_equal "hello", body
+    assert_equal 5, content_length
   end
 
   def test_chunked_request_pause_mid_count
     body = nil
+    content_length = nil
     server_run app: ->(env) {
       body = env['rack.input'].read
+      content_length = env['CONTENT_LENGTH']
       [200, {}, [""]]
     }
 
@@ -561,12 +572,15 @@ def test_chunked_request_pause_mid_count
 
     assert_equal "HTTP/1.1 200 OK\r\nConnection: close\r\nContent-Length: 0\r\n\r\n", data
     assert_equal "hello", body
+    assert_equal 5, content_length
   end
 
   def test_chunked_request_pause_before_count_newline
     body = nil
+    content_length = nil
     server_run app: ->(env) {
       body = env['rack.input'].read
+      content_length = env['CONTENT_LENGTH']
       [200, {}, [""]]
     }
 
@@ -579,12 +593,15 @@ def test_chunked_request_pause_before_count_newline
 
     assert_equal "HTTP/1.1 200 OK\r\nConnection: close\r\nContent-Length: 0\r\n\r\n", data
     assert_equal "hello", body
+    assert_equal 5, content_length
   end
 
   def test_chunked_request_pause_mid_value
     body = nil
+    content_length = nil
     server_run app: ->(env) {
       body = env['rack.input'].read
+      content_length = env['CONTENT_LENGTH']
       [200, {}, [""]]
     }
 
@@ -597,12 +614,15 @@ def test_chunked_request_pause_mid_value
 
     assert_equal "HTTP/1.1 200 OK\r\nConnection: close\r\nContent-Length: 0\r\n\r\n", data
     assert_equal "hello", body
+    assert_equal 5, content_length
   end
 
   def test_chunked_request_pause_between_cr_lf_after_size_of_second_chunk
     body = nil
+    content_length = nil
     server_run app: ->(env)  {
       body = env['rack.input'].read
+      content_length = env['CONTENT_LENGTH']
       [200, {}, [""]]
     }
 
@@ -624,12 +644,15 @@ def test_chunked_request_pause_between_cr_lf_after_size_of_second_chunk
 
     assert_equal "HTTP/1.1 200 OK\r\nConnection: close\r\nContent-Length: 0\r\n\r\n", data
     assert_equal (part1 + 'b'), body
+    assert_equal 4201, content_length
   end
 
   def test_chunked_request_pause_between_closing_cr_lf
     body = nil
+    content_length = nil
     server_run app: ->(env) {
       body = env['rack.input'].read
+      content_length = env['CONTENT_LENGTH']
       [200, {}, [""]]
     }
 
@@ -643,12 +666,15 @@ def test_chunked_request_pause_between_closing_cr_lf
 
     assert_equal "HTTP/1.1 200 OK\r\nConnection: close\r\nContent-Length: 0\r\n\r\n", data
     assert_equal 'hello', body
+    assert_equal 5, content_length
   end
 
   def test_chunked_request_pause_before_closing_cr_lf
     body = nil
+    content_length = nil
     server_run app: ->(env) {
       body = env['rack.input'].read
+      content_length = env['CONTENT_LENGTH']
       [200, {}, [""]]
     }
 
@@ -662,25 +688,31 @@ def test_chunked_request_pause_before_closing_cr_lf
 
     assert_equal "HTTP/1.1 200 OK\r\nConnection: close\r\nContent-Length: 0\r\n\r\n", data
     assert_equal 'hello', body
+    assert_equal 5, content_length
   end
 
   def test_chunked_request_header_case
     body = nil
+    content_length = nil
     server_run app: ->(env) {
       body = env['rack.input'].read
+      content_length = env['CONTENT_LENGTH']
       [200, {}, [""]]
     }
 
     data = send_http_and_read "GET / HTTP/1.1\r\nConnection: close\r\nTransfer-Encoding: Chunked\r\n\r\n1\r\nh\r\n4\r\nello\r\n0\r\n\r\n"
 
     assert_equal "HTTP/1.1 200 OK\r\nConnection: close\r\nContent-Length: 0\r\n\r\n", data
     assert_equal "hello", body
+    assert_equal 5, content_length
   end
 
   def test_chunked_keep_alive
     body = nil
+    content_length = nil
     server_run app: ->(env) {
       body = env['rack.input'].read
+      content_length = env['CONTENT_LENGTH']
       [200, {}, [""]]
     }
 
@@ -690,14 +722,17 @@ def test_chunked_keep_alive
 
     assert_equal ["HTTP/1.1 200 OK", "Content-Length: 0"], h
     assert_equal "hello", body
+    assert_equal 5, content_length
 
     sock.close
   end
 
   def test_chunked_keep_alive_two_back_to_back
     body = nil
+    content_length = nil
     server_run app: ->(env) {
       body = env['rack.input'].read
+      content_length = env['CONTENT_LENGTH']
       [200, {}, [""]]
     }
 
@@ -715,6 +750,7 @@ def test_chunked_keep_alive_two_back_to_back
     h = header(sock)
     assert_equal ["HTTP/1.1 200 OK", "Content-Length: 0"], h
     assert_equal "hello", body
+    assert_equal 5, content_length
     assert_equal true, last_crlf_written
 
     last_crlf_writer.join
@@ -726,16 +762,19 @@ def test_chunked_keep_alive_two_back_to_back
 
     assert_equal ["HTTP/1.1 200 OK", "Content-Length: 0"], h
     assert_equal "goodbye", body
+    assert_equal 7, content_length
 
     sock.close
   end
 
   def test_chunked_keep_alive_two_back_to_back_with_set_remote_address
     body = nil
+    content_length = nil
     remote_addr =nil
     @server = Puma::Server.new @app, @events, { remote_address: :header, remote_address_header: 'HTTP_X_FORWARDED_FOR'}
     server_run app: ->(env) {
       body = env['rack.input'].read
+      content_length = env['CONTENT_LENGTH']
       remote_addr = env['REMOTE_ADDR']
       [200, {}, [""]]
     }
@@ -745,6 +784,7 @@ def test_chunked_keep_alive_two_back_to_back_with_set_remote_address
     h = header sock
     assert_equal ["HTTP/1.1 200 OK", "Content-Length: 0"], h
     assert_equal "hello", body
+    assert_equal 5, content_length
     assert_equal "127.0.0.1", remote_addr
 
     sock << "GET / HTTP/1.1\r\nX-Forwarded-For: 127.0.0.2\r\nConnection: Keep-Alive\r\nTransfer-Encoding: chunked\r\n\r\n4\r\ngood\r\n3\r\nbye\r\n0\r\n\r\n"
@@ -754,6 +794,7 @@ def test_chunked_keep_alive_two_back_to_back_with_set_remote_address
 
     assert_equal ["HTTP/1.1 200 OK", "Content-Length: 0"], h
     assert_equal "goodbye", body
+    assert_equal 7, content_length
     assert_equal "127.0.0.2", remote_addr
 
     sock.close