From 75a08772f0ffddfed0441bfdc7e7f9a5adb3862f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Pawe=C5=82=20Chmielowski?= <pawel@process-one.net>
Date: Fri, 23 Aug 2024 18:07:59 +0200
Subject: [PATCH] Add finish_hanshake function

---
 c_src/fast_tls.c |  8 ++++----
 src/fast_tls.erl | 38 +++++++++++++++++++++++++++++++++++---
 2 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/c_src/fast_tls.c b/c_src/fast_tls.c
index d109db3..f180c25 100644
--- a/c_src/fast_tls.c
+++ b/c_src/fast_tls.c
@@ -1102,7 +1102,7 @@ loop_nif(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
                 if (res == 2) {
                     return err_term;
                 }
-                return return_read_write(env, state, bytes_to_read, enif_make_atom(env, "ok"));
+                return return_read_write(env, state, bytes_to_read, enif_make_atom(env, "init"));
             } else {
                 res = do_send_queue(env, state, &err_term, &to_send);
                 if (res == 2) {
@@ -1114,8 +1114,8 @@ loop_nif(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
                     reason == SSL_R_UNKNOWN_PROTOCOL ||
                     reason == SSL_R_UNEXPECTED_MESSAGE ||
                     reason == SSL_R_WRONG_VERSION_NUMBER ||
-		    reason == SSL_R_HTTP_REQUEST ||
-		    reason == SSL_R_HTTPS_PROXY_REQUEST)
+                    reason == SSL_R_HTTP_REQUEST ||
+                    reason == SSL_R_HTTPS_PROXY_REQUEST)
                     /* Do not report badly formed Client Hello */
                     err_term = ERR_T(enif_make_atom(env, "closed"));
                 else if (state->sni_error)
@@ -1130,7 +1130,7 @@ loop_nif(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
             if (res == 2) {
                 return err_term;
             }
-            return return_read_write(env, state, bytes_to_read, enif_make_atom(env, "ok"));
+            return return_read_write(env, state, bytes_to_read, enif_make_atom(env, "init"));
         }
     }
 
diff --git a/src/fast_tls.erl b/src/fast_tls.erl
index 4c2eceb..a4eac4c 100644
--- a/src/fast_tls.erl
+++ b/src/fast_tls.erl
@@ -41,7 +41,7 @@
          add_certfile/2, get_certfile/1, delete_certfile/1,
          clear_cache/0, get_negotiated_cipher/1,
          get_tls_last_message/2, set_fips_mode/1, get_fips_mode/0,
-         get_tls_cb_exporter/1, p12_to_pem/2]).
+         get_tls_cb_exporter/1, p12_to_pem/2, finish_handshake/2]).
 
 -ifdef(TEST).
 -include_lib("eunit/include/eunit.hrl").
@@ -236,9 +236,9 @@ loop(#tlssock{tcpsock = TCPSocket,
     try loop_nif(Port, ToSend, Received, Length) of
         {error, _} = Err ->
             Err;
-        {ok, <<>>, Decrypted} ->
+        {Tag, <<>>, Decrypted} when Tag == ok; Tag == init ->
             {ok, <<DecBuf/binary, Decrypted/binary>>};
-        {ok, ToWrite, Decrypted} ->
+        {Tag, ToWrite, Decrypted} when Tag == ok; Tag == init ->
             case gen_tcp:send(TCPSocket, ToWrite) of
                 ok ->
                     loop(Socket, <<>>, <<>>, <<DecBuf/binary, Decrypted/binary>>,
@@ -282,6 +282,38 @@ recv_and_loop(#tlssock{tcpsock = TCPSocket} = Socket,
             end
     end.
 
+-spec finish_handshake(tls_socket(), timeout()) ->
+    {error, inet:posix() | binary()} | ok.
+finish_handshake(#tlssock{tcpsock = TCPSocket, tlsport = Port}, Timeout) ->
+    OurLoop = fun OurLoop(Received) ->
+        try loop_nif(Port, <<>>, Received, 0) of
+            {error, _} = Err ->
+                Err;
+            {Tag, ToWrite, <<>>} when Tag == ok; Tag == init ->
+                case gen_tcp:send(TCPSocket, ToWrite) of
+                    ok when Tag == init ->
+                        case gen_tcp:recv(TCPSocket, 0, Timeout) of
+                            {ok, Received2} ->
+                                OurLoop(Received2);
+                            {error, _} = Err ->
+                                Err
+                        end;
+                    ok ->
+                        ok;
+                    {error, _} = Err ->
+                        Err
+                end;
+            {Tag, _ToWrite, _Data} when Tag == ok; Tag == init ->
+                {error, too_much_data_received};
+            {{error, _} = Err, ToWrite, _} ->
+                _ = gen_tcp:send(TCPSocket, ToWrite),
+                Err
+        catch error:badarg ->
+            {error, einval}
+        end
+    end,
+    OurLoop(<<>>).
+
 -spec send(tls_socket(), binary()) ->
     ok | {error, inet:posix() | binary() | timeout}.
 send(Socket, Packet) ->