Add persistence and chunked transfer support

lib/puma/server.rb

@@ -124,31 +124,32 @@ module Puma
 
     def process_client(client)
       begin
-        parser = HttpParser.new
-        env = @proto_env.dup
-        data = client.readpartial(CHUNK_SIZE)
-        nparsed = 0
+        while true
+          parser = HttpParser.new
+          env = @proto_env.dup
+          data = client.readpartial(CHUNK_SIZE)
+          nparsed = 0
 
-        # Assumption: nparsed will always be less since data will get filled
-        # with more after each parsing.  If it doesn't get more then there was
-        # a problem with the read operation on the client socket. 
-        # Effect is to stop processing when the socket can't fill the buffer
-        # for further parsing.
-        while nparsed < data.length
-          nparsed = parser.execute(env, data, nparsed)
+          # Assumption: nparsed will always be less since data will get filled
+          # with more after each parsing.  If it doesn't get more then there was
+          # a problem with the read operation on the client socket. 
+          # Effect is to stop processing when the socket can't fill the buffer
+          # for further parsing.
+          while nparsed < data.length
+            nparsed = parser.execute(env, data, nparsed)
 
-          if parser.finished?
-            handle_request env, client, parser.body
-            break
-          else
-            # Parser is not done, queue up more data to read and continue parsing
-            chunk = client.readpartial(CHUNK_SIZE)
-            break if !chunk or chunk.length == 0  # read failed, stop processing
+            if parser.finished?
+              return unless handle_request env, client, parser.body
+            else
+              # Parser is not done, queue up more data to read and continue parsing
+              chunk = client.readpartial(CHUNK_SIZE)
+              return if !chunk or chunk.length == 0  # read failed, stop processing
 
-            data << chunk
-            if data.length >= MAX_HEADER
-              raise HttpParserError,
-                "HEADER is longer than allowed, aborting client early."
+              data << chunk
+              if data.length >= MAX_HEADER
+                raise HttpParserError,
+                  "HEADER is longer than allowed, aborting client early."
+              end
             end
           end
         end
@@ -207,11 +208,14 @@ module Puma
 
       body = read_body env, client, body
 
-      return unless body
+      return false unless body
 
       env["rack.input"] = body
       env["rack.url_scheme"] =  env["HTTPS"] ? "https" : "http"
 
+      keep_alive = false
+      chunked = false
+
       begin
         begin
           status, headers, res_body = @app.call(env)
@@ -219,11 +223,26 @@ module Puma
           status, headers, res_body = lowlevel_error(e)
         end
 
+        content_length = nil
+
+        if res_body.kind_of? String
+          content_length = res_body.size
+        elsif res_body.kind_of? Array and res_body.size == 1
+          content_length = res_body[0].size
+        end
+
         client.write "HTTP/1.1 "
         client.write status.to_s
         client.write " "
         client.write HTTP_STATUS_CODES[status]
-        client.write "\r\nConnection: close\r\n"
+
+        if content_length
+          client.write "\r\nContent-Length: #{content_length}\r\n"
+          keep_alive = true
+        else
+          client.write "\r\nTransfer-Encoding: chunked\r\n"
+          chunked = true
+        end
 
         colon = ": "
         line_ending = "\r\n"
@@ -240,18 +259,43 @@ module Puma
         client.write line_ending
 
         if res_body.kind_of? String
-          client.write body
+          if chunked
+            client.write res_body.size.to_s
+            client.write line_ending
+            client.write res_body
+            client.write line_ending
+          else
+            client.write res_body
+          end
+
           client.flush
         else
           res_body.each do |part|
-            client.write part
+            if chunked
+              client.write part.size.to_s
+              client.write line_ending
+              client.write part
+              client.write line_ending
+            else
+              client.write part
+            end
+
             client.flush
           end
         end
+
+        if chunked
+          client.write "0"
+          client.write line_ending
+          client.flush
+        end
+
       ensure
         body.close
         res_body.close if res_body.respond_to? :close
       end
+
+      return keep_alive
     end
 
     def read_body(env, client, body)
@@ -328,7 +372,7 @@ module Puma
 
       @ios.each do |io|
         if io.kind_of? TCPServer
-          fixed_name = name.gsub /\./, "-"
+          fixed_name = name.gsub(/\./, "-")
 
           DNSSD.announce io, "puma - #{fixed_name}", "http" do |r|
             @bonjour_registered = true

test/puma.conf

@@ -1 +0,0 @@
-uri "/fromconf", :handler => Puma::Error404Handler.new("test")

test/test_persistent.rb

@@ -0,0 +1,58 @@
+require 'puma'
+require 'test/unit'
+
+class TestPersistent < Test::Unit::TestCase
+  def setup
+    @valid_request = "GET / HTTP/1.1\r\nHost: test.com\r\nContent-Type: text/plain\r\n\r\n"
+    @headers = { "X-Header" => "Works" }
+    @body = ["Hello"]
+    @simple = lambda { |env| [200, @headers, @body] }
+    @server = Puma::Server.new @simple
+    @server.add_tcp_listener "127.0.0.1", 9988
+    @server.run
+
+    @client = TCPSocket.new "127.0.0.1", 9988
+  end
+
+  def teardown
+    @client.close
+    @server.stop(true)
+  end
+
+  def lines(count, s=@client)
+    str = ""
+    count.times { str << s.gets }
+    str
+  end
+
+  def test_one_with_content_length
+    @client << @valid_request
+    sz = @body[0].size.to_s
+
+    assert_equal "HTTP/1.1 200 OK\r\nContent-Length: #{sz}\r\nX-Header: Works\r\n\r\n", lines(4)
+    assert_equal "Hello", @client.read(5)
+  end
+
+  def test_two_back_to_back
+    @client << @valid_request
+    sz = @body[0].size.to_s
+
+    assert_equal "HTTP/1.1 200 OK\r\nContent-Length: #{sz}\r\nX-Header: Works\r\n\r\n", lines(4)
+    assert_equal "Hello", @client.read(5)
+
+    @client << @valid_request
+    sz = @body[0].size.to_s
+
+    assert_equal "HTTP/1.1 200 OK\r\nContent-Length: #{sz}\r\nX-Header: Works\r\n\r\n", lines(4)
+    assert_equal "Hello", @client.read(5)
+  end
+
+  def test_chunked
+    @body << "Chunked"
+
+    @client << @valid_request
+
+    assert_equal "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\nX-Header: Works\r\n\r\n5\r\nHello\r\n7\r\nChunked\r\n0\r\n", lines(9)
+  end
+
+end

test/test_unix_socket.rb

@@ -11,6 +11,8 @@ class TestPumaUnixSocket < Test::Unit::TestCase
 
   def setup
     @server = Puma::Server.new App
+    @server.add_unix_listener Path
+    @server.run
   end
 
   def teardown
@@ -19,14 +21,14 @@ class TestPumaUnixSocket < Test::Unit::TestCase
   end
 
   def test_server
-    @server.add_unix_listener Path
-    @server.run
-
     sock = UNIXSocket.new Path
 
     sock << "GET / HTTP/1.0\r\nHost: blah.com\r\n\r\n"
 
-    assert_equal "HTTP/1.1 200 OK\r\nConnection: close\r\n\r\nWorks",
-                 sock.read
+    expected = "HTTP/1.1 200 OK\r\nContent-Length: 5\r\n\r\nWorks"
+
+    assert_equal expected, sock.read(expected.size)
+
+    sock.close
   end
 end