Introduce Transport:handshake/1,2

This commit deprecates Transport:accept_ack/1 in favor of
a new forward-compatible function. Transport:handshake/1,2
will use ssl:handshake/2,3 from Ranch 2.0 onward.

18 files changed, 181 insertions, 76 deletions

diff --git a/doc/src/guide/protocols.asciidoc b/doc/src/guide/protocols.asciidoc
index 91f4b07..b20b22d 100644
--- a/doc/src/guide/protocols.asciidoc
+++ b/doc/src/guide/protocols.asciidoc
@@ -14,18 +14,18 @@ the call to `ranch:start_listener/5`. This callback must
 return `{ok, Pid}`, with `Pid` the pid of the new process.
 
 The newly started process can then freely initialize itself. However,
-it must call `ranch:accept_ack/1` before doing any socket operation.
+it must call `ranch:handshake/{1,2}` before doing any socket operation.
 This will ensure the connection process is the owner of the socket.
 It expects the listener's name as argument.
 
-.Acknowledge accepting the socket
+.Perform the socket handshake
 
 [source,erlang]
-ok = ranch:accept_ack(Ref).
+{ok, Socket} = ranch:handshake(Ref).
 
 If your protocol code requires specific socket options, you should
 set them while initializing your connection process, after
-calling `ranch:accept_ack/1`. You can use `Transport:setopts/2`
+calling `ranch:handshake/{1,2}`. You can use `Transport:setopts/2`
 for that purpose.
 
 Following is the complete protocol code for the example found
@@ -39,14 +39,14 @@ in `examples/tcp_echo/`.
 -behaviour(ranch_protocol).
 
 -export([start_link/4]).
--export([init/4]).
+-export([init/3]).
 
-start_link(Ref, Socket, Transport, Opts) ->
-	Pid = spawn_link(?MODULE, init, [Ref, Socket, Transport, Opts]),
+start_link(Ref, _Socket, Transport, Opts) ->
+	Pid = spawn_link(?MODULE, init, [Ref, Transport, Opts]),
 	{ok, Pid}.
 
-init(Ref, Socket, Transport, _Opts = []) ->
-	ok = ranch:accept_ack(Ref),
+init(Ref, Transport, _Opts = []) ->
+	{ok, Socket} = ranch:handshake(Ref),
 	loop(Socket, Transport).
 
 loop(Socket, Transport) ->
@@ -64,7 +64,7 @@ loop(Socket, Transport) ->
 Special processes like the ones that use the `gen_statem` or `gen_server`
 behaviours have the particularity of having their `start_link` call not
 return until the `init` function returns. This is problematic, because
-you won't be able to call `ranch:accept_ack/1` from the `init` callback
+you won't be able to call `ranch:handshake/{1,2}` from the `init` callback
 as this would cause a deadlock to happen.
 
 Use the `gen_statem:enter_loop/4` function. It allows you to start your process
@@ -84,12 +84,12 @@ the normal `gen_statem` execution loop.
 -export([init/1]).
 %% Exports of other gen_statem callbacks here.
 
-start_link(Ref, Socket, Transport, Opts) ->
-	{ok, proc_lib:spawn_link(?MODULE, init, [{Ref, Socket, Transport, Opts}])}.
+start_link(Ref, _Socket, Transport, Opts) ->
+	{ok, proc_lib:spawn_link(?MODULE, init, [{Ref, Transport, Opts}])}.
 
-init({Ref, Socket, Transport, _Opts = []}) ->
+init({Ref, Transport, _Opts = []}) ->
 	%% Perform any required state initialization here.
-	ok = ranch:accept_ack(Ref),
+	{ok, Socket} = ranch:handshake(Ref),
 	ok = Transport:setopts(Socket, [{active, once}]),
 	gen_statem:enter_loop(?MODULE, [], state_name, {state_data, Socket, Transport}).
 
