Proper 'Connection' header handling compatible with HTTP 1.0 and HTTP 1.1 protocols

lib/puma/server.rb

@@ -535,7 +535,7 @@ module Puma
         line_ending = LINE_END
         colon = COLON
 
-        if env[HTTP_VERSION] == HTTP_11
+        http_11 = if env[HTTP_VERSION] == HTTP_11
           allow_chunked = true
           keep_alive = env[HTTP_CONNECTION] != CLOSE
           include_keepalive_header = false
@@ -552,6 +552,7 @@ module Puma
 
             no_body ||= status < 200 || STATUS_WITH_NO_ENTITY_BODY[status]
           end
+          true
         else
           allow_chunked = false
           keep_alive = env[HTTP_CONNECTION] == KEEP_ALIVE
@@ -567,6 +568,7 @@ module Puma
 
             no_body ||= status < 200 || STATUS_WITH_NO_ENTITY_BODY[status]
           end
+          false
         end
 
         response_hijack = nil
@@ -593,6 +595,12 @@ module Puma
           end
         end
 
+        if include_keepalive_header
+          lines << CONNECTION_KEEP_ALIVE
+        elsif http_11 && !keep_alive
+          lines << CONNECTION_CLOSE
+        end
+
         if no_body
           if content_length and status != 204
             lines.append CONTENT_LENGTH_S, content_length.to_s, line_ending
@@ -603,12 +611,6 @@ module Puma
           return keep_alive
         end
 
-        if include_keepalive_header
-          lines << CONNECTION_KEEP_ALIVE
-        elsif !keep_alive
-          lines << CONNECTION_CLOSE
-        end
-
         unless response_hijack
           if content_length
             lines.append CONTENT_LENGTH_S, content_length.to_s, line_ending

test/test_persistent.rb

@@ -106,7 +106,7 @@ class TestPersistent < Test::Unit::TestCase
 
     @client << @http10_request
 
-    assert_equal "HTTP/1.0 200 OK\r\nX-Header: Works\r\nConnection: close\r\n\r\n", lines(4)
+    assert_equal "HTTP/1.0 200 OK\r\nX-Header: Works\r\n\r\n", lines(3)
     assert_equal "HelloChunked", @client.read
   end
 
@@ -132,7 +132,7 @@ class TestPersistent < Test::Unit::TestCase
     @client << @http10_request
     sz = @body[0].size.to_s
 
-    assert_equal "HTTP/1.0 200 OK\r\nX-Header: Works\r\nConnection: close\r\nContent-Length: #{sz}\r\n\r\n", lines(5)
+    assert_equal "HTTP/1.0 200 OK\r\nX-Header: Works\r\nContent-Length: #{sz}\r\n\r\n", lines(4)
     assert_equal "Hello", @client.read(5)
   end
 

test/test_puma_server.rb

@@ -216,7 +216,7 @@ class TestPumaServer < Test::Unit::TestCase
 
     data = sock.read
 
-    assert_equal "HTTP/1.0 449 CUSTOM\r\nConnection: close\r\nContent-Length: 0\r\n\r\n", data
+    assert_equal "HTTP/1.0 449 CUSTOM\r\nContent-Length: 0\r\n\r\n", data
   end
 
   def test_custom_http_codes_11
@@ -285,4 +285,116 @@ class TestPumaServer < Test::Unit::TestCase
 
     assert_equal "HTTP/1.1 408 Request Timeout\r\n", data
   end
