On successful accept/connect - also return the final Handshake state

src/enoise.erl

@@ -126,7 +126,7 @@ handshake(Options, Role, ComState) ->
 %% @end
 -spec connect(TcpSock :: gen_tcp:socket(),
               Options :: noise_options()) ->
-                    {ok, noise_socket()} | {error, term()}.
+                    {ok, noise_socket(), enoise_hs_state:state()} | {error, term()}.
 connect(TcpSock, Options) ->
     tcp_handshake(TcpSock, initiator, Options).
 
@@ -236,9 +236,9 @@ tcp_handshake(TcpSock, Role, Options) ->
                           state    => {TcpSock, Active, <<>>} },
 
             case handshake(Options, Role, ComState) of
-                {ok, #{ rx := Rx, tx := Tx }, #{ state := {_, _, Buf} }} ->
+                {ok, #{ rx := Rx, tx := Tx, final_state := FState }, #{ state := {_, _, Buf} }} ->
                     {ok, Pid} = enoise_connection:start_link(TcpSock, Rx, Tx, self(), {Active, Buf}),
-                    {ok, #enoise{ pid = Pid }};
+                    {ok, #enoise{ pid = Pid }, FState};
                 Err = {error, _} ->
                     Err
             end;

src/enoise_hs_state.erl

@@ -8,7 +8,12 @@
 
 -module(enoise_hs_state).
 
--export([finalize/1, init/4, next_message/1, read_message/2, write_message/2]).
+-export([ finalize/1
+        , init/4
+        , next_message/1
+        , read_message/2
+        , remote_keys/1
+        , write_message/2]).
 
 -include("enoise.hrl").
 
@@ -49,12 +54,13 @@ init(Protocol, Role, Prologue, {S, E, RS, RE}) ->
                    ({in, [e]}, HS0)  -> mix_hash(HS0, enoise_keypair:pubkey(RE))
                 end, HS, PreMsgs).
 
-finalize(#noise_hs{ msgs = [], ss = SS, role = Role }) ->
+finalize(HS = #noise_hs{ msgs = [], ss = SS, role = Role }) ->
     {C1, C2} = enoise_sym_state:split(SS),
     HSHash   = enoise_sym_state:h(SS),
+    Final    = #{ hs_hash => HSHash, final_state => HS },
     case Role of
-        initiator -> {ok, #{ tx => C1, rx => C2, hs_hash => HSHash }};
+        initiator -> {ok, Final#{ tx => C1, rx => C2 }};
-        responder -> {ok, #{ rx => C1, tx => C2, hs_hash => HSHash }}
+        responder -> {ok, Final#{ rx => C1, tx => C2 }}
     end;
 finalize(_) ->
     error({bad_state, finalize}).
@@ -72,6 +78,9 @@ read_message(HS = #noise_hs{ msgs = [{in, Msg} | Msgs] }, Message) ->
     {HS1, RestBuf1} = read_message(HS#noise_hs{ msgs = Msgs }, Msg, Message),
     decrypt_and_hash(HS1, RestBuf1).
 
+remote_keys(#noise_hs{ rs = RS }) ->
+    RS.
+
 write_message(HS, [], MsgBuf) ->
     {HS, MsgBuf};
 write_message(HS, [Token | Tokens], MsgBuf0) ->

test/enoise_tests.erl

@@ -101,7 +101,7 @@ noise_test(Conf, SKP, CKP) ->
     {ok, TcpSock} = gen_tcp:connect("localhost", Port, [{active, once}, binary, {reuseaddr, true}], 100),
 
     Opts = [{noise, Protocol}, {s, CKP}] ++ [{rs, SKP} || need_rs(initiator, Conf) ],
-    {ok, EConn} = enoise:connect(TcpSock, Opts),
+    {ok, EConn, _} = enoise:connect(TcpSock, Opts),
 
     ok = enoise:send(EConn, <<"Hello World!">>),
     receive
@@ -131,7 +131,7 @@ echo_srv(Port, Protocol, SKP, CPub) ->
     {ok, TcpSock} = gen_tcp:accept(LSock, 500),
 
     Opts = [{noise, Protocol}, {s, SKP}] ++  [{rs, CPub} || need_rs(responder, Protocol)],
-    {ok, EConn} = enoise:accept(TcpSock, Opts),
+    {ok, EConn, _} = enoise:accept(TcpSock, Opts),
 
     gen_tcp:close(LSock),