Support {active, true} and {active, once} in gen_tcp-wrapper

In the implemented TCP-wrapper (enoise_connection) we now properly support {active, true}
and {active, once} and switching between them (previously no switching was supported).
This commit is contained in:
Hans Svensson
2018-03-13 23:22:42 +01:00
parent 272dcde689
commit 8f3aff4d8b
5 changed files with 152 additions and 186 deletions
+65 -47
View File
@@ -22,9 +22,8 @@
, close/1
, connect/2
, controlling_process/2
, recv/2
, recv/3
, send/2 ]).
, send/2
, set_active/2 ]).
-record(enoise, { pid }).
@@ -43,13 +42,31 @@
| {s, noise_keypair()}
| {re, noise_key()}
| {rs, noise_key()}
| {prologue, binary()}. %% Optional
| {prologue, binary()} %% Optional
| {timeout, integer() | infinity}. %% Optional
-type noise_protocol_option() :: enoise_protocol:protocol() | string() |
binary().
%% Either an instantiated Noise protocol configuration or the name of a Noise
%% configuration (either as a string or a binary string).
-type com_state_state() :: term().
%% The state part of a communiction state
-type recv_msg_fun() :: fun((com_state_state(), integer() | infinity) ->
{ok, binary(), com_state_state()} | {error, term()}).
%% Function that receive a message
-type send_msg_fun() :: fun((com_state_state(), binary()) -> ok).
%% Function that sends a message
-type noise_com_state() :: #{ recv_msg := recv_msg_fun(),
send_msg := send_msg_fun(),
state := term() }.
%% Noise communication state - used to parameterize a handshake. Consists of a
%% send function one receive function and an internal state.
-opaque noise_socket() :: #enoise{}.
%% An abstract Noise socket - holds a reference to a socket that has completed
%% a Noise handshake.
@@ -61,6 +78,13 @@ binary().
%%====================================================================
%% @doc Start an interactive handshake
%% @end
-spec handshake(Options :: noise_options(),
Role :: enoise_hs_state:noise_role()) ->
{in, enoise_hs_state:state()}
| {out, binary(), enoise_hs_state:state()}
| {done, enoise_hs_state:state()}
| {error, term()}.
handshake(Options, Role) ->
HState = create_hstate(Options, Role),
step_handshake(HState, <<>>).
@@ -68,15 +92,24 @@ handshake(Options, Role) ->
step_handshake(HState, Data) ->
do_step_handshake(HState, Data).
%% @doc The main function - performs a Noise handshake
%% @doc Perform a Noise handshake
%% @end
-spec handshake(Options :: noise_options(),
Role :: enoise_hs_state:noise_role(),
ComState :: noise_com_state()) ->
{ok, map(), noise_com_state()} | {error, term()}.
handshake(Options, Role, ComState) ->
HState = create_hstate(Options, Role),
do_handshake(HState, ComState).
Timeout = proplists:get_value(timeout, Options, infinity),
do_handshake(HState, ComState, Timeout).
%% @doc Upgrades a gen_tcp, or equivalent, connected socket to a Noise socket,
%% that is, performs the client-side noise handshake.
%%
%% Note: The TCP socket has to be in mode `{active, true}' or `{active, once}',
%% passive receive is not supported.
%%
%% {@link noise_options()} is a proplist.
%% @end
-spec connect(TcpSock :: gen_tcp:socket(),
@@ -88,6 +121,9 @@ connect(TcpSock, Options) ->
%% @doc Upgrades a gen_tcp, or equivalent, connected socket to a Noise socket,
%% that is, performs the server-side noise handshake.
%%
%% Note: The TCP socket has to be in mode `{active, true}' or `{active, once}',
%% passive receive is not supported.
%%
%% {@link noise_options()} is a proplist.
%% @end
-spec accept(TcpSock :: gen_tcp:socket(),
@@ -102,29 +138,6 @@ accept(TcpSock, Options) ->
send(#enoise{ pid = Pid }, Data) ->
enoise_connection:send(Pid, Data).
%% @equiv recv(Socket, Length, infinity)
-spec recv(Socket :: noise_socket(), Length :: integer()) ->
{ok, binary()} | {error, term()}.
recv(Socket, Length) ->
recv(Socket, Length, infinity).
%% @doc Receives a packet from a socket in passive mode. A closed socket is
%% indicated by return value `{error, closed}'.
%%
%% Argument `Length' denotes the number of bytes to read. If Length = 0, all
%% available bytes are returned. If Length > 0, exactly Length bytes are
%% returned, or an error; possibly discarding less than Length bytes of data
%% when the socket gets closed from the other side.
%%
%% Optional argument `Timeout' specifies a time-out in milliseconds. The
%% default value is `infinity'.
%% @end
-spec recv(Socket :: noise_socket(), Length :: integer(),
Timeout :: integer() | infinity) ->
{ok, binary()} | {error, term()}.
recv(#enoise{ pid = Pid }, Length, Timeout) ->
enoise_connection:recv(Pid, Length, Timeout).
%% @doc Closes a Noise connection.
%% @end
-spec close(NoiseSock :: noise_socket()) -> ok | {error, term()}.
@@ -140,30 +153,38 @@ close(#enoise{ pid = Pid }) ->
controlling_process(#enoise{ pid = Pid }, NewPid) ->
enoise_connection:controlling_process(Pid, NewPid).
%% @doc Set the active option `true | once'. Note that `N' and `false' are
%% not valid options for a Noise socket.
%% @end
-spec set_active(Socket :: noise_socket(), Mode :: true | once) ->
ok | {error, term()}.
set_active(#enoise{ pid = Pid }, ActiveMode) ->
enoise_connection:set_active(Pid, ActiveMode).
%%====================================================================
%% Internal functions
%%====================================================================
do_handshake(HState, ComState) ->
do_handshake(HState, ComState, Timeout) ->
case enoise_hs_state:next_message(HState) of
in ->
case hs_recv_msg(ComState) of
case hs_recv_msg(ComState, Timeout) of
{ok, Data, ComState1} ->
{ok, HState1, _Msg} = enoise_hs_state:read_message(HState, Data),
do_handshake(HState1, ComState1);
do_handshake(HState1, ComState1, Timeout);
Err = {error, _} ->
Err
end;
out ->
{ok, HState1, Msg} = enoise_hs_state:write_message(HState, <<>>),
{ok, ComState1} = hs_send_msg(ComState, Msg),
do_handshake(HState1, ComState1);
do_handshake(HState1, ComState1, Timeout);
done ->
{ok, Res} = enoise_hs_state:finalize(HState),
{ok, Res, ComState}
end.
hs_recv_msg(CS = #{ recv_msg := Recv, state := S }) ->
case Recv(S) of
hs_recv_msg(CS = #{ recv_msg := Recv, state := S }, Timeout) ->
case Recv(S, Timeout) of
{ok, Data, S1} -> {ok, Data, CS#{ state := S1 }};
Err = {error, _} -> Err
end.
@@ -193,7 +214,7 @@ tcp_handshake(TcpSock, Role, Options) ->
case check_gen_tcp(TcpSock) of
ok ->
{ok, [{active, Active}]} = inet:getopts(TcpSock, [active]),
ComState = #{ recv_msg => fun gen_tcp_rcv_msg/1,
ComState = #{ recv_msg => fun gen_tcp_rcv_msg/2,
send_msg => fun gen_tcp_snd_msg/2,
state => {TcpSock, Active, <<>>} },
@@ -228,12 +249,13 @@ create_hstate(Options, Role) ->
Prologue, {S, E, RS, RE}).
check_gen_tcp(TcpSock) ->
{ok, TcpOpts} = inet:getopts(TcpSock, [mode, packet, header, packet_size]),
{ok, TcpOpts} = inet:getopts(TcpSock, [mode, packet, active, header, packet_size]),
Packet = proplists:get_value(packet, TcpOpts, 0),
Active = proplists:get_value(active, TcpOpts, 0),
Header = proplists:get_value(header, TcpOpts, 0),
PSize = proplists:get_value(packet_size, TcpOpts, undefined),
Mode = proplists:get_value(mode, TcpOpts, binary),
case (Packet == 0 orelse Packet == raw)
case (Packet == 0 orelse Packet == raw) andalso (Active == true orelse Active == once)
andalso Header == 0 andalso PSize == 0 andalso Mode == binary of
true ->
gen_tcp:controlling_process(TcpSock, self());
@@ -246,22 +268,18 @@ gen_tcp_snd_msg(S = {TcpSock, _, _}, Msg) ->
ok = gen_tcp:send(TcpSock, <<Len:16, Msg/binary>>),
{ok, S}.
gen_tcp_rcv_msg({TcpSock, true, Buf}) ->
gen_tcp_rcv_msg({TcpSock, Active, Buf}, Timeout) ->
receive {tcp, TcpSock, Data} ->
%% Immediately re-set {active, once}
[ inet:setopts(TcpSock, [{active, once}]) || Active == once ],
case <<Buf/binary, Data/binary>> of
Buf1 = <<Len:16, Rest/binary>> when byte_size(Rest) < Len ->
gen_tcp_rcv_msg({TcpSock, true, Buf1});
gen_tcp_rcv_msg({TcpSock, true, Buf1}, Timeout);
<<Len:16, Rest/binary>> ->
<<Data1:Len/binary, Buf1/binary>> = Rest,
{ok, Data1, {TcpSock, true, Buf1}}
end
after 1000 ->
after Timeout ->
{error, timeout}
end;
gen_tcp_rcv_msg(S = {TcpSock, false, <<>>}) ->
{ok, <<Len:16>>} = gen_tcp:recv(TcpSock, 2, 1000),
case gen_tcp:recv(TcpSock, Len, 1000) of
{ok, Data} -> {ok, Data, S};
Err = {error, _} -> Err
end.