Add ranch:recv_proxy_header/2

This is the function that should be called regardless of
TCP or TLS being used. The proper usage for this function is:

{ok, ProxyInfo} = ranch:recv_proxy_header(Ref, Timeout),
{ok, Socket} = ranch:handshake(Ref),
...

Ranch takes care of everything else under the hood. Transports
now need to have a Transport:recv_proxy_header/2 function. For
ranch_ssl the function gets the port from the sslsocket() record
and then calls ranch_tcp:recv_proxy_header/2 with it.

This means that two undocumented features are currently used for
this, but the interface is really nice so that's a sacrifice
worth doing. Also worth noting is that OTP 22 should have an
alternative for gen_tcp:unrecv/2 so the only real issue is about
the sslsocket() record at the moment.

7 files changed, 41 insertions, 38 deletions

diff --git a/src/ranch.erl b/src/ranch.erl
index 2664fd8..894f78b 100644
--- a/src/ranch.erl
+++ b/src/ranch.erl
@@ -25,6 +25,7 @@
 -export([accept_ack/1]).
 -export([handshake/1]).
 -export([handshake/2]).
+-export([recv_proxy_header/2]).
 -export([remove_connection/1]).
 -export([get_status/1]).
 -export([get_addr/1]).
@@ -256,6 +257,16 @@ handshake(Ref, Opts) ->
 		end
 	end.
 
+-spec recv_proxy_header(ref(), timeout())
+	-> {ok, ranch_proxy_header:proxy_info()}
+	| {error, closed | atom()}
+	| {error, protocol_error, atom()}.
+recv_proxy_header(Ref, Timeout) ->
+	receive HandshakeState={handshake, Ref, Transport, CSocket, _} ->
+		self() ! HandshakeState,
+		Transport:recv_proxy_header(CSocket, Timeout)
+	end.
+
 -spec remove_connection(ref()) -> ok.
 remove_connection(Ref) ->
 	ConnsSup = ranch_server:get_connections_sup(Ref),
diff --git a/src/ranch_ssl.erl b/src/ranch_ssl.erl
index b232039..03eb5ee 100644
--- a/src/ranch_ssl.erl
+++ b/src/ranch_ssl.erl
@@ -30,6 +30,7 @@
 -export([connect/3]).
 -export([connect/4]).
 -export([recv/3]).
+-export([recv_proxy_header/2]).
 -export([send/2]).
 -export([sendfile/2]).
 -export([sendfile/4]).
@@ -169,6 +170,19 @@ connect(Host, Port, Opts, Timeout) when is_integer(Port) ->
 recv(Socket, Length, Timeout) ->
 	ssl:recv(Socket, Length, Timeout).
 
+-spec recv_proxy_header(ssl:sslsocket(), timeout())
+	-> {ok, ranch_proxy_header:proxy_info()}
+	| {error, closed | atom()}
+	| {error, protocol_error, atom()}.
+recv_proxy_header(SSLSocket, Timeout) ->
+	%% There's currently no documented way to perform a TCP recv
+	%% on an sslsocket(), even before the TLS handshake. However
+	%% nothing prevents us from retrieving the TCP socket and using
+	%% it. Since it's an undocumented interface this may however
+	%% make forward-compatibility more difficult.
+	{sslsocket, {gen_tcp, TCPSocket, _, _}, _} = SSLSocket,
+	ranch_tcp:recv_proxy_header(TCPSocket, Timeout).
+
 -spec send(ssl:sslsocket(), iodata()) -> ok | {error, atom()}.
 send(Socket, Packet) ->
 	ssl:send(Socket, Packet).
diff --git a/src/ranch_tcp.erl b/src/ranch_tcp.erl
index ba77308..b7ece5b 100644
--- a/src/ranch_tcp.erl
+++ b/src/ranch_tcp.erl
@@ -133,7 +133,9 @@ recv(Socket, Length, Timeout) ->
 	gen_tcp:recv(Socket, Length, Timeout).
 
 -spec recv_proxy_header(inet:socket(), timeout())
-	-> {ok, any()} | {error, closed | atom()} | {error, protocol_error, atom()}.
+	-> {ok, ranch_proxy_header:proxy_info()}
+	| {error, closed | atom()}
+	| {error, protocol_error, atom()}.
 recv_proxy_header(Socket, Timeout) ->
 	case recv(Socket, 0, Timeout) of
 		{ok, Data} ->
diff --git a/src/ranch_transport.erl b/src/ranch_transport.erl
index 5da5d96..486c6d6 100644
--- a/src/ranch_transport.erl
+++ b/src/ranch_transport.erl
@@ -37,6 +37,10 @@
 	-> {ok, socket()} | {error, atom()}.
 -callback recv(socket(), non_neg_integer(), timeout())
 	-> {ok, any()} | {error, closed | timeout | atom()}.