diff --git a/doc/src/manual/ranch.asciidoc b/doc/src/manual/ranch.asciidoc
index 7fd18d0..2380646 100644
--- a/doc/src/manual/ranch.asciidoc
+++ b/doc/src/manual/ranch.asciidoc
@@ -47,7 +47,7 @@ Unique name used to refer to a listener.
 None of the options are required.
 
 ack_timeout (5000)::
-	Maximum allowed time for the `ranch:accept_ack/1` call to finish.
+	Maximum allowed time for the `ranch:handshake/{1,2}` call to finish.
 connection_type (worker)::
 	Type of process that will handle the connection.
 max_connections (1024)::
@@ -63,13 +63,7 @@ socket::
 
 === accept_ack(Ref) -> ok
 
-Ref = ref():: Listener name.
-
-Acknowledge that the connection is accepted.
-
-This function MUST be used by a connection process to inform
-Ranch that it initialized properly and let it perform any
-additional operations before the socket can be safely used.
+This function is deprecated in favor of `ranch:handshake/1,2`.
 
 === child_spec(Ref, NumAcceptors, Transport, TransOpts, Protocol, ProtoOpts) -> supervisor:child_spec()
 
@@ -127,6 +121,34 @@ TransOpts = any():: Current transport options.
 
 Return the transport options set for the given listener.
 
+=== handshake(Ref) -> {ok, Socket}
+
+Ref = ref():: Listener name.
+Socket = any():: Initialized socket.
+
+Acknowledge that the connection is accepted.
+Returns a socket that is ready to use.
+
+One of the `ranch:handshake/{1,2}` functions MUST be used
+by a connection process to inform Ranch that it initialized
+properly and let it perform any additional operations before
+the socket can be safely used.
+
+=== handshake(Ref, Opts) -> {ok, Socket}
+
+Ref = ref():: Listener name.
+Opts = any():: Initialization options.
+Socket = any():: Initialized socket.
+
+Acknowledge that the connection is accepted.
+Additional options can be provided for socket initialization.
+Returns a socket that is ready to use.
+
+One of the `ranch:handshake/{1,2}` functions MUST be used
+by a connection process to inform Ranch that it initialized
+properly and let it perform any additional operations before
+the socket can be safely used.
+
 === info() -> [{Ref, [{Key, Value}]}]
 
 Ref = ref():: Listener name.
@@ -272,4 +294,4 @@ Operator = '>' | '>=' | '==' | '=<' | '<':: Comparison operator.
 NumConnections = non_neg_integer():: Number of connections to wait for.
 
 Wait until the number of connections on the given listener matches
-the given operator and number of connections.
\ No newline at end of file
+the given operator and number of connections.
diff --git a/doc/src/manual/ranch_transport.asciidoc b/doc/src/manual/ranch_transport.asciidoc
index cbcea42..ed0ad70 100644
--- a/doc/src/manual/ranch_transport.asciidoc
+++ b/doc/src/manual/ranch_transport.asciidoc
@@ -26,22 +26,10 @@ Timeout = timeout():: Accept timeout.
 
 Accept a connection on the given listening socket.
 
-The `accept_ack` callback will be used to initialize the socket
+The `handshake` callback will be used to initialize the socket
 after accepting the connection. This is most useful when the
 transport is not raw TCP, like with SSL for example.
 
-=== accept_ack(CSocket, Timeout) -> ok
-
-CSocket = any():: Socket for this connection.
-Timeout = timeout():: Ack timeout.
-
-Perform post-accept initialization of the connection.
-
-This function will be called by connection processes
-before performing any socket operation. It allows
-transports that require extra initialization to perform
-their task and make the socket ready to use.
-
 === close(Socket) -> ok
 
 Socket = any():: Socket opened with listen/1 or accept/2.
@@ -60,6 +48,20 @@ perform operations on the socket, and that will receive
 messages from the socket when active mode is used. When
 the controlling process dies, the socket is closed.
 
