From 32852f715af9e6eb809284d910e647b8fd6e2ea6 Mon Sep 17 00:00:00 2001
From: MSP-Greg <MSP-Greg@users.noreply.github.com>
Date: Wed, 26 May 2021 09:15:08 -0500
Subject: [PATCH] Add rack_url_scheme to Puma::DSL, allows setting of
 rack.url_scheme header (#2586)

* Add rack_url_scheme to Puma::DSL, allows setting of rack.url_scheme header

* Clarify comment on DSL re: rack URL

* Add tests to test_puma_server.rb

Co-authored-by: Nate Berkopec <nate.berkopec@gmail.com>
---
 lib/puma/binder.rb       |  1 +
 lib/puma/dsl.rb          |  7 +++++++
 lib/puma/request.rb      |  2 +-
 test/test_puma_server.rb | 21 ++++++++++++++++++++-
 4 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/lib/puma/binder.rb b/lib/puma/binder.rb
index f8939ac2e9..0c209e4772 100644
--- a/lib/puma/binder.rb
+++ b/lib/puma/binder.rb
@@ -41,6 +41,7 @@ def initialize(events, conf = Configuration.new)
         "rack.multithread".freeze => conf.options[:max_threads] > 1,
         "rack.multiprocess".freeze => conf.options[:workers] >= 1,
         "rack.run_once".freeze => false,
+        RACK_URL_SCHEME => conf.options[:rack_url_scheme],
         "SCRIPT_NAME".freeze => ENV['SCRIPT_NAME'] || "",
 
         # I'd like to set a default CONTENT_TYPE here but some things
diff --git a/lib/puma/dsl.rb b/lib/puma/dsl.rb
index c4c66813ae..36f0441056 100644
--- a/lib/puma/dsl.rb
+++ b/lib/puma/dsl.rb
@@ -381,6 +381,13 @@ def rackup(path)
       @options[:rackup] ||= path.to_s
     end
 
+    # Allows setting `env['rack.url_scheme']`.
+    # Only necessary if X-Forwarded-Proto is not being set by your proxy
+    # Normal values are 'http' or 'https'.
+    def rack_url_scheme(scheme=nil)
+      @options[:rack_url_scheme] = scheme
+    end
+
     def early_hints(answer=true)
       @options[:early_hints] = answer
     end
diff --git a/lib/puma/request.rb b/lib/puma/request.rb
index 82db731f7a..5b7ab22881 100644
--- a/lib/puma/request.rb
+++ b/lib/puma/request.rb
@@ -51,7 +51,7 @@ def handle_request(client, lines, requests)
       head = env[REQUEST_METHOD] == HEAD
 
       env[RACK_INPUT] = body
-      env[RACK_URL_SCHEME] = default_server_port(env) == PORT_443 ? HTTPS : HTTP
+      env[RACK_URL_SCHEME] ||= default_server_port(env) == PORT_443 ? HTTPS : HTTP
 
       if @early_hints
         env[EARLY_HINTS] = lambda { |headers|
diff --git a/test/test_puma_server.rb b/test/test_puma_server.rb
index 7ed629e4af..7fc64fba98 100644
--- a/test/test_puma_server.rb
+++ b/test/test_puma_server.rb
@@ -304,7 +304,6 @@ def test_doesnt_print_backtrace_in_production
     assert_match(/HTTP\/1.0 500 Internal Server Error/, data)
   end
 
-
   def test_eof_on_connection_close_is_not_logged_as_an_error
     server_run
 
@@ -1312,4 +1311,24 @@ def test_drain_on_shutdown(drain=true)
   def test_not_drain_on_shutdown
     test_drain_on_shutdown false
   end
+
+  def test_rack_url_scheme_dflt
+    server_run
+
+    data = send_http_and_read "GET / HTTP/1.0\r\n\r\n"
+    assert_equal "http", data.split("\r\n").last
+  end
+
+  def test_rack_url_scheme_user
+    @port = UniquePort.call
+    opts = { rack_url_scheme: 'user', binds: ["tcp://#{@host}:#{@port}"] }
+    conf = Puma::Configuration.new(opts).tap(&:clamp)
+    @server = Puma::Server.new @app, @events, conf.options
+    @server.inherit_binder Puma::Binder.new(@events, conf)
+    @server.binder.parse conf.options[:binds], @events
+    @server.run
+
+    data = send_http_and_read "GET / HTTP/1.0\r\n\r\n"
+    assert_equal "user", data.split("\r\n").last
+  end
 end