Add ranch_tcp:recv_proxy_header/2

This uses the undocumented function gen_tcp:unrecv/2.

Tests have been added for both gen_tcp and ssl connections,
including sending data in the same first packet, at least
for gen_tcp (ssl tests may or may not end up buffering some
of the TLS handshake before the recv call, but there's no
guarantees).

4 files changed, 311 insertions, 0 deletions

diff --git a/src/ranch_tcp.erl b/src/ranch_tcp.erl
index d5517a5..ba77308 100644
--- a/src/ranch_tcp.erl
+++ b/src/ranch_tcp.erl
@@ -26,6 +26,7 @@
 -export([connect/3]).
 -export([connect/4]).
 -export([recv/3]).
+-export([recv_proxy_header/2]).
 -export([send/2]).
 -export([sendfile/2]).
 -export([sendfile/4]).
@@ -131,6 +132,28 @@ connect(Host, Port, Opts, Timeout) when is_integer(Port) ->
 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()}.
+recv_proxy_header(Socket, Timeout) ->
+	case recv(Socket, 0, Timeout) of
+		{ok, Data} ->
+			case ranch_proxy_header:parse(Data) of
+				{ok, ProxyInfo, <<>>} ->
+					{ok, ProxyInfo};
+				{ok, ProxyInfo, Rest} ->
+					case gen_tcp:unrecv(Socket, Rest) of
+						ok ->
+							{ok, ProxyInfo};
+						Error ->
+							Error
+					end;
+				{error, HumanReadable} ->
+					{error, protocol_error, HumanReadable}
+			end;
+		Error ->
+			Error
+	end.
+
 -spec send(inet:socket(), iodata()) -> ok | {error, atom()}.
 send(Socket, Packet) ->
 	gen_tcp:send(Socket, Packet).
