fix chacha20

README.md

@@ -19,7 +19,12 @@ by default. It accepts RelayKit UDP relay frames.
 Protocol
 --------
 
-Each UDP datagram is one `RKP1` frame:
+Each UDP datagram is encrypted with ChaCha20-Poly1305. The encrypted datagram
+uses the same combined format as CryptoKit `ChaChaPoly`:
+
+    nonce(12), ciphertext, tag(16)
+
+The decrypted plaintext is one `RKP1` frame:
 
     magic(4) = "RKP1", version(1) = 1, type(1), reserved(2),
     stream_id(8), payload_len(4), payload(payload_len)

apps/relay_server/src/chacha20_cipher.erl

@@ -0,0 +1,62 @@
+%%%-------------------------------------------------------------------
+%% @doc ChaCha20-Poly1305 helpers compatible with CryptoKit ChaChaPoly.
+%% @end
+%%%-------------------------------------------------------------------
+
+-module(chacha20_cipher).
+
+-export([encrypt/2, decrypt/2]).
+
+-define(KEY_BYTES, 32).
+-define(NONCE_BYTES, 12).
+-define(TAG_BYTES, 16).
+-define(AAD, <<>>).
+
+encrypt(PlainText, Key) when is_binary(PlainText) ->
+    case valid_key(Key) of
+        true ->
+            Nonce = crypto:strong_rand_bytes(?NONCE_BYTES),
+            try crypto:crypto_one_time_aead(
+                    chacha20_poly1305, Key, Nonce, PlainText, ?AAD, true
+                ) of
+                {CipherText, Tag} ->
+                    {ok, <<Nonce/binary, CipherText/binary, Tag/binary>>}
+            catch
+                _Class:_Reason ->
+                    {error, crypto_failed}
+            end;
+        false ->
+            {error, invalid_key}
+    end;
+encrypt(_PlainText, _Key) ->
+    {error, invalid_plaintext}.
+
+decrypt(Combined, Key) when is_binary(Combined), byte_size(Combined) >= ?NONCE_BYTES + ?TAG_BYTES ->
+    case valid_key(Key) of
+        true ->
+            <<Nonce:?NONCE_BYTES/binary, Rest/binary>> = Combined,
+            CipherTextBytes = byte_size(Rest) - ?TAG_BYTES,
+            <<CipherText:CipherTextBytes/binary, Tag:?TAG_BYTES/binary>> = Rest,
+            try crypto:crypto_one_time_aead(
+                    chacha20_poly1305, Key, Nonce, CipherText, ?AAD, Tag, false
+                ) of
+                error ->
+                    {error, authentication_failed};
+                PlainText ->
+                    {ok, PlainText}
+            catch
+                _Class:_Reason ->
+                    {error, crypto_failed}
+            end;
+        false ->
+            {error, invalid_key}
+    end;
+decrypt(Combined, _Key) when is_binary(Combined) ->
+    {error, ciphertext_too_short};
+decrypt(_Combined, _Key) ->
+    {error, invalid_ciphertext}.
+
+valid_key(Key) when is_binary(Key), byte_size(Key) =:= ?KEY_BYTES ->
+    true;
+valid_key(_Key) ->
+    false.

apps/relay_server/src/relay_server.app.src

@@ -9,6 +9,7 @@
     {applications, [
         kernel,
         stdlib,
+        crypto,
         esockd,
         sync
     ]},

apps/relay_server/src/relay_server_udp_handler.erl

@@ -13,6 +13,17 @@
 
 -define(TCP_OPTIONS, [binary, {packet, raw}, {active, once}, {nodelay, true}]).
 
+-define(CHACHA20_KEY, <<
+    16#9F, 16#4A, 16#6C, 16#0D,
+    16#8B, 16#21, 16#7E, 16#3F,
+    16#42, 16#D9, 16#AA, 16#78,
+    16#13, 16#C5, 16#E2, 16#B6,
+    16#F0, 16#A9, 16#27, 16#BC,
+    16#6D, 16#31, 16#E8, 16#4C,
+    16#55, 16#FA, 16#10, 16#2A,
+    16#7E, 16#9D, 16#3C, 16#BB
+    >>).
+
 -record(stream, {
     id :: non_neg_integer(),
     socket :: inet:socket(),
@@ -27,6 +38,7 @@
     peer :: {inet:ip_address(), inet:port_number()},
     idle_timeout :: timeout(),
     connect_timeout :: timeout(),
+    cipher_key :: binary(),
     users = #{} :: #{binary() => binary()},
     streams = #{} :: #{non_neg_integer() => stream()},
     sockets = #{} :: #{inet:socket() => non_neg_integer()}
@@ -45,6 +57,7 @@ init([{udp, Server, _Sock}, Peer, IdleTimeout, ConnectTimeout, Users]) ->
         peer = Peer,
         idle_timeout = IdleTimeout,
         connect_timeout = ConnectTimeout,
+        cipher_key = ?CHACHA20_KEY,
         users = normalize_users(Users)
     }, IdleTimeout}.
 
@@ -58,7 +71,7 @@ handle_info({datagram, Server, <<"stop">>}, State = #state{server = Server, peer
     logger:debug("UDP peer stopped: ~s", [esockd:format(Peer)]),
     {stop, normal, State};
 handle_info({datagram, Server, Packet}, State = #state{server = Server}) ->
-    handle_datagram(Packet, State);
+    handle_encrypted_datagram(Packet, State);
 
 handle_info({tcp, Socket, Data}, State) ->
     handle_remote_data(Socket, Data, State);
@@ -80,6 +93,15 @@ terminate(_Reason, #state{streams = Streams}) ->
 code_change(_OldVsn, State, _Extra) ->
     {ok, State}.
 
+handle_encrypted_datagram(Packet, State = #state{idle_timeout = IdleTimeout, cipher_key = Key}) ->
+    case chacha20_cipher:decrypt(Packet, Key) of
+        {ok, PlainPacket} ->
+            handle_datagram(PlainPacket, State);
+        {error, Reason} ->
+            logger:debug("UDP relay ignored undecryptable datagram: ~p", [Reason]),
+            {noreply, State, IdleTimeout}
+    end.
+
 handle_datagram(Packet, State = #state{idle_timeout = IdleTimeout}) ->
     case relay_server_udp_protocol:decode(Packet) of
         {ok, {open, StreamId, Payload}} ->
@@ -230,8 +252,17 @@ remove_stream(StreamId, State = #state{streams = Streams, sockets = Sockets}) ->
 close_remote(#stream{socket = Socket}) ->
     gen_tcp:close(Socket).
 
-send_frame(Type, StreamId, Payload, #state{server = Server, peer = Peer}) ->
-    Server ! {datagram, Peer, relay_server_udp_protocol:encode(Type, StreamId, Payload)}.
+send_frame(Type, StreamId, Payload, #state{server = Server, peer = Peer, cipher_key = Key}) ->
+    PlainPacket = relay_server_udp_protocol:encode(Type, StreamId, Payload),
+    case chacha20_cipher:encrypt(PlainPacket, Key) of
+        {ok, Packet} ->
+            Server ! {datagram, Peer, Packet},
+            ok;
+        {error, Reason} ->
+            logger:error("UDP relay failed to encrypt ~p frame for stream ~p: ~p",
+                [Type, StreamId, Reason]),
+            {error, Reason}
+    end.
 
 send_error(StreamId, Message, State) ->
     send_frame(error, StreamId, Message, State).