From d2720842a63dc7bc6ac01d9e2866bfa78cb39aa5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lo=C3=AFc=20Hoguin?= <essen@ninenines.eu>
Date: Tue, 23 Oct 2018 12:10:53 +0200
Subject: 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.
---
 src/ranch.erl           | 11 +++++++++++
 src/ranch_ssl.erl       | 14 ++++++++++++++
 src/ranch_tcp.erl       |  4 +++-
 src/ranch_transport.erl |  4 ++++
 4 files changed, 32 insertions(+), 1 deletion(-)

(limited to 'src')

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()}.
-- 
cgit v1.2.3