+=== handshake(CSocket0, Options, Timeout) -> {ok, CSocket1}
+
+CSocket0 = any():: Uninitialized socket for this connection.
+Options = any():: Options for initialization.
+Timeout = timeout():: Handshake timeout.
+CSocket1 = any():: Initialized socket for this connection.
+
+Perform post-accept initialization of the connection.
+
+This function will be called by connection processes
+before performing any socket operation. It allows
+transports that require extra initialization to perform
+their task and return a socket that is ready to use.
+
 === listen(TransOpts) -> {ok, LSocket} | {error, atom()}
 
 TransOpts = any():: Transport options.
diff --git a/examples/tcp_echo/src/echo_protocol.erl b/examples/tcp_echo/src/echo_protocol.erl
index 6e1df40..d8944c6 100644
--- a/examples/tcp_echo/src/echo_protocol.erl
+++ b/examples/tcp_echo/src/echo_protocol.erl
@@ -4,14 +4,14 @@
 -behaviour(ranch_protocol).
 
 -export([start_link/4]).
--export([init/4]).
+-export([init/3]).
 
-start_link(Ref, Socket, Transport, Opts) ->
-	Pid = spawn_link(?MODULE, init, [Ref, Socket, Transport, Opts]),
+start_link(Ref, _Socket, Transport, Opts) ->
+	Pid = spawn_link(?MODULE, init, [Ref, Transport, Opts]),
 	{ok, Pid}.
 
-init(Ref, Socket, Transport, _Opts = []) ->
-	ok = ranch:accept_ack(Ref),
+init(Ref, Transport, _Opts = []) ->
+	{ok, Socket} = ranch:handshake(Ref),
 	loop(Socket, Transport).
 
 loop(Socket, Transport) ->
diff --git a/examples/tcp_reverse/src/reverse_protocol.erl b/examples/tcp_reverse/src/reverse_protocol.erl
index a37ca46..91a5b95 100644
--- a/examples/tcp_reverse/src/reverse_protocol.erl
+++ b/examples/tcp_reverse/src/reverse_protocol.erl
@@ -20,16 +20,16 @@
 
 %% API.
 
-start_link(Ref, Socket, Transport, Opts) ->
-	{ok, proc_lib:spawn_link(?MODULE, init, [{Ref, Socket, Transport, Opts}])}.
+start_link(Ref, _Socket, Transport, Opts) ->
+	{ok, proc_lib:spawn_link(?MODULE, init, [{Ref, Transport, Opts}])}.
 
 %% gen_statem.
 
 callback_mode() ->
 	state_functions.
 
