From ffde489e53734f9a2061f50aada5c936cfcdd369 Mon Sep 17 00:00:00 2001
From: Luis Rascao <luis.rascao@miniclip.com>
Date: Tue, 27 Apr 2021 15:37:48 +0100
Subject: [PATCH] Fix rekey, improve coverage

ChaChaPoly key is expected to be 256 bits long. It's safe to disregard
the MAC portion.
---
 src/enoise_crypto.erl              | 4 ++++
 test/enoise_chiper_state_tests.erl | 9 +++++++++
 test/enoise_crypto_tests.erl       | 7 +++++++
 3 files changed, 20 insertions(+)

diff --git a/src/enoise_crypto.erl b/src/enoise_crypto.erl
index 8789310..80a727d 100644
--- a/src/enoise_crypto.erl
+++ b/src/enoise_crypto.erl
@@ -56,6 +56,10 @@ hkdf(Hash, Key, Data) ->
 
 -spec rekey(Cipher :: enoise_cipher_state:noise_cipher(),
             Key :: binary()) -> binary() | {error, term()}.
+rekey('ChaChaPoly', K0) ->
+    KLen = enacl:aead_chacha20poly1305_ietf_KEYBYTES(),
+    <<K:KLen/binary, _/binary>> = encrypt('ChaChaPoly', K0, ?MAX_NONCE, <<>>, <<0:(32*8)>>),
+    K;
 rekey(Cipher, K) ->
     encrypt(Cipher, K, ?MAX_NONCE, <<>>, <<0:(32*8)>>).
 
diff --git a/test/enoise_chiper_state_tests.erl b/test/enoise_chiper_state_tests.erl
index f5b49f0..64fbfed 100644
--- a/test/enoise_chiper_state_tests.erl
+++ b/test/enoise_chiper_state_tests.erl
@@ -26,5 +26,14 @@ chachapoly_test() ->
         enoise_cipher_state:decrypt_with_ad(CS1, AD, <<CipherText/binary, MAC/binary>>),
 
     ?assertMatch(PlainText, PlainText0),
+
+    % rekey test
+    CS4 = enoise_cipher_state:rekey(CS1),
+    {ok, _CS5, <<CipherText1:CTLen/binary, MAC1:MACLen/binary>>} =
+        enoise_cipher_state:encrypt_with_ad(CS4, AD, PlainText),
+    {ok, _CS6, <<PlainText1:PTLen/binary>>} =
+        enoise_cipher_state:decrypt_with_ad(CS4, AD, <<CipherText1/binary, MAC1/binary>>),
+    ?assertMatch(PlainText, PlainText1),
+
     ok.
 
diff --git a/test/enoise_crypto_tests.erl b/test/enoise_crypto_tests.erl
index d52e2eb..14be99a 100644
--- a/test/enoise_crypto_tests.erl
+++ b/test/enoise_crypto_tests.erl
@@ -44,6 +44,13 @@ chachapoly_test() ->
         enoise_crypto:decrypt('ChaChaPoly', Key, Nonce, AD, <<CipherText/binary, MAC/binary>>),
 
     ?assertMatch(PlainText, PlainText0),
+
+    Key1 = enoise_crypto:rekey('ChaChaPoly', Key),
+    <<CipherText1:CTLen/binary, MAC1:MACLen/binary>> =
+        enoise_crypto:encrypt('ChaChaPoly', Key1, Nonce, AD, PlainText),
+    <<PlainText1:PTLen/binary>> =
+        enoise_crypto:decrypt('ChaChaPoly', Key1, Nonce, AD, <<CipherText1/binary, MAC1/binary>>),
+    ?assertMatch(PlainText, PlainText1),
     ok.
 
 blake2b_test() ->