Refactor handshake flow control

This commit is contained in:
Hans Svensson
2018-03-02 15:42:43 +01:00
parent 4414c0d7a2
commit bc8ebc7ec6
2 changed files with 36 additions and 64 deletions
+22 -53
View File
@@ -32,35 +32,23 @@
%%====================================================================
%% API functions
%%====================================================================
connect(Address, Port, Options) ->
connect(Address, Port, Options, infinity).
connect(TcpSock, Options) ->
do_handshake(TcpSock, initiator, Options).
accept(TcpSock, Options) ->
do_handshake(TcpSock, responder, Options).
connect(Address, Port, Options, Timeout) ->
case initiate_handshake(initiator, Options) of
{ok, HS} ->
TcpOpts = enoise_opts:tcp_opts(Options),
case gen_tcp:connect(Address, Port, TcpOpts, Timeout) of
{ok, TcpSock} ->
do_handshake(TcpSock, HS, Options);
Err = {error, _Reason} ->
Err
end;
Err = {error, _Reason} ->
Err
end.
send(E = #enoise{ tcp_sock = TcpSock, rx = RX0 }, Msg0) ->
{ok, RX1, Msg1} = enoise_cipher_state:encrypt_with_ad(RX0, <<>>, Msg0),
send(E = #enoise{ tcp_sock = TcpSock, tx = TX0 }, Msg0) ->
{ok, TX1, Msg1} = enoise_cipher_state:encrypt_with_ad(TX0, <<>>, Msg0),
gen_tcp:send(TcpSock, <<(byte_size(Msg1)):16, Msg1/binary>>),
E#enoise{ rx = RX1 }.
E#enoise{ tx = TX1 }.
recv(E = #enoise{ tcp_sock = TcpSock, tx = TX0 }) ->
recv(E = #enoise{ tcp_sock = TcpSock, rx = RX0 }) ->
receive {tcp, TcpSock, <<Size:16, Data/binary>>} ->
Size = byte_size(Data),
{ok, TX1, Msg1} = enoise_cipher_state:decrypt_with_ad(TX0, <<>>, Data),
{E#enoise{ tx = TX1 }, Msg1}
after 1000 -> error(timeout) end.
{ok, RX1, Msg1} = enoise_cipher_state:decrypt_with_ad(RX0, <<>>, Data),
{E#enoise{ rx = RX1 }, Msg1}
after 5000 -> error(timeout) end.
close(#enoise{ tcp_sock = TcpSock }) ->
gen_tcp:close(TcpSock).
@@ -69,7 +57,7 @@ close(#enoise{ tcp_sock = TcpSock }) ->
%%====================================================================
%% Internal functions
%%====================================================================
initiate_handshake(Role, Options) ->
do_handshake(TcpSock, Role, Options) ->
Prologue = proplists:get_value(prologue, Options, <<>>),
NoiseProtocol = proplists:get_value(noise, Options),
@@ -79,41 +67,22 @@ initiate_handshake(Role, Options) ->
RE = proplists:get_value(re, Options, undefined),
HSState = enoise_hs_state:init(NoiseProtocol, Role, Prologue, {S, E, RS, RE}),
{ok, HSState}.
do_handshake(TcpSock, HState, Options) ->
PreComm = proplists:get_value(pre_comm, Options, <<>>), %% TODO: Not standard!
gen_tcp:send(TcpSock, PreComm),
do_handshake(TcpSock, HState).
do_handshake(TcpSock, HSState).
do_handshake(TcpSock, HState) ->
case enoise_hs_state:next_message(HState) of
in ->
receive {tcp, TcpSock, Data} ->
case enoise_hs_state:read_message(HState, Data) of
{ok, HState1, _Msg} ->
do_handshake(TcpSock, HState1);
{done, _HState1, _Msg, {C1, C2}} ->
{ok, #enoise{ tcp_sock = TcpSock, rx = C1, tx = C2 }}
end
after 1000 ->
error(timeout)
end;
{ok, HState1, _Msg} = enoise_hs_state:read_message(HState, Data),
do_handshake(TcpSock, HState1)
after 1000 -> error(timeout) end;
out ->
case enoise_hs_state:write_message(HState, <<>>) of
{ok, HState1, Msg} ->
io:format("Sending: ~p\n", [add_len(Msg)]),
gen_tcp:send(TcpSock, add_len(Msg)),
do_handshake(TcpSock, HState1);
{done, _HState1, Msg, {C1, C2}} ->
io:format("Sending: ~p\n", [add_len(Msg)]),
gen_tcp:send(TcpSock, add_len(Msg)),
{ok, #enoise{ tcp_sock = TcpSock, rx = C1, tx = C2 }}
end
{ok, HState1, Msg} = enoise_hs_state:write_message(HState, <<>>),
gen_tcp:send(TcpSock, add_len(Msg)),
do_handshake(TcpSock, HState1);
done ->
{ok, #{ rx := Rx, tx := Tx }} = enoise_hs_state:finalize(HState),
{ok, #enoise{ tcp_sock = TcpSock, rx = Rx, tx = Tx }}
end.
add_len(Msg) ->