diff --git a/test/proxy_header_SUITE.erl b/test/proxy_header_SUITE.erl
new file mode 100644
index 0000000..b2308f8
--- /dev/null
+++ b/test/proxy_header_SUITE.erl
@@ -0,0 +1,236 @@
+%% Copyright (c) 2018, Loïc Hoguin <[email protected]>
+%%
+%% Permission to use, copy, modify, and/or distribute this software for any
+%% purpose with or without fee is hereby granted, provided that the above
+%% copyright notice and this permission notice appear in all copies.
+%%
+%% THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
+%% WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+%% MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
+%% ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+%% WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
+%% ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
+%% OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+
+-module(proxy_header_SUITE).
+-compile(export_all).
+-compile(nowarn_export_all).
+
+-import(ct_helper, [doc/1]).
+-import(ct_helper, [name/0]).
+
+%% ct.
+
+all() ->
+	ct_helper:all(?MODULE).
+
+%% Tests.
+
+recv_v1_proxy_header_tcp(_) ->
+	doc("Confirm we can read the proxy header at the start of the connection."),
+	Name = name(),
+	ProxyInfo = #{
+		version => 1,
+		command => proxy,
+		transport_family => ipv4,
+		transport_protocol => stream,
+		src_address => {127, 0, 0, 1},
+		src_port => 444,
+		dest_address => {192, 168, 0, 1},
+		dest_port => 443
+	},
+	do_proxy_header_tcp(Name, ProxyInfo, <<>>, <<"TCP Ranch is working!">>).
+
+recv_v1_proxy_header_tcp_extra_data(_) ->
+	doc("Confirm we can read the proxy header at the start of the connection "
+		"and that the extra data in the first packet can be read afterwards."),
+	Name = name(),
+	ProxyInfo = #{
+		version => 1,
+		command => proxy,
+		transport_family => ipv4,
+		transport_protocol => stream,
+		src_address => {127, 0, 0, 1},
+		src_port => 444,
+		dest_address => {192, 168, 0, 1},
+		dest_port => 443
+	},
+	do_proxy_header_tcp(Name, ProxyInfo, <<"HELLO">>, <<"TCP Ranch is working!">>).
+
+recv_v2_proxy_header_tcp(_) ->
+	doc("Confirm we can read the proxy header at the start of the connection."),
+	Name = name(),
+	ProxyInfo = #{
+		version => 2,
+		command => proxy,
+		transport_family => ipv4,
+		transport_protocol => stream,
+		src_address => {127, 0, 0, 1},
+		src_port => 444,
+		dest_address => {192, 168, 0, 1},
+		dest_port => 443
+	},
+	do_proxy_header_tcp(Name, ProxyInfo, <<>>, <<"TCP Ranch is working!">>).
+
+recv_v2_proxy_header_tcp_extra_data(_) ->
+	doc("Confirm we can read the proxy header at the start of the connection "
+		"and that the extra data in the first packet can be read afterwards."),
+	Name = name(),
+	ProxyInfo = #{
+		version => 2,
+		command => proxy,
+		transport_family => ipv4,
+		transport_protocol => stream,
+		src_address => {127, 0, 0, 1},
+		src_port => 444,
+		dest_address => {192, 168, 0, 1},
+		dest_port => 443
+	},
+	do_proxy_header_tcp(Name, ProxyInfo, <<"HELLO">>, <<"TCP Ranch is working!">>).
+
+recv_v2_local_header_tcp(_) ->
+	doc("Confirm we can read the proxy header at the start of the connection."),
+	Name = name(),
+	ProxyInfo = #{
+		version => 2,
+		command => local
+	},
+	do_proxy_header_tcp(Name, ProxyInfo, <<>>, <<"TCP Ranch is working!">>).
+
+recv_v2_local_header_tcp_extra_data(_) ->
+	doc("Confirm we can read the proxy header at the start of the connection "
+		"and that the extra data in the first packet can be read afterwards."),
+	Name = name(),
+	ProxyInfo = #{
+		version => 2,
+		command => local
+	},
+	do_proxy_header_tcp(Name, ProxyInfo, <<"HELLO">>, <<"TCP Ranch is working!">>).
+
+do_proxy_header_tcp(Name, ProxyInfo, Data1, Data2) ->
+	{ok, _} = ranch:start_listener(Name,
+		ranch_tcp, #{},
+		proxy_protocol, []),
+	Port = ranch:get_port(Name),
+	{ok, Socket} = gen_tcp:connect("localhost", Port, [binary, {active, false}, {packet, raw}]),
+	ok = gen_tcp:send(Socket, [ranch_proxy_header:header(ProxyInfo), Data1]),
+	receive
+		{proxy_protocol, ProxyInfo} ->
+			ok
+	after 2000 ->
+		error(timeout)
+	end,
+	ok = gen_tcp:send(Socket, Data2),
+	Len1 = byte_size(Data1),
+	Len2 = byte_size(Data2),
+	{ok, <<Data1:Len1/binary, Data2/binary>>} = gen_tcp:recv(Socket, Len1 + Len2, 1000),
+	ok = ranch:stop_listener(Name),
+	ok.
+
+recv_v1_proxy_header_ssl(_) ->
+	doc("Confirm we can read the proxy header at the start of the connection."),
+	Name = name(),
+	ProxyInfo = #{
+		version => 1,
+		command => proxy,
+		transport_family => ipv4,
+		transport_protocol => stream,
+		src_address => {127, 0, 0, 1},
+		src_port => 444,
+		dest_address => {192, 168, 0, 1},
+		dest_port => 443
+	},
+	do_proxy_header_ssl(Name, ProxyInfo, <<>>, <<"TCP Ranch is working!">>).
+
+recv_v1_proxy_header_ssl_extra_data(_) ->
+	doc("Confirm we can read the proxy header at the start of the connection "
+		"and that the extra data in the first packet can be read afterwards."),
+	Name = name(),
+	ProxyInfo = #{
+		version => 1,
+		command => proxy,
+		transport_family => ipv4,
+		transport_protocol => stream,
+		src_address => {127, 0, 0, 1},
+		src_port => 444,
+		dest_address => {192, 168, 0, 1},
+		dest_port => 443
+	},
+	do_proxy_header_ssl(Name, ProxyInfo, <<"HELLO">>, <<"TCP Ranch is working!">>).
+
+recv_v2_proxy_header_ssl(_) ->
+	doc("Confirm we can read the proxy header at the start of the connection."),
+	Name = name(),
+	ProxyInfo = #{
+		version => 2,
+		command => proxy,
+		transport_family => ipv4,
+		transport_protocol => stream,
+		src_address => {127, 0, 0, 1},
+		src_port => 444,
+		dest_address => {192, 168, 0, 1},
+		dest_port => 443
+	},
+	do_proxy_header_ssl(Name, ProxyInfo, <<>>, <<"TCP Ranch is working!">>).
+
+recv_v2_proxy_header_ssl_extra_data(_) ->
+	doc("Confirm we can read the proxy header at the start of the connection "
+		"and that the extra data in the first packet can be read afterwards."),
+	Name = name(),
+	ProxyInfo = #{
+		version => 2,
+		command => proxy,
+		transport_family => ipv4,
+		transport_protocol => stream,
+		src_address => {127, 0, 0, 1},
+		src_port => 444,
+		dest_address => {192, 168, 0, 1},
+		dest_port => 443
+	},
+	do_proxy_header_ssl(Name, ProxyInfo, <<"HELLO">>, <<"TCP Ranch is working!">>).
+
+recv_v2_local_header_ssl(_) ->
+	doc("Confirm we can read the proxy header at the start of the connection."),
+	Name = name(),
+	ProxyInfo = #{
+		version => 2,
+		command => local
+	},
+	do_proxy_header_ssl(Name, ProxyInfo, <<>>, <<"TCP Ranch is working!">>).
+
+recv_v2_local_header_ssl_extra_data(_) ->
+	doc("Confirm we can read the proxy header at the start of the connection "
+		"and that the extra data in the first packet can be read afterwards."),
+	Name = name(),
+	ProxyInfo = #{
+		version => 2,
+		command => local
+	},
+	do_proxy_header_ssl(Name, ProxyInfo, <<"HELLO">>, <<"TCP Ranch is working!">>).
+
+do_proxy_header_ssl(Name, ProxyInfo, Data1, Data2) ->
+	{ok, _} = ranch:start_listener(Name,
+		ranch_tcp, #{},
+		proxy_protocol_ssl, []),
+	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} ->
+			ok
+	after 2000 ->
+		error(timeout)
+	end,
+	ok = ssl:send(Socket, Data2),
+	Len1 = byte_size(Data1),
+	Len2 = byte_size(Data2),
+	{ok, <<Data1:Len1/binary, Data2/binary>>} = ssl:recv(Socket, Len1 + Len2, 1000),
+	ok = ranch:stop_listener(Name),
+	ok.
diff --git a/test/proxy_protocol.erl b/test/proxy_protocol.erl
new file mode 100644
index 0000000..3b18349
--- /dev/null
+++ b/test/proxy_protocol.erl
@@ -0,0 +1,25 @@
+-module(proxy_protocol).
+-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},
+	loop(Socket, Transport).
+
+loop(Socket, Transport) ->
+	case Transport:recv(Socket, 0, 5000) of
+		{ok, Data} ->
+			Transport:send(Socket, Data),
+			loop(Socket, Transport);
+		_ ->
+			ok = Transport:close(Socket)
+	end.
diff --git a/test/proxy_protocol_ssl.erl b/test/proxy_protocol_ssl.erl
new file mode 100644
index 0000000..a25b297
--- /dev/null
+++ b/test/proxy_protocol_ssl.erl
@@ -0,0 +1,27 @@
+-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.