-init({Ref, Socket, Transport, _Opts = []}) ->
-	ok = ranch:accept_ack(Ref),
+init({Ref, Transport, _Opts = []}) ->
+	{ok, Socket} = ranch:handshake(Ref),
 	ok = Transport:setopts(Socket, [{active, once}, {packet, line}]),
 	gen_statem:enter_loop(?MODULE, [], connected,
 		#state{socket=Socket, transport=Transport},
diff --git a/src/ranch.erl b/src/ranch.erl
index 9cd3809..287b1e3 100644
--- a/src/ranch.erl
+++ b/src/ranch.erl
@@ -22,6 +22,8 @@
 -export([child_spec/5]).
 -export([child_spec/6]).
 -export([accept_ack/1]).
+-export([handshake/1]).
+-export([handshake/2]).
 -export([remove_connection/1]).
 -export([get_status/1]).
 -export([get_addr/1]).
@@ -166,10 +168,20 @@ child_spec(Ref, NumAcceptors, Transport, TransOpts, Protocol, ProtoOpts)
 
 -spec accept_ack(ref()) -> ok.
 accept_ack(Ref) ->
-	receive {shoot, Ref, Transport, Socket, AckTimeout} ->
+	receive {handshake, Ref, Transport, Socket, AckTimeout} ->
 		Transport:accept_ack(Socket, AckTimeout)
 	end.
 
+-spec handshake(ref()) -> {ok, ranch_transport:socket()}.
+handshake(Ref) ->
+	handshake(Ref, []).
+
+-spec handshake(ref(), any()) -> {ok, ranch_transport:socket()}.
+handshake(Ref, Opts) ->
+	receive {handshake, Ref, Transport, Socket, AckTimeout} ->
+		Transport:handshake(Socket, Opts, AckTimeout)
+	end.
+
 -spec remove_connection(ref()) -> ok.
 remove_connection(Ref) ->
 	ConnsSup = ranch_server:get_connections_sup(Ref),
diff --git a/src/ranch_conns_sup.erl b/src/ranch_conns_sup.erl
index 812a6c3..72de2d7 100644
--- a/src/ranch_conns_sup.erl
+++ b/src/ranch_conns_sup.erl
@@ -115,9 +115,9 @@ loop(State=#state{parent=Parent, ref=Ref, conn_type=ConnType,
 		{?MODULE, start_protocol, To, Socket} ->
 			try Protocol:start_link(Ref, Socket, Transport, Opts) of
 				{ok, Pid} ->
-					shoot(State, CurConns, NbChildren, Sleepers, To, Socket, Pid, Pid);
+					handshake(State, CurConns, NbChildren, Sleepers, To, Socket, Pid, Pid);
 				{ok, SupPid, ProtocolPid} when ConnType =:= supervisor ->
-					shoot(State, CurConns, NbChildren, Sleepers, To, Socket, SupPid, ProtocolPid);
+					handshake(State, CurConns, NbChildren, Sleepers, To, Socket, SupPid, ProtocolPid);
 				Ret ->
 					To ! self(),
 					error_logger:error_msg(
@@ -219,11 +219,11 @@ loop(State=#state{parent=Parent, ref=Ref, conn_type=ConnType,
 			loop(State, CurConns, NbChildren, Sleepers)
 	end.
 
-shoot(State=#state{ref=Ref, transport=Transport, ack_timeout=AckTimeout, max_conns=MaxConns},
+handshake(State=#state{ref=Ref, transport=Transport, ack_timeout=AckTimeout, max_conns=MaxConns},
 		CurConns, NbChildren, Sleepers, To, Socket, SupPid, ProtocolPid) ->
 	case Transport:controlling_process(Socket, ProtocolPid) of
 		ok ->
-			ProtocolPid ! {shoot, Ref, Transport, Socket, AckTimeout},
+			ProtocolPid ! {handshake, Ref, Transport, Socket, AckTimeout},
 			put(SupPid, active),
 			CurConns2 = CurConns + 1,
 			if CurConns2 < MaxConns ->
diff --git a/src/ranch_ssl.erl b/src/ranch_ssl.erl
index ac96e45..d651e7a 100644
--- a/src/ranch_ssl.erl
+++ b/src/ranch_ssl.erl
@@ -15,6 +15,10 @@
 -module(ranch_ssl).
 -behaviour(ranch_transport).
 
+-ifdef(OTP_RELEASE).
+-compile({nowarn_deprecated_function, [{ssl, ssl_accept, 3}]}).
+-endif.
+
 -export([name/0]).
 -export([secure/0]).
 -export([messages/0]).
@@ -22,6 +26,7 @@
 -export([disallowed_listen_options/0]).
 -export([accept/2]).
 -export([accept_ack/2]).
+-export([handshake/3]).
 -export([connect/3]).
 -export([connect/4]).
 -export([recv/3]).
@@ -128,9 +133,14 @@ accept(LSocket, Timeout) ->
 
 -spec accept_ack(ssl:sslsocket(), timeout()) -> ok.
 accept_ack(CSocket, Timeout) ->
-	case ssl:ssl_accept(CSocket, Timeout) of
+	{ok, _} = handshake(CSocket, [], Timeout),
+	ok.
+
+-spec handshake(ssl:sslsocket(), opts(), timeout()) -> {ok, ssl:sslsocket()}.
+handshake(CSocket, Opts, Timeout) ->
+	case ssl:ssl_accept(CSocket, Opts, Timeout) of
 		ok ->
-			ok;
+			{ok, CSocket};
 		%% Garbage was most likely sent to the socket, don't error out.
 		{error, {tls_alert, _}} ->
 			ok = close(CSocket),
diff --git a/src/ranch_tcp.erl b/src/ranch_tcp.erl
index bc0c75d..58cbb37 100644
--- a/src/ranch_tcp.erl
+++ b/src/ranch_tcp.erl
@@ -22,6 +22,7 @@
 -export([disallowed_listen_options/0]).
 -export([accept/2]).
 -export([accept_ack/2]).
+-export([handshake/3]).
 -export([connect/3]).
 -export([connect/4]).
 -export([recv/3]).
@@ -100,9 +101,14 @@ accept(LSocket, Timeout) ->
 	gen_tcp:accept(LSocket, Timeout).
 
 -spec accept_ack(inet:socket(), timeout()) -> ok.
-accept_ack(_, _) ->
+accept_ack(CSocket, Timeout) ->
+	{ok, _} = handshake(CSocket, [], Timeout),
 	ok.
 
+-spec handshake(inet:socket(), opts(), timeout()) -> {ok, inet:socket()}.
+handshake(CSocket, _, _) ->
+	{ok, CSocket}.
+
 %% @todo Probably filter Opts?
 -spec connect(inet:ip_address() | inet:hostname(),
 	inet:port_number(), any())
diff --git a/src/ranch_transport.erl b/src/ranch_transport.erl
index b671cf5..9d3083f 100644
--- a/src/ranch_transport.erl
+++ b/src/ranch_transport.erl
@@ -17,6 +17,8 @@
 -export([sendfile/6]).
 
 -type socket() :: any().
+-export_type([socket/0]).
+
 -type opts() :: any().
 -type stats() :: any().
 -type sendfile_opts() :: [{chunk_size, non_neg_integer()}].
@@ -28,7 +30,7 @@
 -callback listen(opts()) -> {ok, socket()} | {error, atom()}.
 -callback accept(socket(), timeout())
 	-> {ok, socket()} | {error, closed | timeout | atom()}.
--callback accept_ack(socket(), timeout()) -> ok.
+-callback handshake(socket(), opts(), timeout()) -> {ok, socket()}.
 -callback connect(string(), inet:port_number(), opts())
 	-> {ok, socket()} | {error, atom()}.
 -callback connect(string(), inet:port_number(), opts(), timeout())
diff --git a/test/accept_ack_protocol.erl b/test/accept_ack_protocol.erl
new file mode 100644
index 0000000..3c36fc7
--- /dev/null
+++ b/test/accept_ack_protocol.erl
@@ -0,0 +1,22 @@
+-module(accept_ack_protocol).
+-behaviour(ranch_protocol).
+
+-export([start_link/4]).
+-export([init/4]).
+
+start_link(Ref, Socket, Transport, Opts) ->
+	Pid = spawn_link(?MODULE, init, [Ref, Socket, Transport, Opts]),
+	{ok, Pid}.
+
+init(Ref, Socket, Transport, _Opts = []) ->
+	ok = ranch:accept_ack(Ref),
+	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/acceptor_SUITE.erl b/test/acceptor_SUITE.erl
index f1d707d..0be5d40 100644
--- a/test/acceptor_SUITE.erl
+++ b/test/acceptor_SUITE.erl
@@ -32,6 +32,7 @@ groups() ->
 		tcp_active_echo,
 		tcp_echo,
 		tcp_graceful,
+		tcp_accept_ack,
 		tcp_inherit_options,
 		tcp_max_connections,
 		tcp_max_connections_and_beyond,
@@ -50,6 +51,7 @@ groups() ->
 		ssl_active_echo,
 		ssl_echo,
 		ssl_graceful,
+		ssl_accept_ack,
 		ssl_sni_echo,
 		ssl_sni_fail,
 		ssl_getopts_capability,
@@ -489,6 +491,20 @@ ssl_graceful(_) ->
 	{'EXIT', _} = begin catch ranch:get_port(Name) end,
 	ok.
 
+ssl_accept_ack(_) ->
+	doc("Ensure accept_ack works with SSL transport."),
+	Name = name(),
+	Opts = ct_helper:get_certs_from_ets(),
+	{ok, _} = ranch:start_listener(Name, ranch_ssl, Opts, accept_ack_protocol, []),
+	Port = ranch:get_port(Name),
+	{ok, Socket} = ssl:connect("localhost", Port, [binary, {active, false}, {packet, raw}]),
+	ok = ssl:send(Socket, <<"SSL transport accept_ack is working!">>),
+	{ok, <<"SSL transport accept_ack is working!">>} = ssl:recv(Socket, 36, 1000),
+	ok = ranch:stop_listener(Name),
+	{error, closed} = ssl:recv(Socket, 0, 1000),
+	{'EXIT', _} = begin catch ranch:get_port(Name) end,
+	ok.
+
 ssl_getopts_capability(_) ->
 	doc("Ensure getopts/2 capability."),
 	Name=name(),
@@ -641,6 +657,19 @@ tcp_graceful(_) ->
 	{'EXIT', _} = begin catch ranch:get_port(Name) end,
 	ok.
 
+tcp_accept_ack(_) ->
+	doc("Ensure accept_ack works with TCP transport."),
+	Name = name(),
+	{ok, _} = ranch:start_listener(Name, ranch_tcp, [], accept_ack_protocol, []),
+	Port = ranch:get_port(Name),
+	{ok, Socket} = gen_tcp:connect("localhost", Port, [binary, {active, false}, {packet, raw}]),
+	ok = gen_tcp:send(Socket, <<"TCP transport accept_ack is working!">>),
+	{ok, <<"TCP transport accept_ack is working!">>} = gen_tcp:recv(Socket, 36, 1000),
+	ok = ranch:stop_listener(Name),
+	{error, closed} = gen_tcp:recv(Socket, 0, 1000),
+	{'EXIT', _} = begin catch ranch:get_port(Name) end,
+	ok.
+
 tcp_inherit_options(_) ->
 	doc("Ensure TCP options are inherited in the protocol."),
 	Name = name(),
diff --git a/test/active_echo_protocol.erl b/test/active_echo_protocol.erl
index 6767141..37dd6db 100644
--- a/test/active_echo_protocol.erl
+++ b/test/active_echo_protocol.erl
@@ -2,14 +2,14 @@
 -behaviour(ranch_protocol).
 
 -export([start_link/4]).
--export([init/4]).
+-export([init/3]).
 
-start_link(Ref, Socket, Transport, Opts) ->
-	Pid = spawn_link(?MODULE, init, [Ref, Socket, Transport, Opts]),
+start_link(Ref, _Socket, Transport, Opts) ->
+	Pid = spawn_link(?MODULE, init, [Ref, Transport, Opts]),
 	{ok, Pid}.
 
-init(Ref, Socket, Transport, _Opts = []) ->
-	ok = ranch:accept_ack(Ref),
+init(Ref, Transport, _Opts = []) ->
+	{ok, Socket} = ranch:handshake(Ref),
 	loop(Socket, Transport).
 
 loop(Socket, Transport) ->
diff --git a/test/echo_protocol.erl b/test/echo_protocol.erl
index 8c454bc..640f5c6 100644
--- a/test/echo_protocol.erl
+++ b/test/echo_protocol.erl
@@ -2,14 +2,14 @@
 -behaviour(ranch_protocol).
 
 -export([start_link/4]).
--export([init/4]).
+-export([init/3]).
 
-start_link(Ref, Socket, Transport, Opts) ->
-	Pid = spawn_link(?MODULE, init, [Ref, Socket, Transport, Opts]),
+start_link(Ref, _Socket, Transport, Opts) ->
+	Pid = spawn_link(?MODULE, init, [Ref, Transport, Opts]),
 	{ok, Pid}.
 
-init(Ref, Socket, Transport, _Opts = []) ->
-	ok = ranch:accept_ack(Ref),
+init(Ref, Transport, _Opts = []) ->
+	{ok, Socket} = ranch:handshake(Ref),
 	loop(Socket, Transport).
 
 loop(Socket, Transport) ->
diff --git a/test/remove_conn_and_wait_protocol.erl b/test/remove_conn_and_wait_protocol.erl
index db688ff..59cedfc 100644
--- a/test/remove_conn_and_wait_protocol.erl
+++ b/test/remove_conn_and_wait_protocol.erl
@@ -9,7 +9,7 @@ start_link(Ref, _, _, [{remove, MaybeRemove, Timeout}]) ->
 	{ok, Pid}.
 
 init(Ref, MaybeRemove, Timeout) ->
-	ranch:accept_ack(Ref),
+	{ok, _} = ranch:handshake(Ref),
 	case MaybeRemove of
 		true ->
 			ranch:remove_connection(Ref);
diff --git a/test/sendfile_SUITE.erl b/test/sendfile_SUITE.erl
index 7bec5ac..5a6d2cb 100644
--- a/test/sendfile_SUITE.erl
+++ b/test/sendfile_SUITE.erl
@@ -280,7 +280,7 @@ sockets(Config) ->
 	end,
 	_ = spawn_link(Fun),
 	{ok, Server} = Transport:accept(LSocket, 500),
-	ok = Transport:accept_ack(Server, 500),
+	{ok, _} = Transport:handshake(Server, [], 500),
 	receive
 		{ok, Client} ->
 			ok = Transport:close(LSocket),
diff --git a/test/transport_capabilities_protocol.erl b/test/transport_capabilities_protocol.erl
index 0f8bf9d..ba5024a 100644
--- a/test/transport_capabilities_protocol.erl
+++ b/test/transport_capabilities_protocol.erl
@@ -2,14 +2,14 @@
 -behaviour(ranch_protocol).
 
 -export([start_link/4]).
--export([init/4]).
+-export([init/3]).
 
-start_link(Ref, Socket, Transport, Opts) ->
-	Pid = spawn_link(?MODULE, init, [Ref, Socket, Transport, Opts]),
+start_link(Ref, _Socket, Transport, Opts) ->
+	Pid = spawn_link(?MODULE, init, [Ref, Transport, Opts]),
 	{ok, Pid}.
 
-init(Ref, Socket, Transport, _Opts = []) ->
-	ok = ranch:accept_ack(Ref),
+init(Ref, Transport, _Opts = []) ->
+	{ok, Socket} = ranch:handshake(Ref),
 	loop(Socket, Transport).
 
 loop(Socket, Transport) ->
diff --git a/test/trap_exit_protocol.erl b/test/trap_exit_protocol.erl
index a0c4329..6100075 100644
--- a/test/trap_exit_protocol.erl
+++ b/test/trap_exit_protocol.erl
@@ -2,15 +2,15 @@
 -behaviour(ranch_protocol).
 
 -export([start_link/4]).
--export([init/4]).
+-export([init/3]).
 
-start_link(Ref, Socket, Transport, Opts) ->
-	Pid = spawn_link(?MODULE, init, [Ref, Socket, Transport, Opts]),
+start_link(Ref, _Socket, Transport, Opts) ->
+	Pid = spawn_link(?MODULE, init, [Ref, Transport, Opts]),
 	{ok, Pid}.
 
-init(Ref, Socket, Transport, _Opts = []) ->
+init(Ref, Transport, _Opts = []) ->
 	process_flag(trap_exit, true),
-	ok = ranch:accept_ack(Ref),
+	{ok, Socket} = ranch:handshake(Ref),
 	loop(Socket, Transport).
 
 loop(Socket, Transport) ->