From d44e7a16f7d2823cc658e39b5d953ba0850e47ba Mon Sep 17 00:00:00 2001
From: juhlig <j.uhlig@mailingwork.de>
Date: Tue, 17 Sep 2019 17:55:25 +0200
Subject: Enable multiple steps handshake

Also fix some Protocol:start_link/4 into start_link/3
left over in the documentation.
---
 src/ranch.erl | 74 ++++++++++++++++++++++++++++++++++++++++++++---------------
 1 file changed, 55 insertions(+), 19 deletions(-)

(limited to 'src/ranch.erl')

diff --git a/src/ranch.erl b/src/ranch.erl
index 08eebfc..89ca21e 100644
--- a/src/ranch.erl
+++ b/src/ranch.erl
@@ -22,6 +22,9 @@
 -export([child_spec/5]).
 -export([handshake/1]).
 -export([handshake/2]).
+-export([handshake_continue/1]).
+-export([handshake_continue/2]).
+-export([handshake_cancel/1]).
 -export([recv_proxy_header/2]).
 -export([remove_connection/1]).
 -export([get_status/1]).
@@ -197,28 +200,61 @@ child_spec(Ref, Transport, TransOpts0, Protocol, ProtoOpts) ->
 		Ref, Transport, TransOpts, Protocol, ProtoOpts
 	]}, type => supervisor}.
 
--spec handshake(ref()) -> {ok, ranch_transport:socket()}.
+-spec handshake(ref()) -> {ok, ranch_transport:socket()} | {continue, any()}.
 handshake(Ref) ->
-	handshake(Ref, []).
+	handshake1(Ref, undefined).
 
--spec handshake(ref(), any()) -> {ok, ranch_transport:socket()}.
+-spec handshake(ref(), any()) -> {ok, ranch_transport:socket()} | {continue, any()}.
 handshake(Ref, Opts) ->
-	receive {handshake, Ref, Transport, CSocket, HandshakeTimeout} ->
-		case Transport:handshake(CSocket, Opts, HandshakeTimeout) of
-			OK = {ok, _} ->
-				OK;
-			%% Garbage was most likely sent to the socket, don't error out.
-			{error, {tls_alert, _}} ->
-				ok = Transport:close(CSocket),
-				exit(normal);
-			%% Socket most likely stopped responding, don't error out.
-			{error, Reason} when Reason =:= timeout; Reason =:= closed ->
-				ok = Transport:close(CSocket),
-				exit(normal);
-			{error, Reason} ->
-				ok = Transport:close(CSocket),
-				error(Reason)
-		end
+	handshake1(Ref, {opts, Opts}).
+
+handshake1(Ref, Opts) ->
+	receive {handshake, Ref, Transport, CSocket, Timeout} ->
+		Handshake = handshake_transport(Transport, handshake, CSocket, Opts, Timeout),
+		handshake_result(Handshake, Ref, Transport, CSocket, Timeout)
+	end.
+
+-spec handshake_continue(ref()) -> {ok, ranch_transport:socket()}.
+handshake_continue(Ref) ->
+	handshake_continue1(Ref, undefined).
+
+-spec handshake_continue(ref(), any()) -> {ok, ranch_transport:socket()}.
+handshake_continue(Ref, Opts) ->
+	handshake_continue1(Ref, {opts, Opts}).
+
+handshake_continue1(Ref, Opts) ->
+	receive {handshake_continue, Ref, Transport, CSocket, Timeout} ->
+		Handshake = handshake_transport(Transport, handshake_continue, CSocket, Opts, Timeout),
+		handshake_result(Handshake, Ref, Transport, CSocket, Timeout)
+	end.
+
+handshake_transport(Transport, Fun, CSocket, undefined, Timeout) ->
+	Transport:Fun(CSocket, Timeout);
+handshake_transport(Transport, Fun, CSocket, {opts, Opts}, Timeout) ->
+	Transport:Fun(CSocket, Opts, Timeout).
+
+handshake_result(Result, Ref, Transport, CSocket, Timeout) ->
+	case Result of
+		OK = {ok, _} ->
+			OK;
+		{ok, CSocket2, Info} ->
+			self() ! {handshake_continue, Ref, Transport, CSocket2, Timeout},
+			{continue, Info};
+		{error, {tls_alert, _}} ->
+			ok = Transport:close(CSocket),
+			exit(normal);
+		{error, Reason} when Reason =:= timeout; Reason =:= closed ->
+			ok = Transport:close(CSocket),
+			exit(normal);
+		{error, Reason} ->
+			ok = Transport:close(CSocket),
+			error(Reason)
+	end.
+
+-spec handshake_cancel(ref()) -> ok.
+handshake_cancel(Ref) ->
+	receive {handshake_continue, Ref, Transport, CSocket, _} ->
+		Transport:handshake_cancel(CSocket)
 	end.
 
 %% Unlike handshake/2 this function always return errors because
-- 
cgit v1.2.3