+
+  def test_http_11_keep_alive_with_body
+    @server.app = proc { |env| [200, {"Content-Type" => "plain/text"}, ["hello"]] }
+
+    @server.add_tcp_listener @host, @port
+    @server.run
+
+    sock = TCPSocket.new @host, @port
+    sock << "GET / HTTP/1.1\r\nConnection: Keep-Alive\r\n\r\n"
+
+    data = sock.read
+
+    assert_equal "HTTP/1.1 200 OK\r\nContent-Type: plain/text\r\nContent-Length: 5\r\n\r\nhello", data
+  end
+
+  def test_http_11_close_with_body
+    @server.app = proc { |env| [200, {"Content-Type" => "plain/text"}, ["hello"]] }
+
+    @server.add_tcp_listener @host, @port
+    @server.run
+
+    sock = TCPSocket.new @host, @port
+    sock << "GET / HTTP/1.1\r\nConnection: close\r\n\r\n"
+
+    data = sock.read
+
+    assert_equal "HTTP/1.1 200 OK\r\nContent-Type: plain/text\r\nConnection: close\r\nContent-Length: 5\r\n\r\nhello", data
+  end
+
+  def test_http_11_keep_alive_without_body
+    @server.app = proc { |env| [204, {}, []] }
+
+    @server.add_tcp_listener @host, @port
+    @server.run
+
+    sock = TCPSocket.new @host, @port
+    sock << "GET / HTTP/1.1\r\nConnection: Keep-Alive\r\n\r\n"
+
+    data = sock.read
+
+    assert_equal "HTTP/1.1 204 No Content\r\n\r\n", data
+  end
+
+  def test_http_11_close_without_body
+    @server.app = proc { |env| [204, {}, []] }
+
+    @server.add_tcp_listener @host, @port
+    @server.run
+
+    sock = TCPSocket.new @host, @port
+    sock << "GET / HTTP/1.1\r\nConnection: close\r\n\r\n"
+
+    data = sock.read
+
+    assert_equal "HTTP/1.1 204 No Content\r\nConnection: close\r\n\r\n", data
+  end
+
+  def test_http_10_keep_alive_with_body
+    @server.app = proc { |env| [200, {"Content-Type" => "plain/text"}, ["hello"]] }
+
+    @server.add_tcp_listener @host, @port
+    @server.run
+
+    sock = TCPSocket.new @host, @port
+    sock << "GET / HTTP/1.0\r\nConnection: Keep-Alive\r\n\r\n"
+
+    data = sock.read
+
+    assert_equal "HTTP/1.0 200 OK\r\nContent-Type: plain/text\r\nConnection: Keep-Alive\r\nContent-Length: 5\r\n\r\nhello", data
+  end
+
+  def test_http_10_close_with_body
+    @server.app = proc { |env| [200, {"Content-Type" => "plain/text"}, ["hello"]] }
+
+    @server.add_tcp_listener @host, @port
+    @server.run
+
+    sock = TCPSocket.new @host, @port
+    sock << "GET / HTTP/1.0\r\nConnection: close\r\n\r\n"
+
+    data = sock.read
+
+    assert_equal "HTTP/1.0 200 OK\r\nContent-Type: plain/text\r\nContent-Length: 5\r\n\r\nhello", data
+  end
+
+  def test_http_10_keep_alive_without_body
+    @server.app = proc { |env| [204, {}, []] }
+
+    @server.add_tcp_listener @host, @port
+    @server.run
+
+    sock = TCPSocket.new @host, @port
+    sock << "GET / HTTP/1.0\r\nConnection: Keep-Alive\r\n\r\n"
+
+    data = sock.read
+
+    assert_equal "HTTP/1.0 204 No Content\r\nConnection: Keep-Alive\r\n\r\n", data
+  end
+
+  def test_http_10_close_without_body
+    @server.app = proc { |env| [204, {}, []] }
+
+    @server.add_tcp_listener @host, @port
+    @server.run
+
+    sock = TCPSocket.new @host, @port
+    sock << "GET / HTTP/1.0\r\nConnection: close\r\n\r\n"
+
+    data = sock.read
+
+    assert_equal "HTTP/1.0 204 No Content\r\n\r\n", data
+  end
 end

test/test_unix_socket.rb

@@ -8,31 +8,31 @@ require 'socket'
 # (or Windows)
 unless defined?(JRUBY_VERSION) || RbConfig::CONFIG["host_os"] =~ /mingw|mswin/
   class TestPumaUnixSocket < Test::Unit::TestCase
-  
+
     App = lambda { |env| [200, {}, ["Works"]] }
-  
+
     Path = "test/puma.sock"
-  
+
     def setup
       @server = Puma::Server.new App
       @server.add_unix_listener Path
       @server.run
     end
-  
+
     def teardown
       @server.stop(true)
       File.unlink Path if File.exist? Path
     end
-  
+
     def test_server
       sock = UNIXSocket.new Path
-  
+
       sock << "GET / HTTP/1.0\r\nHost: blah.com\r\n\r\n"
-  
-      expected = "HTTP/1.0 200 OK\r\nConnection: close\r\nContent-Length: 5\r\n\r\nWorks"
-  
+
+      expected = "HTTP/1.0 200 OK\r\nContent-Length: 5\r\n\r\nWorks"
+
       assert_equal expected, sock.read(expected.size)
-  
+
       sock.close
     end
   end