From c55bf5c156de2eae566ac3e0d1c04fc5b6a2a28a 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`.

If `ENV['rack.early_hints']` is not set then early hints is not
supported by the webserver. Early hints is off by default.
---
 lib/puma/cli.rb          |  4 +++
 lib/puma/const.rb        |  2 ++
 lib/puma/dsl.rb          |  4 +++
 lib/puma/runner.rb       |  4 +++
 lib/puma/server.rb       | 18 ++++++++++++-
 test/test_puma_server.rb | 55 ++++++++++++++++++++++++++++++++++++++++
 6 files changed, 86 insertions(+), 1 deletion(-)

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 597e191bfa..e25f8cab6a 100644
--- a/lib/puma/const.rb
+++ b/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
diff --git a/lib/puma/dsl.rb b/lib/puma/dsl.rb
index e2e142246c..0773f6fb45 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(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
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..1efbd06160 100644
--- a/lib/puma/server.rb
+++ b/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,22 @@ 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
+        }
+      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 1764016f38..a8d1a40267 100644
--- a/test/test_puma_server.rb
+++ b/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", as: "style" }, { link: "/script.js" }])
+     [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, {}, []] }