zNoise/src/znoise.erl

%%% @copyright 2026, QPQ AG
%%% @copyright 2018, Aeternity Anstalt
%%%
%%% @doc
%%% Interface to the Noise protocol: https://noiseprotocol.org
%%%
%%% This is a fork of the `enoise' project: https://git.qpq.swiss/QPQ-AG/enoise
%%%
%%% For convenience there is also an API to use Noise over TCP (i.e. `gen_tcp')
%%% and after "upgrading" a `gen_tcp'-socket into a `znoise'-socket it has a
%%% similar API as `gen_tcp'.
%%% @end

-module(znoise).
-vsn("0.1.0").
-author("Craig Everett <craigeverett@qpq.swiss>").
-author("Hans Svensson <hanssv@gmail.com>").
-copyright("QPQ AG <info@qpq.swiss>").
-license("ISC").

%% Main function with generic Noise handshake
-export([handshake/2, handshake/3, step_handshake/2]).

%% API exports - Mainly mimicing gen_tcp
-export([accept/2,
         close/1,
         connect/2,
         controlling_process/2,
         send/2,
         set_active/2]).

-record(znoise,
        {pid = none | pid()}).

-type key() :: binary().
-type keypair() :: znoise_keypair:keypair().

-type options() :: [option()].
%% A list of Noise options is a proplist, it *must* contain a value `noise'
%% that describes which Noise configuration to use. It is possible to give
%% a `prologue' to the protocol. And for the protocol to work, the correct
%% configuration of pre-defined keys (`s', `e', `rs', `re') should also be
%% provided.

-type option() :: {noise,    protocol_option()}     %% Required
                | {e,        keypair()}             %% Mandatary depending on `noise'
                | {s,        keypair()}
                | {re,       key()}
                | {rs,       key()}
                | {prologue, binary()}              %% Optional
                | {timeout,  integer() | infinity}. %% Optional

-type protocol_option() :: znoise_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 timeout() :: pos_integer() | infinity.
-type recv_return() :: {ok, binary(), com_state_state()}
                     | {error, term()}).
-type recv_msg_fun() :: fun((com_state_state(), timeout()) -> recv_return()).

-type send_msg_fun() :: fun((com_state_state(), binary()) -> ok).

-type 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.

-type split_state() :: znoise_hs_state:noise_split_state().
%% Return value from the final `split' operation. Provides a CipherState for
%% receiving and a CipherState transmission. Also includes the final handshake
%% hash for channel binding.

-opaque socket() :: #znoise{}.
%% An abstract Noise socket - holds a reference to a socket that has completed
%% a Noise handshake.

-export_type([socket/0]).



%%% API functions

-spec handshake(Options, Role) -> Outcome
    when Options :: options(),
         Role    :: znoise_hs_state:noise_role(),
         Outcome :: {ok, znoise_hs_state:state()}
                  | {error, term()}.
%% @doc
%% Start an interactive handshake

handshake(Options, Role) ->
    create_hstate(Options, Role).


-spec handshake(Options, Role, ComState) -> Outcome
    when Options  :: options(),
         Role     :: znoise_hs_state:noise_role(),
         ComState :: com_state(),
         Outcome  :: {ok, split_state(), com_state()}
                   | {error, term()}.
%% @doc
%% Perform a Noise handshake

handshake(Options, Role, ComState) ->
    case create_hstate(Options, Role) of
        {ok, HState} ->
            Timeout = proplists:get_value(timeout, Options, infinity),
            do_handshake(HState, ComState, Timeout);
        Err = {error, _} ->
            Err
    end.


-spec step_handshake(HState, Data) -> Next
    when HState :: znoise_hs_state:state(),
         Data   :: {rcvd, binary()}
                 | {send, binary()},
         Next   :: {ok, send, binary(), znoise_hs_state:state()}
                 | {ok, rcvd, binary(), znoise_hs_state:state()}
                 | {ok, done, split_state()}
                 | {error, term()}.
%% @doc
%% Do a step (one of `{send, Payload}', `{rcvd, EncryptedData}', or `done')

step_handshake(HState, Data) ->
    do_step_handshake(HState, Data).


-spec connect(TcpSock, Options) -> Outcome
    when TcpSock :: gen_tcp:socket(),
         Options :: options(),
         Outcome :: {ok, socket(), znoise_hs_state:state()}
                  | {error, term()}.
%% @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 options()} is a proplist.

connect(TcpSock, Options) ->
    tcp_handshake(TcpSock, initiator, Options).


-spec accept(TcpSock, Options) -> Outcome
    when TcpSock :: gen_tcp:socket(),
         Options :: options(),
         Outcome :: {ok, socket(), znoise_hs_state:state()}
                  | {error, term()}.
%% @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 options()} is a proplist.

accept(TcpSock, Options) ->
    tcp_handshake(TcpSock, responder, Options).


-spec send(Socket, Data) -> Outcome
    when Socket  :: socket(),
         Data    :: binary(),
         Outcome :: ok | {error, term()}.
%% @doc
%% Writes `Data' to `Socket'

