From d5914573c90fe6f8b9422f397e28c0358ce4aa3b Mon Sep 17 00:00:00 2001
From: Evan Phoenix <evan@phx.io>
Date: Wed, 6 Jan 2016 10:12:09 -0800
Subject: [PATCH] Add 'set_remote_address' config option

---
 lib/puma/client.rb        | 19 +++++++++++++++++++
 lib/puma/configuration.rb |  1 +
 lib/puma/const.rb         |  1 +
 lib/puma/dsl.rb           | 38 ++++++++++++++++++++++++++++++++++++++
 lib/puma/server.rb        | 20 ++++++++++++++++++--
 5 files changed, 77 insertions(+), 2 deletions(-)

diff --git a/lib/puma/client.rb b/lib/puma/client.rb
index e842cee9..71a37c0a 100644
--- a/lib/puma/client.rb
+++ b/lib/puma/client.rb
@@ -45,11 +45,18 @@ module Puma
 
       @requests_served = 0
       @hijacked = false
+
+      @peerip = nil
+      @remote_addr_header = nil
     end
 
     attr_reader :env, :to_io, :body, :io, :timeout_at, :ready, :hijacked,
                 :tempfile
 
+    attr_writer :peerip
+
+    attr_accessor :remote_addr_header
+
     def inspect
       "#<Puma::Client:0x#{object_id.to_s(16)} @ready=#{@ready.inspect}>"
     end
@@ -297,5 +304,17 @@ module Puma
       rescue StandardError
       end
     end
+
+    def peerip
+      return @peerip if @peerip
+
+      if @remote_addr_header
+        hdr = @env[@remote_addr_header] || LOCALHOST_ADDR
+        @peerip = hdr
+        return hdr
+      end
+
+      @peerip ||= @io.peeraddr.last
+    end
   end
 end
diff --git a/lib/puma/configuration.rb b/lib/puma/configuration.rb
index 475de7a8..f6becf52 100644
--- a/lib/puma/configuration.rb
+++ b/lib/puma/configuration.rb
@@ -29,6 +29,7 @@ module Puma
       @conf[:worker_timeout] ||= DefaultWorkerTimeout
       @conf[:worker_boot_timeout] ||= @conf[:worker_timeout]
       @conf[:worker_shutdown_timeout] ||= DefaultWorkerShutdownTimeout
+      @conf[:remote_address] ||= :socket
 
       @options = {}
     end
diff --git a/lib/puma/const.rb b/lib/puma/const.rb
index a3ac44cd..579914f8 100644
--- a/lib/puma/const.rb
+++ b/lib/puma/const.rb
@@ -183,6 +183,7 @@ module Puma
     PORT_443 = "443".freeze
     LOCALHOST = "localhost".freeze
     LOCALHOST_IP = "127.0.0.1".freeze
+    LOCALHOST_ADDR = "127.0.0.1:0".freeze
 
     SERVER_PROTOCOL = "SERVER_PROTOCOL".freeze
     HTTP_11 = "HTTP/1.1".freeze
diff --git a/lib/puma/dsl.rb b/lib/puma/dsl.rb
index 7c40d8b3..7a0f6d0f 100644
--- a/lib/puma/dsl.rb
+++ b/lib/puma/dsl.rb
@@ -297,5 +297,43 @@ module Puma
     def shutdown_debug(val=true)
       @options[:shutdown_debug] = val
     end
+
+    # Control how the remote address of the connection is set. This
+    # is configurable because to calculate the true socket peer address
+    # a kernel syscall is required which for very fast rack handlers
+    # slows down the handling significantly.
+    #
+    # There are 4 possible values:
+    #
+    # * :socket (the default) - read the peername from the socket using the
+    #           syscall. This is the normal behavior.
+    # * :localhost - set the remote address to "127.0.0.1"
+    # * header: http_header - set the remote address to the value of the
+    #                          provided http header. For instance:
+    #                          `set_remote_address header: "X-Real-IP"`
+    # * Any string - this allows you to hardcode remote address to any value
+    #                you wish. Because puma never uses this field anyway, it's
+    #                format is entirely in your hands.
+    def set_remote_address(val=:socket)
+      case val
+      when :socket
+        @options[:remote_address] = val
+      when :localhost
+        @options[:remote_address] = :value
+        @options[:remote_address_value] = "127.0.0.1".freeze
+      when String
+        @options[:remote_address] = :value
+        @options[:remote_address_value] = val
+      when Hash
+        if hdr = val[:header]
+          @options[:remote_address] = :header
+          @options[:remote_address_header] = "HTTP_" + hdr.upcase.gsub("-", "_")
+        else
+          raise "Invalid value for set_remote_address - #{val.inspect}"
+        end
+      else
+        raise "Invalid value for set_remote_address - #{val}"
+      end
+    end
   end
 end
diff --git a/lib/puma/server.rb b/lib/puma/server.rb
index 0a17e925..7ebe0eed 100644
--- a/lib/puma/server.rb
+++ b/lib/puma/server.rb
@@ -307,6 +307,16 @@ module Puma
         pool = @thread_pool
         queue_requests = @queue_requests
 
+        remote_addr_value = nil
+        remote_addr_header = nil
+
+        case @options[:remote_address]
+        when :value
+          remote_addr_value = @options[:remote_address_value]
+        when :header
+          remote_addr_header = @options[:remote_address_header]
+        end
+
         while @status == :run
           begin
             ios = IO.select sockets
@@ -317,6 +327,12 @@ module Puma
                 begin
                   if io = sock.accept_nonblock
                     client = Client.new io, @binder.env(sock)
+                    if remote_addr_value
+                      client.peerip = remote_addr_value
+                    elsif remote_addr_header
+                      client.remote_addr_header = remote_addr_header
+                    end
+
                     pool << client
                     pool.wait_until_not_full unless queue_requests
                   end
@@ -481,7 +497,7 @@ module Puma
 
       unless env.key?(REMOTE_ADDR)
         begin
-          addr = client.peeraddr.last
+          addr = client.peerip
         rescue Errno::ENOTCONN
           # Client disconnects can result in an inability to get the
           # peeraddr from the socket; default to localhost.
@@ -513,7 +529,7 @@ module Puma
       env = req.env
       client = req.io
 
-      normalize_env env, client
+      normalize_env env, req
 
       env[PUMA_SOCKET] = client