3 files changed, 9 insertions, 37 deletions
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.