Add early hints feature

lib/puma/cli.rb

@@ -181,6 +181,10 @@ def setup_options
             user_config.tcp_mode!
           end
 
+          o.on "--early-hints", "Enable early hints support" do
+            user_config.early_hints
+          end
+
           o.on "-V", "--version", "Print the version information" do
             puts "puma version #{Puma::Const::VERSION}"
             exit 0

lib/puma/const.rb

@@ -221,5 +221,7 @@ module Const
     HIJACK_P = "rack.hijack?".freeze
     HIJACK = "rack.hijack".freeze
     HIJACK_IO = "rack.hijack_io".freeze
+
+    EARLY_HINTS = "rack.early_hints".freeze
   end
 end

lib/puma/dsl.rb

@@ -241,6 +241,10 @@ def tcp_mode!
       @options[:mode] = :tcp
     end
 
+    def early_hints(answer=true)
+      @options[:early_hints] = answer
+    end
+
     # Redirect STDOUT and STDERR to files specified.
     def stdout_redirect(stdout=nil, stderr=nil, append=false)
       @options[:redirect_stdout] = stdout

lib/puma/runner.rb

@@ -161,6 +161,10 @@ def start_server
         server.tcp_mode!
       end
 
+      if @options[:early_hints]
+        server.early_hints = true
+      end
+
       unless development?
         server.leak_stack_on_error = false
       end

lib/puma/server.rb

@@ -81,7 +81,7 @@ def initialize(app, events=Events.stdio, options={})
       @precheck_closing = true
     end
 
-    attr_accessor :binder, :leak_stack_on_error
+    attr_accessor :binder, :leak_stack_on_error, :early_hints
 
     forward :add_tcp_listener,  :@binder
     forward :add_ssl_listener,  :@binder
@@ -595,6 +595,24 @@ def handle_request(req, lines)
       env[RACK_INPUT] = body
       env[RACK_URL_SCHEME] =  env[HTTPS_KEY] ? HTTPS : HTTP
 
+      if @early_hints
+        env[EARLY_HINTS] = lambda { |headers|
+          fast_write client, "HTTP/1.1 103 Early Hints\r\n".freeze
+
+          headers.each_pair do |k, vs|
+            if vs.respond_to?(:to_s) && !vs.to_s.empty?
+              vs.to_s.split(NEWLINE).each do |v|
+                fast_write client, "#{k}: #{v}\r\n"
+              end
+            else
+              fast_write client, "#{k}: #{v}\r\n"
+            end
+          end
+
+          fast_write client, "\r\n".freeze
+        }
+      end
+
       # A rack extension. If the app writes #call'ables to this
       # array, we will invoke them when the request is done.
       #

test/test_puma_server.rb

@@ -159,6 +159,61 @@ def test_GET_with_empty_body_has_sane_chunking
     assert_equal "HTTP/1.0 200 OK\r\nContent-Length: 0\r\n\r\n", data
   end
 
+  def test_early_hints_works
+    @server.app = proc { |env|
+     env['rack.early_hints'].call("Link" => "</style.css>; rel=preload; as=style\n</script.js>; rel=preload")
+     [200, { "X-Hello" => "World" }, ["Hello world!"]]
+    }
+
+    @server.add_tcp_listener @host, @port
+    @server.early_hints = true
+    @server.run
+
+    sock = TCPSocket.new @host, @server.connected_port
+    sock << "HEAD / HTTP/1.0\r\n\r\n"
+
+    data = sock.read
+
+    expected_data = (<<EOF
+HTTP/1.1 103 Early Hints
+Link: </style.css>; rel=preload; as=style
+Link: </script.js>; rel=preload
+
+HTTP/1.0 200 OK
+X-Hello: World
+Content-Length: 12
+EOF
+).split("\n").join("\r\n") + "\r\n\r\n"
+
+    assert_equal true, @server.early_hints
+    assert_equal expected_data, data
+  end
+
+  def test_early_hints_is_off_by_default
+    @server.app = proc { |env|
+     assert_nil env['rack.early_hints']
+     [200, { "X-Hello" => "World" }, ["Hello world!"]]
+    }
+
+    @server.add_tcp_listener @host, @port
+    @server.run
+
+    sock = TCPSocket.new @host, @server.connected_port
+    sock << "HEAD / HTTP/1.0\r\n\r\n"
+
+    data = sock.read
+
+    expected_data = (<<EOF
+HTTP/1.0 200 OK
+X-Hello: World
+Content-Length: 12
+EOF
+).split("\n").join("\r\n") + "\r\n\r\n"
+
+    assert_nil @server.early_hints
+    assert_equal expected_data, data
+  end
+
   def test_GET_with_no_body_has_sane_chunking
     @server.app = proc { |env| [200, {}, []] }