Merge pull request #1747 from rabbitmq/ranch_proxy_header

Use the built-in Ranch PROXY protocol support

3 files changed, 34 insertions, 29 deletions

diff --git a/src/rabbit_connection_sup.erl b/src/rabbit_connection_sup.erl
index 4dcfa8dc8a..7ec5262448 100644
--- a/src/rabbit_connection_sup.erl
+++ b/src/rabbit_connection_sup.erl
@@ -42,7 +42,7 @@
 
 %%--------------------------------------------------------------------------
 
-start_link(Ref, Sock, _Transport, _Opts) ->
+start_link(Ref, _Sock, _Transport, _Opts) ->
     {ok, SupPid} = supervisor2:start_link(?MODULE, []),
     %% We need to get channels in the hierarchy here so they get shut
     %% down after the reader, so the reader gets a chance to terminate
@@ -62,7 +62,7 @@ start_link(Ref, Sock, _Transport, _Opts) ->
     {ok, ReaderPid} =
         supervisor2:start_child(
           SupPid,
-          {reader, {rabbit_reader, start_link, [HelperSup, Ref, Sock]},
+          {reader, {rabbit_reader, start_link, [HelperSup, Ref]},
            intrinsic, ?WORKER_WAIT, worker, [rabbit_reader]}),
     {ok, SupPid, ReaderPid}.
 
diff --git a/src/rabbit_networking.erl b/src/rabbit_networking.erl
index 256d424740..cf431ee04f 100644
--- a/src/rabbit_networking.erl
+++ b/src/rabbit_networking.erl
@@ -36,7 +36,7 @@
          connection_info_all/0, connection_info_all/1,
          emit_connection_info_all/4, emit_connection_info_local/3,
          close_connection/2, accept_ack/2,
-         tcp_host/1]).
+         handshake/2, tcp_host/1]).
 
 %% Used by TCP-based transports, e.g. STOMP adapter
 -export([tcp_listener_addresses/1, tcp_listener_spec/9,
@@ -121,7 +121,6 @@ boot() ->
     %% Failures will throw exceptions
     _ = boot_listeners(fun boot_tcp/1, application:get_env(rabbit, num_tcp_acceptors, 10), "TCP"),
     _ = boot_listeners(fun boot_tls/1, application:get_env(rabbit, num_ssl_acceptors, 10), "TLS"),
-    _ = maybe_start_proxy_protocol(),
     ok.
 
 boot_listeners(Fun, NumAcceptors, Type) ->
@@ -190,12 +189,6 @@ log_poodle_fail(Context) ->
       "'rabbit' section of your configuration file.~n",
       [rabbit_misc:otp_release(), Context]).
 
-maybe_start_proxy_protocol() ->
-    case application:get_env(rabbit, proxy_protocol, false) of
-        false -> ok;
-        true  -> application:start(ranch_proxy_protocol)
-    end.
-
 fix_ssl_options(Config) ->
     rabbit_ssl_options:fix(Config).
 
@@ -263,12 +256,9 @@ start_listener0(Address, NumAcceptors, Protocol, Label, Opts) ->
     end.
 
 transport(Protocol) ->
-    ProxyProtocol = application:get_env(rabbit, proxy_protocol, false),
-    case {Protocol, ProxyProtocol} of
-        {amqp, false}       -> ranch_tcp;
-        {amqp, true}        -> ranch_proxy;
-        {'amqp/ssl', false} -> ranch_ssl;
-        {'amqp/ssl', true}  -> ranch_proxy_ssl
+    case Protocol of
+        amqp       -> ranch_tcp;
+        'amqp/ssl' -> ranch_ssl
     end.
 
 
@@ -368,16 +358,31 @@ close_connection(Pid, Explanation) ->
             ok
     end.
 
+handshake(Ref, ProxyProtocol) ->
+    case ProxyProtocol of
+        true ->
+            {ok, ProxyInfo} = ranch:recv_proxy_header(Ref, 1000),
+            {ok, Sock} = ranch:handshake(Ref),
+            tune_buffer_size(Sock),
+            ok = file_handle_cache:obtain(),
+            {ok, {rabbit_proxy_socket, Sock, ProxyInfo}};
+        false ->
+            ranch:handshake(Ref)
+    end.
+
 accept_ack(Ref, Sock) ->
     ok = ranch:accept_ack(Ref),
-    case tune_buffer_size(Sock) of
+    tune_buffer_size(Sock),
+    ok = file_handle_cache:obtain().
+
+tune_buffer_size(Sock) ->
+    case tune_buffer_size1(Sock) of
         ok         -> ok;
         {error, _} -> rabbit_net:fast_close(Sock),
                       exit(normal)
-    end,
-    ok = file_handle_cache:obtain().
+    end.
 
-tune_buffer_size(Sock) ->
+tune_buffer_size1(Sock) ->
     case rabbit_net:getopts(Sock, [sndbuf, recbuf, buffer]) of
         {ok, BufSizes} -> BufSz = lists:max([Sz || {_Opt, Sz} <- BufSizes]),
                           rabbit_net:setopts(Sock, [{buffer, BufSz}]);
diff --git a/src/rabbit_reader.erl b/src/rabbit_reader.erl
index 91002d0b94..8370d08ed9 100644
--- a/src/rabbit_reader.erl
+++ b/src/rabbit_reader.erl
@@ -57,12 +57,12 @@
 -include("rabbit_framing.hrl").
 -include("rabbit.hrl").
 
--export([start_link/3, info_keys/0, info/1, info/2,
+-export([start_link/2, info_keys/0, info/1, info/2,
          shutdown/2]).
 
 -export([system_continue/3, system_terminate/4, system_code_change/4]).
 
--export([init/4, mainloop/4, recvloop/4]).
+-export([init/3, mainloop/4, recvloop/4]).
 
 -export([conserve_resources/3, server_properties/1]).
 
@@ -157,7 +157,7 @@
 
 %%--------------------------------------------------------------------------
 
--spec start_link(pid(), any(), rabbit_net:socket()) -> rabbit_types:ok(pid()).
+-spec start_link(pid(), any()) -> rabbit_types:ok(pid()).
 -spec info_keys() -> rabbit_types:info_keys().
 -spec info(pid()) -> rabbit_types:infos().
 -spec info(pid(), rabbit_types:info_keys()) -> rabbit_types:infos().
@@ -170,7 +170,7 @@
           rabbit_framing:amqp_table().
 
 %% These specs only exists to add no_return() to keep dialyzer happy
--spec init(pid(), pid(), any(), rabbit_net:socket()) -> no_return().
+-spec init(pid(), pid(), any()) -> no_return().
 -spec start_connection(pid(), pid(), any(), rabbit_net:socket()) ->
           no_return().
 
@@ -181,18 +181,18 @@
 
 %%--------------------------------------------------------------------------
 
-start_link(HelperSup, Ref, Sock) ->
-    Pid = proc_lib:spawn_link(?MODULE, init, [self(), HelperSup, Ref, Sock]),
+start_link(HelperSup, Ref) ->
+    Pid = proc_lib:spawn_link(?MODULE, init, [self(), HelperSup, Ref]),
 
     {ok, Pid}.
 
 shutdown(Pid, Explanation) ->
     gen_server:call(Pid, {shutdown, Explanation}, infinity).
 
-init(Parent, HelperSup, Ref, Sock) ->
+init(Parent, HelperSup, Ref) ->
     ?LG_PROCESS_TYPE(reader),
-    RealSocket = rabbit_net:unwrap_socket(Sock),
-    rabbit_networking:accept_ack(Ref, RealSocket),
+    {ok, Sock} = rabbit_networking:handshake(Ref,
+        application:get_env(rabbit, proxy_protocol, false)),
     Deb = sys:debug_options([]),
     start_connection(Parent, HelperSup, Deb, Sock).