send(#znoise{ pid = Pid }, Data) ->
    znoise_tcp:send(Pid, Data).


-spec close(NoiseSock) -> Outcome
    when NoiseSock :: socket()
         Outcome   :: ok | {error, term()}.
%% @doc
%% Closes a Noise connection.

close(#znoise{ pid = Pid }) ->
    znoise_tcp:close(Pid).


-spec controlling_process(Socket, PID) -> Outcome
    when Socket  :: socket(),
         PID     :: pid(),
         Outcome :: ok | {error, term()}.
%% @doc
%% Assigns a new controlling process to the Noise socket. A controlling
%% process is the owner of an Noise socket, and receives all messages from the
%% socket.

controlling_process(#znoise{pid = PID}, NewPID) ->
    znoise_tcp:controlling_process(PID, NewPID).


-spec set_active(Socket, Mode) -> Outcome
    when Socket  :: socket(),
         Mode    :: true | once,
         Outcome :: ok | {error, term()}.
%% @doc
%% Set the active option `true | once'. Note that `N' and `false' are
%% not valid options for a Noise socket.

set_active(#znoise{ pid = Pid }, ActiveMode) ->
    znoise_tcp:set_active(Pid, ActiveMode).



%%% Internal functions

do_handshake(HState, ComState, Timeout) ->
    case znoise_hs_state:next_message(HState) of
        in ->
            case hs_recv_msg(ComState, Timeout) of
                {ok, Data, ComState1} ->
                    case znoise_hs_state:read_message(HState, Data) of
                        {ok, HState1, _Msg} ->
                            do_handshake(HState1, ComState1, Timeout);
                        Err = {error, _} ->
                            Err
                    end;
                Err = {error, _} ->
                    Err
            end;
        out ->
            {ok, HState1, Msg} = znoise_hs_state:write_message(HState, <<>>),
            case hs_send_msg(ComState, Msg) of
                {ok, ComState1} ->
                    do_handshake(HState1, ComState1, Timeout);
                Err = {error, _} ->
                    Err
            end;
        done ->
            {ok, Res} = znoise_hs_state:finalize(HState),
            {ok, Res, ComState}
    end.

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.

hs_send_msg(CS = #{ send_msg := Send, state := S }, Data) ->
    case Send(S, Data) of
        {ok, S1}         -> {ok, CS#{ state := S1 }};
        Err = {error, _} -> Err
    end.

do_step_handshake(HState, Data) ->
    case {znoise_hs_state:next_message(HState), Data} of
        {in, {rcvd, Encrypted}} ->
            case znoise_hs_state:read_message(HState, Encrypted) of
                {ok, HState1, Msg} ->
                    {ok, rcvd, Msg, HState1};
                Err = {error, _} ->
                    Err
            end;
        {out, {send, Payload}} ->
            {ok, HState1, Msg} = znoise_hs_state:write_message(HState, Payload),
            {ok, send, Msg, HState1};
        {done, done} ->
            {ok, Res} = znoise_hs_state:finalize(HState),
            {ok, done, Res};
        {Next, _} ->
            {error, {invalid_step, expected, Next, got, Data}}
    end.

%% -- gen_tcp specific functions ---------------------------------------------
tcp_handshake(TcpSock, Role, Options) ->
    case check_gen_tcp(TcpSock) of
        ok ->
            case inet:getopts(TcpSock, [active]) of
                {ok, [{active, Active}]} ->
                    do_tcp_handshake(Options, Role, TcpSock, Active);
                Err = {error, _} ->
                    Err
            end;
        Err = {error, _} ->
            Err
    end.


do_tcp_handshake(Options, Role, TcpSock, Active) ->
    ComState = #{recv_msg => fun gen_tcp_rcv_msg/2,
                 send_msg => fun gen_tcp_snd_msg/2,
                 state    => {TcpSock, Active, <<>>}},
    case handshake(Options, Role, ComState) of
        {ok, #{rx := Rx, tx := Tx, final_state := FState}, #{state := {_, _, Buf}}} ->
            case znoise_tcp:start_link(TcpSock, Rx, Tx, self(), {Active, Buf}) of
                {ok, Pid} -> {ok, #znoise{ pid = Pid }, FState};
                Error     -> Error
            end;
        Error ->
            Error
    end.


create_hstate(Options, Role) ->
    Prologue = proplists:get_value(prologue, Options, <<>>),
    Noise    = proplists:get_value(noise, Options),
    Protocol =
        case is_binary(Noise) orelse is_list(Noise) of
            true  -> znoise_protocol:from_name(X);
            false -> Noise
        end,
    DH = znoise_protocol:dh(Protocol),
    S  = proplists:get_value(s, Options, undefined),
    E  = proplists:get_value(e, Options, undefined),
    RS = remote_keypair(DH, proplists:get_value(rs, Options, undefined)),
    RE = remote_keypair(DH, proplists:get_value(re, Options, undefined)),
    znoise_hs_state:init(Protocol, Role, Prologue, {S, E, RS, RE}).


check_gen_tcp(TcpSock) ->
    case inet:getopts(TcpSock, [mode, packet, active, header, packet_size]) of
        {ok, TcpOpts} ->
            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)
                    andalso (Active == true orelse Active == once)
                    andalso Header  == 0
                    andalso PSize   == 0
                    andalso Mode    == binary of
                true ->
                    gen_tcp:controlling_process(TcpSock, self());
                false ->
                    {error, {invalid_tcp_options, TcpOpts}}
            end;
        Error ->
            Error
    end.


gen_tcp_snd_msg(S = {TcpSock, _, _}, Msg) ->
    Len = byte_size(Msg),
    case gen_tcp:send(TcpSock, <<Len:16, Msg/binary>>) of
        ok    -> {ok, S};
        Error -> Error
    end.


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}, Timeout);
                <<Len:16, Rest/binary>> ->
                    <<Data1:Len/binary, Buf1/binary>> = Rest,
                    {ok, Data1, {TcpSock, true, Buf1}}
            end
    after Timeout ->
        {error, timeout}
    end.


remote_keypair(_DH, undefined) ->
    undefined;
remote_keypair(DH, RemotePub) when is_binary(RemotePub) ->
    znoise_keypair:new(DH, RemotePub).