From 59764b6ee7e863298a4f655903df6a6ca4f10c29 Mon Sep 17 00:00:00 2001
From: eileencodes <eileencodes@gmail.com>
Date: Tue, 29 Aug 2017 13:14:23 -0400
Subject: [PATCH] Add early hints feature

This commit adds the early hints lambda to the headers hash. Users can
call it to emit the early hints headers. For example:

```
class Server
  def call env
    if env["REQUEST_PATH"] == "/"
      env['rack.early_hints'].call([{ link: "/style.css", as: "style" }, { link: "/script.js" }])
      [200, { "X-Hello" => "World" }, ["Hello world!"]]
    else
      [200, { "X-Hello" => "World" }, ["NEAT!"]]
    end
  end
end

run Server.new
```

In this example, the server sends stylesheet and javascript early hints
if the proxy supports it, it will send H2 pushes to the client.

Of course not every proxy server supports early hints, so to enable the
early hints feature with puma you have to pass the configuration variable,
`--early-hints`.
---
 lib/puma/cli.rb          |  4 ++++
 lib/puma/const.rb        |  2 ++
 lib/puma/dsl.rb          |  4 ++++
 lib/puma/runner.rb       |  4 ++++
 lib/puma/server.rb       | 23 +++++++++++++++++++++++
 test/test_puma_server.rb | 26 ++++++++++++++++++++++++++
 6 files changed, 63 insertions(+)

diff --git a/lib/puma/cli.rb b/lib/puma/cli.rb
index 7d747d9c5c..bcfeb33050 100644
--- a/lib/puma/cli.rb
+++ b/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
diff --git a/lib/puma/const.rb b/lib/puma/const.rb
index bb85eefdd0..a8751cea63 100644
--- a/lib/puma/const.rb
+++ b/lib/puma/const.rb
@@ -220,5 +220,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
diff --git a/lib/puma/dsl.rb b/lib/puma/dsl.rb
index e2e142246c..772638dce2 100644
--- a/lib/puma/dsl.rb
+++ b/lib/puma/dsl.rb
@@ -241,6 +241,10 @@ def tcp_mode!
       @options[:mode] = :tcp
     end
 
+    def early_hints!
+      @options[:early_hints] = true
+    end
+
     # Redirect STDOUT and STDERR to files specified.
     def stdout_redirect(stdout=nil, stderr=nil, append=false)
       @options[:redirect_stdout] = stdout
diff --git a/lib/puma/runner.rb b/lib/puma/runner.rb
index cd556cff1c..f57f127058 100644
--- a/lib/puma/runner.rb
+++ b/lib/puma/runner.rb
@@ -161,6 +161,10 @@ def start_server
         server.tcp_mode!
       end
 
+      if @options[:early_hints]
+        server.early_hints!
+      end
+
       unless development?
         server.leak_stack_on_error = false
       end
diff --git a/lib/puma/server.rb b/lib/puma/server.rb
index 51cd810811..8990c63061 100644
--- a/lib/puma/server.rb
+++ b/lib/puma/server.rb
@@ -77,6 +77,7 @@ def initialize(app, events=Events.stdio, options={})
       ENV['RACK_ENV'] ||= "development"
 
       @mode = :http
+      @early_hints = false
 
       @precheck_closing = true
     end
@@ -97,6 +98,10 @@ def tcp_mode!
       @mode = :tcp
     end
 
+    def early_hints!
+      @early_hints = true
+    end
+
     # On Linux, use TCP_CORK to better control how the TCP stack
     # packetizes our stream. This improves both latency and throughput.
     #
@@ -595,6 +600,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 { |links|
+          fast_write client, "HTTP/1.1 103 Early Hints\r\n".freeze
+
+          links.each do |link|
+            fast_write client, "Link: <#{link[:link]}>; rel=preload"
+            if as = link[:as]
+              fast_write client, "; as=#{as}"
+            end
+            fast_write client, "\r\n".freeze
+          end
+
+          fast_write client, "\r\n".freeze
+        }
+      else
+        env[EARLY_HINTS] = lambda { |_| }
+      end
+
       # A rack extension. If the app writes #call'ables to this
       # array, we will invoke them when the request is done.
       #
diff --git a/test/test_puma_server.rb b/test/test_puma_server.rb
index ad573ca729..3f34d11048 100644
--- a/test/test_puma_server.rb
+++ b/test/test_puma_server.rb
@@ -159,6 +159,32 @@ 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", as: "style" }, { link: "/script.js" }])
+     [200, { "X-Hello" => "World" }, ["Hello world!"]]
+    }
+
+    @server.add_tcp_listener @host, @port
+    @server.early_hints!
+    @server.run
+
+    sock = TCPSocket.new @host, @server.connected_port
+    sock << "HEAD / HTTP/1.0\r\n\r\n"
+
+    data = sock.read
+
+    assert_equal <<~eos.split("\n").join("\r\n") + "\r\n\r\n", data
+    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
+    eos
+  end
+
   def test_GET_with_no_body_has_sane_chunking
     @server.app = proc { |env| [200, {}, []] }