add support for the PROXY protocol (v1 only)

lib/puma/client.rb

@@ -56,6 +56,7 @@ def initialize(io, env=nil)
       @parser = HttpParser.new
       @parsed_bytes = 0
       @read_header = true
+      @read_proxy = false
       @ready = false
 
       @body = nil
@@ -71,6 +72,7 @@ def initialize(io, env=nil)
       @peerip = nil
       @listener = nil
       @remote_addr_header = nil
+      @expect_proxy_proto = false
 
       @body_remain = 0
 
@@ -106,7 +108,7 @@ def call
 
     # @!attribute [r] in_data_phase
     def in_data_phase
-      !@read_header
+      !(@read_header || @read_proxy)
     end
 
     def set_timeout(val)
@@ -121,6 +123,7 @@ def timeout
     def reset(fast_check=true)
       @parser.reset
       @read_header = true
+      @read_proxy = !!@expect_proxy_proto
       @env = @proto_env.dup
       @body = nil
       @tempfile = nil
@@ -163,7 +166,7 @@ def close
     end
 
     def try_to_finish
-      return read_body unless @read_header
+      return read_body if in_data_phase
 
       begin
         data = @io.read_nonblock(CHUNK_SIZE)
@@ -188,6 +191,24 @@ def try_to_finish
         @buffer = data
       end
 
+      if @read_proxy
+        if @expect_proxy_proto == :v1
+          if @buffer.include? "\r\n"
+            if md = PROXY_PROTOCOL_V1_REGEX.match(@buffer)
+              if md[1]
+                @peerip = md[1].split(" ")[0]
+              end
+              @buffer = md.post_match
+            end
+            # if the buffer has a \r\n but doesn't have a PROXY protocol
+            # request, this is just HTTP from a non-PROXY client; move on
+            @read_proxy = false
+          else
+            return false
+          end
+        end
+      end
+
       @parsed_bytes = @parser.execute(@env, @buffer, @parsed_bytes)
 
       if @parser.finished?
@@ -244,6 +265,17 @@ def can_close?
       @parsed_bytes == 0
     end
 
+    def expect_proxy_proto=(val)
+      if val
+        if @read_header
+          @read_proxy = true
+        end
+      else
+        @read_proxy = false
+      end
+      @expect_proxy_proto = val
+    end
+
     private
 
     def setup_body

lib/puma/const.rb

@@ -247,5 +247,7 @@ module Const
 
     # Banned keys of response header
     BANNED_HEADER_KEY = /\A(rack\.|status\z)/.freeze
+
+    PROXY_PROTOCOL_V1_REGEX = /^PROXY (?:TCP4|TCP6|UNKNOWN) ([^\r]+)\r\n/.freeze
   end
 end

lib/puma/dsl.rb

@@ -828,7 +828,10 @@ def wait_for_less_busy_worker(val=0.005)
     #    `set_remote_address header: "X-Real-IP"`.
     #    Only the first word (as separated by spaces or comma) is used, allowing
     #    headers such as X-Forwarded-For to be used as well.
-    # 4. **\<Any string\>** - this allows you to hardcode remote address to any value
+    # 4. **proxy_protocol: v1**- set the remote address to the value read from the
+    #    HAproxy PROXY protocol, version 1. If the request does not have the PROXY
+    #    protocol attached to it, will fall back to :socket
+    # 5. **\<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.
     #
@@ -846,6 +849,13 @@ def set_remote_address(val=:socket)
         if hdr = val[:header]
           @options[:remote_address] = :header
           @options[:remote_address_header] = "HTTP_" + hdr.upcase.tr("-", "_")
+        elsif protocol_version = val[:proxy_protocol]
+          @options[:remote_address] = :proxy_protocol
+          protocol_version = protocol_version.downcase.to_sym
+          if ! [:v1].include?(protocol_version)
+            raise "Invalid value for proxy_protocol - #{protocol_version.inspect}"
+          end
+          @options[:remote_address_proxy_protocol] = protocol_version
         else
           raise "Invalid value for set_remote_address - #{val.inspect}"
         end

lib/puma/server.rb

@@ -321,6 +321,8 @@ def handle_servers
           remote_addr_value = @options[:remote_address_value]
         when :header
           remote_addr_header = @options[:remote_address_header]
+        when :proxy_protocol
+          remote_addr_proxy_protocol = @options[:remote_address_proxy_protocol]
         end
 
         while @status == :run || (drain && shutting_down?)
@@ -346,6 +348,8 @@ def handle_servers
                   client.peerip = remote_addr_value
                 elsif remote_addr_header
                   client.remote_addr_header = remote_addr_header
+                elsif remote_addr_proxy_protocol
+                  client.expect_proxy_proto = remote_addr_proxy_protocol
                 end
                 pool << client
               end

test/test_puma_server.rb

@@ -2,6 +2,7 @@
 require "puma/events"
 require "net/http"
 require "nio"
+require "ipaddr"
 
 class TestPumaServer < Minitest::Test
   parallelize_me! unless JRUBY_HEAD
@@ -48,6 +49,14 @@ def send_http(req)
     new_connection << req
   end
 
+  def send_proxy_v1_http(req, remote_ip)
+    addr = IPAddr.new(remote_ip)
+    family = addr.ipv4? ? "TCP4" : "TCP6"
+    target = addr.ipv4? ? "127.0.0.1" : "::1"
+    conn = new_connection
+    conn << ("PROXY #{family} #{remote_ip} #{target} 10000 80\r\n" + req)
+  end
+
   def new_connection
     TCPSocket.new(@host, @port).tap {|sock| @ios << sock}
   end
@@ -1017,6 +1026,18 @@ def test_newline_splits_in_early_hint
     assert_match "X-header: first line\r\nX-header: second line\r\n", data
   end
 
+  def test_proxy_protocol
+    server_run(remote_address: :proxy_protocol, remote_address_proxy_protocol: :v1) do |env|
+      [200, {}, [env["REMOTE_ADDR"]]]
+    end
+
+    remote_addr = send_proxy_v1_http("GET / HTTP/1.0\r\n\r\n", "1.2.3.4").read.split("\r\n").last
+    assert_equal '1.2.3.4', remote_addr
+
+    remote_addr = send_proxy_v1_http("GET / HTTP/1.0\r\n\r\n", "fd00::1").read.split("\r\n").last
+    assert_equal 'fd00::1', remote_addr
+  end
+
   # To comply with the Rack spec, we have to split header field values
   # containing newlines into multiple headers.
   def assert_does_not_allow_http_injection(app, opts = {})