WIP

src/znoise.erl

@@ -18,7 +18,7 @@
 -copyright("QPQ AG <info@qpq.swiss>").
 -license("ISC").
 
-%% Main function with generic Noise handshake
+%% Generic Noise handshake
 -export([handshake/2, handshake/3, step_handshake/2]).
 
 %% API exports - Mainly mimicing gen_tcp

src/znoise_cipher.erl

@@ -0,0 +1,123 @@
+%%% @copyright 2026, QPQ AG
+%%% @copyright 2018, Aeternity Anstalt
+%%%
+%%% @doc
+%%% Module encapsulating a Noise Cipher state
+%%% @end
+
+-module(znoise_cipher).
+-vsn("0.1.0").
+-author("Craig Everett <craigeverett@qpq.swiss>").
+-author("Hans Svensson <hanssv@gmail.com>").
+-copyright("QPQ AG <info@qpq.swiss>").
+-license("ISC").
+
+-export([cipher/1,
+         decrypt_with_ad/3,
+         encrypt_with_ad/3,
+         has_key/1,
+         init/2,
+         key/1,
+         rekey/1,
+         set_key/2,
+         set_nonce/2]).
+
+-include("znoise.hrl").
+
+-type cipher() :: 'ChaChaPoly' | 'AESGCM'.
+-type nonce()  :: non_neg_integer().
+-type key()    :: empty | binary().
+
+% cs: "cipher state"
+-record(cs,
+        {k      = empty        :: key(),
+         n      = 0            :: nonce(),
+         cipher = 'ChaChaPoly' :: cipher()}).
+
+-opaque state() :: #cs{}.
+
+-export_type([cipher/0, state/0]).
+
+
+-spec init(key(), cipher()) -> state().
+
+init(Key, Cipher) ->
+    #cs{k = Key, n = 0, cipher = Cipher}.
+
+
+-spec set_key(CState, NewKey) -> NewCState
+    when CState    :: state(),
+         NewKey    :: key(),
+         NewCState :: state().
+
+set_key(CState, NewKey) ->
+    CState#cs{k = NewKey, n = 0}.
+
+
+-spec has_key(state()) -> boolean().
+
+has_key(#cs{k = Key}) ->
+    Key =/= empty.
+
+
+-spec set_nonce(CState, NewNonce) -> NewCState
+    when CState    :: state(),
+         NewNonce  :: nonce(),
+         NewCState :: state().
+
+set_nonce(CState, Nonce) ->
+    CState#cs{n = Nonce}.
+
+
+-spec encrypt_with_ad(CState, AD, PlainText) -> Outcome
+    when CState     :: state(),
+         AD         :: binary(),
+         PlainText  :: binary(),
+         Outcome    :: {ok, NewCState, CipherText} | {error, Reason},
+         NewCState  :: state(),
+         CipherText :: binary(),
+         Reason     :: term().
+
+encrypt_with_ad(CState = #cs{k = empty}, _AD, PlainText) ->
+    {ok, CState, PlainText};
+encrypt_with_ad(CState = #cs{k = K, n = N, cipher = Cipher}, AD, PlainText) ->
+    CipherText = znoise_crypto:encrypt(Cipher, K, N, AD, PlainText),
+    {ok, CState#cs{n = N + 1}, CipherText}.
+
+
+-spec decrypt_with_ad(CState, AD, CipherText) -> Outcome
+    when CState     :: state(),
+         AD         :: binary(),
+         CipherText :: binary(),
+         Outcome    :: {ok, NewCState, PlainText} | {error, Reason},
+         NewCState  :: state(),
+         PlainText  :: binary(),
+         Reason     :: term().
+
+decrypt_with_ad(CState = #cs{k = empty}, _AD, CipherText) ->
+    {ok, CState, CipherText};
+decrypt_with_ad(CState = #cs{k = K, n = N, cipher = Cipher}, AD, CipherText) ->
+    case znoise_crypto:decrypt(Cipher, K, N, AD, CipherText) of
+        PlainText when is_binary(PlainText) ->
+            {ok, CState#cs{n = N + 1}, PlainText};
+        Error ->
+            Error
+    end.
+
+-spec rekey(CState :: state()) -> state().
+
+rekey(CState = #cs{k = empty}) ->
+    CState;
+rekey(CState = #cs{k = K, cipher = Cipher}) ->
+    CState#cs{k = znoise_crypto:rekey(Cipher, K)}.
+
+
+-spec cipher(CState :: state()) -> cipher().
+
+cipher(#cs{cipher = Cipher}) ->
+    Cipher.
+
+-spec key(CState :: state()) -> key().
+
+key(#cs{k = K}) ->
+    K.

src/znoise_cipher_state.erl

@@ -1,88 +0,0 @@
-%%% @copyright 2026, QPQ AG
-%%% @copyright 2018, Aeternity Anstalt
-%%%
-%%% @doc
-%%% Module encapsulating a Noise Cipher state
-%%% @end
-
--module(znoise_cipher_state).
--vsn("0.1.0").
--author("Craig Everett <craigeverett@qpq.swiss>").
--author("Hans Svensson <hanssv@gmail.com>").
--copyright("QPQ AG <info@qpq.swiss>").
--license("ISC").
-
--export([ cipher/1
-        , decrypt_with_ad/3
-        , encrypt_with_ad/3
-        , has_key/1
-        , init/2
-        , key/1
-        , rekey/1
-        , set_key/2
-        , set_nonce/2
-        ]).
-
--include("znoise.hrl").
-
--type noise_cipher() :: 'ChaChaPoly' | 'AESGCM'.
--type nonce()        :: non_neg_integer().
--type key()          :: empty | binary().
-
--record(noise_cs, { k = empty             :: key()
-                  , n = 0                 :: nonce()
-                  , cipher = 'ChaChaPoly' :: noise_cipher() }).
-
--opaque state() :: #noise_cs{}.
-
--export_type([noise_cipher/0, state/0]).
-
--spec init(Key :: key(), Cipher :: noise_cipher()) -> state().
-init(Key, Cipher) ->
-    #noise_cs{ k = Key, n = 0, cipher = Cipher }.
-
--spec set_key(CState :: state(), NewKey :: key()) -> state().
-set_key(CState, NewKey) ->
-    CState#noise_cs{ k = NewKey, n = 0 }.
-
--spec has_key(CState :: state()) -> boolean().
-has_key(#noise_cs{ k = Key }) ->
-    Key =/= empty.
-
--spec set_nonce(CState :: state(), NewNonce :: nonce()) -> state().
-set_nonce(CState = #noise_cs{}, Nonce) ->
-    CState#noise_cs{ n = Nonce }.
-
--spec encrypt_with_ad(CState :: state(), AD :: binary(), PlainText :: binary()) ->
-                {ok, state(), binary()} | {error, term()}.
-encrypt_with_ad(CState = #noise_cs{ k = empty }, _AD, PlainText) ->
-    {ok, CState, PlainText};
-encrypt_with_ad(CState = #noise_cs{ k = K, n = N, cipher = Cipher }, AD, PlainText) ->
-    CipherText = znoise_crypto:encrypt(Cipher, K, N, AD, PlainText),
-    {ok, CState#noise_cs{ n = N+1 }, CipherText}.
-
--spec decrypt_with_ad(CState :: state(), AD :: binary(), CipherText :: binary()) ->
-                {ok, state(), binary()} | {error, term()}.
-decrypt_with_ad(CState = #noise_cs{ k = empty }, _AD, CipherText) ->
-    {ok, CState, CipherText};
-decrypt_with_ad(CState = #noise_cs{ k = K, n = N, cipher = Cipher }, AD, CipherText) ->
-    case znoise_crypto:decrypt(Cipher, K, N, AD, CipherText) of
-        PlainText when is_binary(PlainText) ->
-            {ok, CState#noise_cs{ n = N+1 }, PlainText};
-        Err = {error, _} ->
-            Err
-    end.
-
--spec rekey(CState :: state()) -> state().
-rekey(CState = #noise_cs{ k = empty }) ->
-    CState;
-rekey(CState = #noise_cs{ k = K, cipher = Cipher }) ->
-    CState#noise_cs{ k = znoise_crypto:rekey(Cipher, K) }.
-
--spec cipher(CState :: state()) -> noise_cipher().
-cipher(#noise_cs{ cipher = Cipher }) ->
-    Cipher.
-
--spec key(CState :: state()) -> key().
-key(#noise_cs{ k = K }) ->
-    K.

src/znoise_crypto.erl

@@ -63,7 +63,7 @@ hkdf(Hash, Key, Data) ->
     Output3 = hmac(Hash, TempKey, <<Output2/binary, 3:8>>),
     [Output1, Output2, Output3].
 
--spec rekey(Cipher :: znoise_cipher_state:noise_cipher(), Key :: binary()) -> binary().
+-spec rekey(Cipher :: znoise_cipher:noise_cipher(), Key :: binary()) -> binary().
 rekey('ChaChaPoly', K0) ->
     KLen = 32,
     <<K:KLen/binary, _/binary>> = encrypt('ChaChaPoly', K0, ?MAX_NONCE, <<>>, <<0:(32*8)>>),
@@ -71,13 +71,13 @@ rekey('ChaChaPoly', K0) ->
 rekey(Cipher, K) ->
     encrypt(Cipher, K, ?MAX_NONCE, <<>>, <<0:(32*8)>>).
 
--spec encrypt(Cipher :: znoise_cipher_state:noise_cipher(), Key :: binary(),
+-spec encrypt(Cipher :: znoise_cipher:noise_cipher(), Key :: binary(),
               Nonce :: non_neg_integer(), Ad :: binary(), PlainText :: binary()) -> binary().
 encrypt(Cipher, K, N, Ad, PlainText) ->
     {CText, CTag} = crypto:crypto_one_time_aead(cipher(Cipher), K, nonce(Cipher, N), PlainText, Ad, true),
     <<CText/binary, CTag/binary>>.
 
--spec decrypt(Cipher ::znoise_cipher_state:noise_cipher(), Key :: binary(),
+-spec decrypt(Cipher ::znoise_cipher:noise_cipher(), Key :: binary(),
               Nonce :: non_neg_integer(), AD :: binary(),
               CipherText :: binary()) -> binary() | {error, term()}.
 decrypt(Cipher, K, N, Ad, CipherText0) ->

src/znoise_hs_state.erl

@@ -25,8 +25,8 @@
 -type noise_dh()    :: dh25519 | dh448.
 -type noise_token() :: s | e | ee | ss | es | se.
 -type keypair()     :: znoise_keypair:keypair().
--type noise_split_state() :: #{rx          := znoise_cipher_state:state(),
+-type noise_split_state() :: #{rx          := znoise_cipher:state(),
-                               tx          := znoise_cipher_state:state(),
+                               tx          := znoise_cipher:state(),
                                hs_hash     := binary(),
                                final_state => state() }.
 
@@ -180,7 +180,7 @@ dh(#noise_hs{ dh = DH }, Key1, Key2) ->
 
 has_key(#noise_hs{ ss = SS }) ->
     CS = znoise_sym_state:cipher_state(SS),
-    znoise_cipher_state:has_key(CS).
+    znoise_cipher:has_key(CS).
 
 mix_key(HS = #noise_hs{ ss = SS0 }, Data) ->
     HS#noise_hs{ ss = znoise_sym_state:mix_key(SS0, Data) }.

src/znoise_keypair.erl

@@ -12,30 +12,32 @@
 -copyright("QPQ AG <info@qpq.swiss>").
 -license("ISC").
 
--export([ key_type/1
+-export([key_type/1,
-        , new/1
+         new/1,
-        , new/2
+         new/2,
-        , new/3
+         new/3,
-        , pubkey/1
+         pubkey/1,
-        , seckey/1
+         seckey/1]).
-        ]).
 
 -type key_type() :: dh25519 | dh448.
 
--record(kp, { type :: key_type()
+-record(kp,
-            , sec  :: binary() | undefined
+        {type :: key_type(),
-            , pub  :: binary() }).
+         sec  :: binary() | undefined,
+         pub  :: binary()}).
 
 -opaque keypair() :: #kp{}.
 %% Abstract keypair holding a secret key/public key pair and its type.
 
 -export_type([keypair/0]).
 
-%% @doc Generate a new keypair of type `Type'.
 -spec new(Type :: key_type()) -> keypair().
+%% @doc Generate a new keypair of type `Type'.
+
 new(Type) ->
     {Pub, Sec} = new_key_pair(Type),
-    #kp{ type = Type, sec = Sec, pub = Pub }.
+    #kp{type = Type, sec = Sec, pub = Pub}.
+
 
 %% @doc Create a new keypair of type `Type'. If `Public' is `undefined'
 %% it will be computed from the `Secret' (using the curve/algorithm

src/znoise_protocol.erl

@@ -12,15 +12,15 @@
 -copyright("QPQ AG <info@qpq.swiss>").
 -license("ISC").
 
--export([ cipher/1
+-export([cipher/1,
-        , dh/1
+         dh/1,
-        , from_name/1
+         from_name/1,
-        , hash/1
+         hash/1,
-        , msgs/2
+         msgs/2,
-        , pattern/1
+         pattern/1,
-        , pre_msgs/2
+         pre_msgs/2,
-        , supported/0
+         supported/0,
-        , to_name/1]).
+         to_name/1]).
 
 -ifdef(TEST).
 -export([to_name/4, from_name_pattern/1, to_name_pattern/1]).
@@ -30,17 +30,16 @@
 -type noise_msg()     :: {in | out, [znoise_hs_state:noise_token()]}.
 
 -record(noise_protocol,
-        { hs_pattern = noiseNN      :: noise_pattern()
+        {hs_pattern = noiseNN      :: noise_pattern()
-        , dh         = dh25519      :: znoise_hs_state:noise_dh()
+         dh         = dh25519      :: znoise_hs_state:noise_dh()
-        , cipher     = 'ChaChaPoly' :: znoise_cipher_state:noise_cipher()
+         cipher     = 'ChaChaPoly' :: znoise_cipher:noise_cipher()
-        , hash       = blake2b      :: znoise_sym_state:noise_hash()
+         hash       = blake2b      :: znoise_sym_state:noise_hash()}).
-        }).
 
 -opaque protocol() :: #noise_protocol{}.
 
 -export_type([noise_msg/0, noise_pattern/0, protocol/0]).
 
--spec cipher(Protocol :: protocol()) -> znoise_cipher_state:noise_cipher().
+-spec cipher(Protocol :: protocol()) -> znoise_cipher:noise_cipher().
 cipher(#noise_protocol{ cipher = Cipher }) ->
     Cipher.
 

src/znoise_sym_state.erl

@@ -28,7 +28,7 @@
 
 -type noise_hash() :: sha256 | sha512 | blake2s | blake2b.
 
--record(noise_ss, { cs             :: znoise_cipher_state:state()
+-record(noise_ss, { cs             :: znoise_cipher:state()
                   , ck = <<>>      :: binary()
                   , h  = <<>>      :: binary()
                   , hash = blake2b :: noise_hash() }).
@@ -50,13 +50,13 @@ init(Protocol) ->
     #noise_ss{ h = H1
              , ck = H1
              , hash = Hash
-             , cs = znoise_cipher_state:init(empty, Cipher) }.
+             , cs = znoise_cipher:init(empty, Cipher) }.
 
 -spec mix_key(SState :: state(), InputKeyMaterial :: binary()) -> state().
 mix_key(SState = #noise_ss{ hash = Hash, ck = CK0, cs = CS0 }, InputKeyMaterial) ->
     [CK1, <<TempK:32/binary, _/binary>> | _] =
         znoise_crypto:hkdf(Hash, CK0, InputKeyMaterial),
-    CS1 = znoise_cipher_state:set_key(CS0, TempK),
+    CS1 = znoise_cipher:set_key(CS0, TempK),
     SState#noise_ss{ ck = CK1, cs = CS1 }.
 
 -spec mix_hash(SState :: state(), Data :: binary()) -> state().
@@ -68,32 +68,32 @@ mix_hash(SState = #noise_ss{ hash = Hash, h = H0 }, Data) ->
 mix_key_and_hash(SState = #noise_ss{ hash = Hash, ck = CK0, cs = CS0 }, InputKeyMaterial) ->
     [CK1, TempH, <<TempK:32/binary, _/binary>>] =
         znoise_crypto:hkdf(Hash, CK0, InputKeyMaterial),
-    CS1 = znoise_cipher_state:set_key(CS0, TempK),
+    CS1 = znoise_cipher:set_key(CS0, TempK),
     mix_hash(SState#noise_ss{ ck = CK1, cs = CS1 }, TempH).
 
 -spec encrypt_and_hash(SState :: state(), PlainText :: binary()) -> {ok, state(), binary()}.
 encrypt_and_hash(SState = #noise_ss{ cs = CS0, h = H }, PlainText) ->
-    {ok, CS1, CipherText} = znoise_cipher_state:encrypt_with_ad(CS0, H, PlainText),
+    {ok, CS1, CipherText} = znoise_cipher:encrypt_with_ad(CS0, H, PlainText),
     {ok, mix_hash(SState#noise_ss{ cs = CS1 }, CipherText), CipherText}.
 
 -spec decrypt_and_hash(SState :: state(), CipherText :: binary()) ->
                 {ok, state(), binary()} | {error, term()}.
 decrypt_and_hash(SState = #noise_ss{ cs = CS0, h = H }, CipherText) ->
-    case znoise_cipher_state:decrypt_with_ad(CS0, H, CipherText) of
+    case znoise_cipher:decrypt_with_ad(CS0, H, CipherText) of
         Err = {error, _} ->
             Err;
         {ok, CS1, PlainText} ->
             {ok, mix_hash(SState#noise_ss{ cs = CS1 }, CipherText), PlainText}
     end.
 
--spec split(SState :: state()) -> {znoise_cipher_state:state(), znoise_cipher_state:state()}.
+-spec split(SState :: state()) -> {znoise_cipher:state(), znoise_cipher:state()}.
 split(#noise_ss{ hash = Hash, ck = CK, cs = CS }) ->
     [<<TempK1:32/binary, _/binary>>, <<TempK2:32/binary, _/binary>>, _] =
          znoise_crypto:hkdf(Hash, CK, <<>>),
-    {znoise_cipher_state:set_key(CS, TempK1),
+    {znoise_cipher:set_key(CS, TempK1),
-     znoise_cipher_state:set_key(CS, TempK2)}.
+     znoise_cipher:set_key(CS, TempK2)}.
 
--spec cipher_state(SState :: state()) -> znoise_cipher_state:state().
+-spec cipher_state(SState :: state()) -> znoise_cipher:state().
 cipher_state(#noise_ss{ cs = CS }) ->
     CS.
 

src/znoise_tcp.erl

@@ -4,8 +4,8 @@
 %%% @doc
 %%% A gen_server for holding a Noise connection over gen_tcp.
 %%%
-%%% Some care is needed since the underlying transmission is broken up
+%%% Currently only "raw" mode is supported, but a `{packet, N}'
-%%% into Noise packets, so we need some buffering.
+%%% option would be very convenient in the future.
 %%% @end
 
 -module(znoise_tcp).
@@ -15,229 +15,175 @@
 -copyright("QPQ AG <info@qpq.swiss>").
 -license("ISC").
 
--export([controlling_process/2,
+-export([start_listen/2,
-         close/1,
+         start_connect/2,
          send/2,
-         set_active/2,
+         close/1]).
-         start_link/5]).
 
-%% gen_server
+-include("$zx_include/zx_logger.hrl").
--export([init/1, handle_call/3, handle_cast/2, handle_info/2,
-         terminate/2, code_change/3]).
 
--record(znoise,
-        {pid}).
 
+% TODO: provide gen_tcp's packet option or 'raw'
 -record(s,
-        {rx         = ,
+        {rx   = none :: none | znoise_cipher:state(),
-         tx         = ,
+         tx   = none :: none | znoise_cipher:state(),
-         owner      = none :: none | pid(),
+         sock = none :: none | gen_tcp:socket(),
-         owner_ref  = none :: none | reference(),
+         mode = raw  :: raw | {packet, 1..4},
-         tcp_sock   = none :: none | gen_tcp:socket(),
+         buff = <<>> :: binary()}).
-         active     = once :: true | {once, boolean()},
-         msgbuf     = []   :: list(),
-         rawbuf     = <<>> :: binary()}).
 
 
-start_link(TcpSock, Rx, Tx, Owner, {Active0, Buf}) ->
+%%% Initializers
-    Active =
+
-        case Active0 of
+-spec start_listen(ListenSock, Options, Noise) -> Outcome
-            true -> true;
+    when ListenSock ::
-            once -> {once, false}
+         Options    ::
-        end,
+         Noise      :: 
-    State =
+         Outcome    :: {ok, pid()} | {error, Reason :: term()}.
-        #s{rx       = Rx,
+
-           tx       = Tx,
+start_listen(ListenSock, Options, Noise) ->
-           owner    = Owner,
+    proc_lib:start_link(?MODULE, init_listen, [self(), ListenSock, Options, Noise]).
-           tcp_sock = TcpSock,
+
-           active   = Active},
+
-    case gen_server:start_link(?MODULE, [State], []) of
+init_listen(Parent, ListenSock, Options, Noise) ->
-        {ok, Pid} ->
+    Debug = sys:debug_options(proplists:get_value(debug, Options, [])),
-            case gen_tcp:controlling_process(TcpSock, Pid) of
+    ok = proc_lib:init_ack(Parent, {ok, self()}),
-                ok ->
+    listen(Parent, Debug, ListenSock, Options, Noise).
-                    % Changing controlling process require a bit of
+
-                    % fiddling with already received and delivered content...
+
-                    ok =
+listen(Parent, Debug, ListenSock, Options, Noise) ->
-                        case Buf =/= <<>> of
+
-                            true  -> Pid ! {tcp, TcpSock, Buf};
+
-                            false -> ok
+start_connect(Host, Options, Noise) ->
-                        end,
+    proc_lib:start_link(?MODULE, init_connect, [self(), Host, Options, Noise]).
-                    flush_tcp(Pid, TcpSock),
+
-                    {ok, Pid};
+
-                Error ->
+init_connect(Parent, Host, Options, Noise) ->
-                    close(Pid),
+    Debug = sys:debug_options(proplists:get_value(debug, Options, [])),
-                    Error
+    ok = proc_lib:init_ack(Parent, {ok, self()}),
-            end;
+    connect(Parent, Debug, Host, Options, Noise).
-        Error
+
-            Error
+
+connect(Parent, Debug, Host, Options, Noise) ->
+
+
+
+%%% Interface
+
+-spec send(Conn, Data) -> Outcome
+    when Conn    :: pid(),
+         Data    :: binary() | {file, file:filename()},
+         Outcome :: ok | {error, Reason :: term()}.
+
+send(Conn, Data) ->
+    call(Conn, {send, Data}).
+
+
+-spec close(Conn) -> Outcome
+    case Conn    :: pid(),
+         Outcome :: ok | {error, Reason :: term()}.
+
+close(Conn) ->
+    call(Conn, close).
+
+
+%% Non-system calls should only ever come from the parent
+
+call(Conn, Data) ->
+    call(Conn, Data, 5000).
+
+call(Conn, Data, Timeout) ->
+    Ref = make_ref(),
+    Conn ! {call, Ref, Data},
+    receive
+        {resp, Ref, Result} -> Result
+        after Timeout       -> {error, timeout}
     end.
 
 
--spec send(Noise :: pid(), Data :: binary()) -> ok | {error, term()}.
+-spec loop(Parent, Debug, State) -> no_return()
+    when Parent :: pid(),
+         Debug  :: [sys:dbg_opt()],
+         State  :: state().
 
-send(Noise, Data) ->
+loop(Parent, Debug, State = #s{sock = Sock}) ->
-    gen_server:call(Noise, {send, Data}).
+    case inet:setopts(Sock, [{active, once}]) of
+        ok              -> wait(Parent, Debug, State);
+        {error, einval} -> exit(normal)
+    end.
+
+wait(Parent, Debug, State = #s{sock = Sock}) ->
+    receive
+        {tcp, Sock, Binary} ->
+            Newtate = get_bytes(Parent, Binary, State),
+            loop(Parent, Debug, NewState);
+        {call, Ref, {send, Data}} ->
+            {Result, NewState} = send_bytes(Data, State),
+            Parent ! {resp, Ref, Result},
+            wait(Parent, Debug, NewState);
+        {call, Ref, close} ->
+            Result = gen_tcp:close(Sock),
+            Parent ! {resp, Ref, Result},
+            exit(normal);
+        {tcp_closed, Sock} ->
+            Parent ! {noise_closed, self()},
+            ok = tell(info, "Socket closed!"),
+            exit(normal);
+        {tcp_error, Sock, Info} ->
+            Parent ! {noise_error, Info},
+            ok = tell(info, "Socket error: ~tp", [Info]),
+            exit(normal);
+        {system, From, Request} ->
+            sys:handle_system_msg(Request, From, Parent, ?MODULE, Debug, State);
+        Unexpected ->
+            ok = tell(info, "~p Unexpected message: ~tp", [self(), Unexpected]),
+            wait(Parent, Debug, State)
+    end.
 
 
--spec set_active(Noise :: pid(), Active :: true | once) -> ok | {error, term()}.
+system_continue(Parent, Debug, State) ->
-
+    wait(Parent, Debug, State).
-set_active(Noise, Active) ->
-    gen_server:call(Noise, {active, self(), Active}).
 
 
--spec close(Noise :: pid()) -> ok | {error, term()}.
+system_terminate(Reason, _Parent, _Debug, _State) ->
-
+    exit(Reason).
-close(Noise) ->
-    gen_server:call(Noise, close).
 
 
--spec controlling_process(Noise :: pid(), NewPid :: pid()) -> ok | {error, term()}.
+system_get_state(State) ->
-
+    {ok, State#s{rx = nope, tx = nope}}.
-controlling_process(Noise, NewPid) ->
-    gen_server:call(Noise, {controlling_process, self(), NewPid}, 100).
 
 
-%% gen_server
+system_replace_state(SusFun, State) ->
-
+    ok = tell(info, "Attempt to run system_replace_state/2 with ~tp", [SusFun]),
-init([#s{owner = Owner} = State]) ->
+    {ok, State, State}.
-    OwnerRef = erlang:monitor(process, Owner),
-    {ok, State#s{owner_ref = OwnerRef}}.
 
 
-handle_call(close, _, State) ->
+system_code_change(State, _Module, _OldVsn, _Extra) ->
-    {stop, normal, ok, State};
-handle_call(_Call, _, State = #s{tcp_sock = closed}) ->
-    {reply, {error, closed}, State};
-handle_call({send, Data}, _, State) ->
-    {Result, NewState} = handle_send(State, Data),
-    {reply, Result, NewState};
-handle_call({controlling_process, OldPID, NewPID}, _, State) ->
-    {Result, NewState} = handle_control_change(State, OldPID, NewPID),
-    {reply, Result, NewState};
-handle_call({active, PID, NewActive}, _, State) ->
-    {Result, NewState} = handle_active(State, PID, NewActive),
-    {reply, Result, NewState}.
-
-
-handle_cast(_, State) ->
-    {noreply, State}.
-
-
-handle_info({tcp, TS, Data}, State = #s{tcp_sock = TS, owner = O}) ->
-    try
-        {NextState = #s{msgbuf = Buf}, Msgs} = handle_data(State, Data),
-        NewState = handle_msgs(NextState#s{msgbuf = Buf ++ Msgs}),
-        set_active(NewState),
-        {noreply, NewState}
-    catch error:{znoise_error, _} ->
-        %% We are not likely to recover, but leave the decision to upstream
-        O ! {znoise_error, TS, decrypt_error},
-        {noreply, State}
-    end;
-handle_info({tcp_closed, TS}, State = #s{tcp_sock = TS, owner = O}) ->
-    O ! {tcp_closed, TS},
-    {noreply, State#s{tcp_sock = closed}};
-handle_info({'DOWN', OwnerRef, process, _, normal},
-            State = #s{tcp_sock = TS, owner_ref = OwnerRef}) ->
-    close_tcp(TS),
-    {stop, normal, State#s{tcp_sock = closed, owner_ref = undefined}};
-handle_info({'DOWN', _, _, _, _}, State) ->
-    %% Ignore non-normal monitor messages - we are linked.
-    {noreply, State};
-handle_info(_Msg, State) ->
-    {noreply, State}.
-
-
-terminate(_, #s{tcp_sock = TcpSock, owner_ref = ORef}) ->
-    [ gen_tcp:close(TcpSock) || TcpSock /= closed ],
-    [ erlang:demonitor(ORef, [flush]) || ORef /= undefined ],
-    ok.
-
-code_change(_OldVsn, State, _Extra) ->
     {ok, State}.
 
 
-
 %%% Handlers
 
-handle_control_change(State = #s{owner = PID, owner_ref = OldRef}, PID, NewPID) ->
+read_bytes(Parent, <<Len:16, CT:Len/binary, Rest/binary>>, State = #s{buff = <<>>, rx = RX}) ->
-    NewRef = erlang:monitor(process, NewPID),
+    case znoise_cipher:decrypt_with_ad(RX, <<>>, CT) of
-    erlang:demonitor(OldRef, [flush]),
+        {ok, NewRX, PT} ->
-    {ok, State#s{owner = NewPID, owner_ref = NewRef}};
+            Parent ! {noise, PT},
-handle_control_change(State, _, _) ->
+            read_bytes(Parent, Rest, State#s{rx = NewRX});
-    {{error, not_owner}, State}.
+        {error, _} ->
-
+            Parent ! {noise_error, decrypt_input_failed}
-
+            exit(normal)
-handle_active(State = #s{owner = PID, tcp_sock = TcpSock}, PID, Active) ->
-    case Active of
-        true ->
-            inet:setopts(TcpSock, [{active, true}]),
-            {ok, handle_msgs(State#s{active = true})};
-        once ->
-            NewState = handle_msgs(State#s{active = {once, false}}),
-            set_active(NewState),
-            {ok, NewState}
     end;
-handle_active(State, _, _) ->
+read_bytes(Parent, Received, State = #s{buff = <<>>}) ->
-    {{error, not_owner}, State}.
+    State#s{buff = Received};
-
+read_bytes(Parent, Received, State = #s{buff = Buff}) ->
-
+    case <<Buff/binary, Received/binary>> of
-handle_data(State = #s{rawbuf = Buf, rx = RX}, Data) ->
+        <<>> -> State;
-    case <<Buf/binary, Data/binary>> of
+        Data -> read_bytes(Parent, Data, State#s{buff = <<>>})
-        B = <<Len:16, Rest/binary>> when Len > byte_size(Rest) ->
-            {State#s{rawbuf = B}, []}; %% Not a full Noise message - save it
-        <<Len:16, Rest/binary>> ->
-            <<Msg:Len/binary, Rest2/binary>> = Rest,
-            case znoise_cipher_state:decrypt_with_ad(RX, <<>>, Msg) of
-                {ok, NewRX, NewMsg} ->
-                    {NewState, Msgs} = handle_data(State#s{rawbuf = Rest2, rx = NewRX}, <<>>),
-                    {NewState, [NewMsg | Msgs]};
-                {error, _} ->
-                    error({znoise_error, decrypt_input_failed})
-            end;
-        EmptyOrSingleByte ->
-            {State#s{rawbuf = EmptyOrSingleByte}, []}
     end.
 
 
-handle_msgs(State = #s{msgbuf = []}) ->
+send_bytes(Data, State = #s{sock = Sock, tx = TX}) ->
-    State;
+    {ok, NewTX, Msg} = znoise_cipher:encrypt_with_ad(TX, <<>>, Data),
-handle_msgs(State = #s{msgbuf = Msgs, active = true, owner = Owner}) ->
-    [ Owner ! {noise, #znoise{pid = self()}, Msg} || Msg <- Msgs ],
-    State#s{msgbuf = []};
-handle_msgs(State = #s{msgbuf = [Msg | Msgs], active = {once, Delivered}, owner = Owner}) ->
-    case Delivered of
-        true  ->
-            State;
-        false ->
-            Owner ! {noise, #znoise{pid = self()}, Msg},
-            State#s{msgbuf = Msgs, active = {once, true}}
-    end.
-
-
-handle_send(State = #s{tcp_sock = TcpSock, tx = TX}, Data) ->
-    {ok, MewTX, Msg} = znoise_cipher_state:encrypt_with_ad(TX, <<>>, Data),
     case gen_tcp:send(TcpSock, <<(byte_size(Msg)):16, Msg/binary>>) of
-        ok    -> {ok, State#s{tx = MewTX}};
+        ok    -> {ok, State#s{tx = NewTX}};
         Error -> {Error, State}
     end.
-
-
-set_active(#s{msgbuf = [], active = {once, _}, tcp_sock = TcpSock}) ->
-    inet:setopts(TcpSock, [{active, once}]);
-set_active(_) ->
-    ok.
-
-
-flush_tcp(Pid, TcpSock) ->
-    receive {tcp, TcpSock, Data} ->
-        Pid ! {tcp, TcpSock, Data},
-        flush_tcp(Pid, TcpSock)
-    after 1 -> ok
-    end.
-
-
-close_tcp(closed) ->
-    ok;
-close_tcp(Sock) ->
-    gen_tcp:close(Sock).