+-callback recv_proxy_header(socket(), timeout())
+	-> {ok, ranch_proxy_header:proxy_info()}
+	| {error, closed | atom()}
+	| {error, protocol_error, atom()}.
 -callback send(socket(), iodata()) -> ok | {error, atom()}.
 -callback sendfile(socket(), file:name_all() | file:fd())
 	-> {ok, non_neg_integer()} | {error, atom()}.
diff --git a/test/proxy_header_SUITE.erl b/test/proxy_header_SUITE.erl
index b2308f8..6a80635 100644
--- a/test/proxy_header_SUITE.erl
+++ b/test/proxy_header_SUITE.erl
@@ -209,21 +209,17 @@ recv_v2_local_header_ssl_extra_data(_) ->
 	do_proxy_header_ssl(Name, ProxyInfo, <<"HELLO">>, <<"TCP Ranch is working!">>).
 
 do_proxy_header_ssl(Name, ProxyInfo, Data1, Data2) ->
+	Opts = ct_helper:get_certs_from_ets(),
 	{ok, _} = ranch:start_listener(Name,
-		ranch_tcp, #{},
-		proxy_protocol_ssl, []),
+		ranch_ssl, Opts,
+		proxy_protocol, []),
 	Port = ranch:get_port(Name),
 	{ok, Socket0} = gen_tcp:connect("localhost", Port, [binary, {active, false}, {packet, raw}]),
 	ok = gen_tcp:send(Socket0, [ranch_proxy_header:header(ProxyInfo)]),
-	%% This timeout is necessary to avoid a race condition when trying
-	%% to obtain the pid of the test case from the protocol. The race
-	%% condition is due to the TLS upgrade which changes the process
-	%% owning the socket.
-	timer:sleep(100),
 	{ok, Socket} = ssl:connect(Socket0, [], 1000),
 	ok = ssl:send(Socket, Data1),
 	receive
-		{proxy_protocol_ssl, ProxyInfo} ->
+		{proxy_protocol, ProxyInfo} ->
 			ok
 	after 2000 ->
 		error(timeout)
diff --git a/test/proxy_protocol.erl b/test/proxy_protocol.erl
index 6a161c7..9c679e3 100644
--- a/test/proxy_protocol.erl
+++ b/test/proxy_protocol.erl
@@ -9,9 +9,12 @@ start_link(Ref, _Socket, Transport, Opts) ->
 	{ok, Pid}.
 
 init(Ref, Transport, _Opts = []) ->
+	{ok, ProxyInfo} = ranch:recv_proxy_header(Ref, 1000),
 	{ok, Socket} = ranch:handshake(Ref),
-	{ok, ProxyInfo} = Transport:recv_proxy_header(Socket, 1000),
-	Pid = ct_helper:get_remote_pid_tcp(Socket),
+	Pid = case Transport of
+		ranch_tcp -> ct_helper:get_remote_pid_tcp(Socket);
+		ranch_ssl -> ct_helper:get_remote_pid_tls(Socket)
+	end,
 	Pid ! {?MODULE, ProxyInfo},
 	loop(Socket, Transport).
 
diff --git a/test/proxy_protocol_ssl.erl b/test/proxy_protocol_ssl.erl
deleted file mode 100644
index d850800..0000000
--- a/test/proxy_protocol_ssl.erl
+++ /dev/null
@@ -1,27 +0,0 @@
--module(proxy_protocol_ssl).
--behaviour(ranch_protocol).
-
--export([start_link/4]).
--export([init/3]).
-
-start_link(Ref, _Socket, Transport, Opts) ->
-	Pid = spawn_link(?MODULE, init, [Ref, Transport, Opts]),
-	{ok, Pid}.
-
-init(Ref, Transport, _Opts = []) ->
-	{ok, Socket} = ranch:handshake(Ref),
-	{ok, ProxyInfo} = Transport:recv_proxy_header(Socket, 1000),
-	Pid = ct_helper:get_remote_pid_tcp(Socket),
-	Pid ! {?MODULE, ProxyInfo},
-	Opts = ct_helper:get_certs_from_ets(),
-	{ok, SslSocket} = ranch_ssl:handshake(Socket, Opts, 1000),
-	loop(SslSocket, ranch_ssl).
-
-loop(Socket, Transport) ->
-	case Transport:recv(Socket, 0, 5000) of
-		{ok, Data} ->
-			_ = Transport:send(Socket, Data),
-			loop(Socket, Transport);
-		_ ->
-			ok = Transport:close(Socket)
-	end.