From 5ef64557b5b4c92224d68d1445cdb7846a76a0be Mon Sep 17 00:00:00 2001
From: Nelson Vides <nelson.vides@erlang-solutions.com>
Date: Tue, 20 Apr 2021 14:58:58 +0200
Subject: Exit gracefully on {error,closed} when reading the PROXY header

LH: Simplified the test a little.
---
 src/cowboy_clear.erl        | 16 +++++++++-------
 src/cowboy_tls.erl          | 16 +++++++++-------
 test/proxy_header_SUITE.erl | 24 ++++++++++++++++++++++++
 3 files changed, 42 insertions(+), 14 deletions(-)

diff --git a/src/cowboy_clear.erl b/src/cowboy_clear.erl
index 4f3a234..e4ca9da 100644
--- a/src/cowboy_clear.erl
+++ b/src/cowboy_clear.erl
@@ -33,13 +33,7 @@ start_link(Ref, Transport, Opts) ->
 
 -spec connection_process(pid(), ranch:ref(), module(), cowboy:opts()) -> ok.
 connection_process(Parent, Ref, Transport, Opts) ->
-	ProxyInfo = case maps:get(proxy_header, Opts, false) of
-		true ->
-			{ok, ProxyInfo0} = ranch:recv_proxy_header(Ref, 1000),
-			ProxyInfo0;
-		false ->
-			undefined
-	end,
+	ProxyInfo = get_proxy_info(Ref, Opts),
 	{ok, Socket} = ranch:handshake(Ref),
 	%% Use cowboy_http2 directly only when 'http' is missing.
 	%% Otherwise switch to cowboy_http2 from cowboy_http.
@@ -58,3 +52,11 @@ init(Parent, Ref, Socket, Transport, ProxyInfo, Opts, Protocol) ->
 		supervisor -> process_flag(trap_exit, true)
 	end,
 	Protocol:init(Parent, Ref, Socket, Transport, ProxyInfo, Opts).
+
+get_proxy_info(Ref, #{proxy_header := true}) ->
+	case ranch:recv_proxy_header(Ref, 1000) of
+		{ok, ProxyInfo} -> ProxyInfo;
+		{error, closed} -> exit({shutdown, closed})
+	end;
+get_proxy_info(_, _) ->
+	undefined.
diff --git a/src/cowboy_tls.erl b/src/cowboy_tls.erl
index c049ecb..4385cbc 100644
--- a/src/cowboy_tls.erl
+++ b/src/cowboy_tls.erl
@@ -33,13 +33,7 @@ start_link(Ref, Transport, Opts) ->
 
 -spec connection_process(pid(), ranch:ref(), module(), cowboy:opts()) -> ok.
 connection_process(Parent, Ref, Transport, Opts) ->
-	ProxyInfo = case maps:get(proxy_header, Opts, false) of
-		true ->
-			{ok, ProxyInfo0} = ranch:recv_proxy_header(Ref, 1000),
-			ProxyInfo0;
-		false ->
-			undefined
-	end,
+	ProxyInfo = get_proxy_info(Ref, Opts),
 	{ok, Socket} = ranch:handshake(Ref),
 	case ssl:negotiated_protocol(Socket) of
 		{ok, <<"h2">>} ->
@@ -54,3 +48,11 @@ init(Parent, Ref, Socket, Transport, ProxyInfo, Opts, Protocol) ->
 		supervisor -> process_flag(trap_exit, true)
 	end,
 	Protocol:init(Parent, Ref, Socket, Transport, ProxyInfo, Opts).
+
+get_proxy_info(Ref, #{proxy_header := true}) ->
+	case ranch:recv_proxy_header(Ref, 1000) of
+		{ok, ProxyInfo} -> ProxyInfo;
+		{error, closed} -> exit({shutdown, closed})
+	end;
+get_proxy_info(_, _) ->
+	undefined.
diff --git a/test/proxy_header_SUITE.erl b/test/proxy_header_SUITE.erl
index ce8b6cb..9d1ca2f 100644
--- a/test/proxy_header_SUITE.erl
+++ b/test/proxy_header_SUITE.erl
@@ -71,6 +71,30 @@ init_dispatch() ->
 
 %% Tests.
 
+fail_gracefully_on_disconnect(Config) ->
+	doc("Probing a port must not generate a crash"),
+	{ok, Socket} = gen_tcp:connect("localhost", config(port, Config),
+		[binary, {active, false}, {packet, raw}]),
+	timer:sleep(50),
+	Pid = case config(type, Config) of
+		tcp -> ct_helper:get_remote_pid_tcp(Socket);
+		%% We connect to a TLS port using a TCP socket so we need
+		%% to first obtain the remote pid of the TCP socket, which
+		%% is a TLS socket on the server, and then get the real
+		%% remote pid from its state.
+		ssl -> ct_helper:get_remote_pid_tls_state(ct_helper:get_remote_pid_tcp(Socket))
+	end,
+	Ref = erlang:monitor(process, Pid),
+	gen_tcp:close(Socket),
+	receive
+		{'DOWN', Ref, process, Pid, {shutdown, closed}} ->
+			ok;
+		{'DOWN', Ref, process, Pid, Reason} ->
+			error(Reason)
+	after 500 ->
+		error(timeout)
+	end.
+
 v1_proxy_header(Config) ->
 	doc("Confirm we can read the proxy header at the start of the connection."),
 	ProxyInfo = #{
-- 
cgit v1.2.3