Refactor FS chacha into its own module

protocol/src/chacha20poly1305.rs

@@ -1,18 +1,42 @@
 mod chacha20;
 mod poly1305;
 
-use crate::error;
 pub(crate) use chacha20::ChaCha20;
 use poly1305::Poly1305;
 
-use error::ChaCha20Poly1305DecryptionError;
-use error::ChaCha20Poly1305EncryptionError;
-
-use alloc::string::ToString;
+use alloc::fmt;
 
 // Zero array for padding slices.
 const ZEROES: [u8; 16] = [0u8; 16];
 
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+pub enum Error {
+    UnauthenticatedAdditionalData,
+    CiphertextTooShort,
+    IncorrectBuffer,
+}
+
+impl fmt::Display for Error {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Error::UnauthenticatedAdditionalData => write!(f, "Unauthenticated aad."),
+            Error::CiphertextTooShort => write!(f, "Ciphertext must be at least 16 bytes."),
+            Error::IncorrectBuffer => write!(f, "The buffer provided was incorrect. Ensure the buffer is 16 bytes longer than the message."),
+        }
+    }
+}
+
+#[cfg(feature = "std")]
+impl std::error::Error for Error {
+    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
+        match self {
+            Error::UnauthenticatedAdditionalData => None,
+            Error::CiphertextTooShort => None,
+            Error::IncorrectBuffer => None,
+        }
+    }
+}
+
 #[derive(Debug)]
 pub struct ChaCha20Poly1305 {
     key: [u8; 32],
@@ -29,9 +53,9 @@ impl ChaCha20Poly1305 {
         plaintext: &'a mut [u8],
         aad: Option<&'a [u8]>,
         buffer: &'a mut [u8],
-    ) -> Result<&'a [u8], ChaCha20Poly1305EncryptionError> {
+    ) -> Result<&'a [u8], Error> {
         if plaintext.len() + 16 != buffer.len() {
-            return Err(ChaCha20Poly1305EncryptionError::IncorrectBuffer("The buffer provided was incorrect. Ensure the buffer is 16 bytes longer than the message.".to_string()));
+            return Err(Error::IncorrectBuffer);
         }
         let mut chacha = ChaCha20::new_from_block(self.key, self.nonce, 1);
         chacha.apply_keystream(plaintext);
@@ -81,7 +105,7 @@ impl ChaCha20Poly1305 {
         self,
         ciphertext: &'a mut [u8],
         aad: Option<&'a [u8]>,
-    ) -> Result<&'a [u8], ChaCha20Poly1305DecryptionError> {
+    ) -> Result<&'a [u8], Error> {
         let mut chacha = ChaCha20::new_from_block(self.key, self.nonce, 0);
         let keystream = chacha.get_keystream(0);
         let mut poly = Poly1305::new(
@@ -118,16 +142,10 @@ impl ChaCha20Poly1305 {
                 chacha.apply_keystream(received_msg);
                 Ok(received_msg)
             } else {
-                Err(
-                    ChaCha20Poly1305DecryptionError::UnauthenticatedAdditionalData(
-                        "Computed tag did not match.".to_string(),
-                    ),
-                )
+                Err(Error::UnauthenticatedAdditionalData)
             }
         } else {
-            Err(ChaCha20Poly1305DecryptionError::CiphertextTooShort(
-                "Ciphertext must be at least 16 bytes.".to_string(),
-            ))
+            Err(Error::CiphertextTooShort)
         }
     }
 }

protocol/src/fschacha20poly1305.rs

@@ -0,0 +1,167 @@
+use alloc::{fmt, vec::Vec};
+
+use crate::chacha20poly1305::{ChaCha20, ChaCha20Poly1305};
+
+const CHACHA_BLOCKS_USED: u32 = 3;
+pub(crate) const REKEY_INTERVAL: u32 = 224;
+const REKEY_INITIAL_NONCE: [u8; 4] = [0xFF, 0xFF, 0xFF, 0xFF];
+
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+pub enum Error {
+    Encryption,
+    Decryption,
+}
+
+impl fmt::Display for Error {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Error::Encryption => write!(f, "Unable to encrypt"),
+            Error::Decryption => write!(f, "Unable to dycrypt"),
+        }
+    }
+}
+
+#[cfg(feature = "std")]
+impl std::error::Error for Error {
+    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
+        match self {
+            Error::Encryption => None,
+            Error::Decryption => None,
+        }
+    }
+}
+
+pub enum CryptType {
+    Encrypt,
+    Decrypt,
+}
+
+/// A wrapper over ChaCha20Poly1305 AEAD stream cipher which handles automatically changing
+/// nonces and re-keying.
+///
+/// FSChaCha20Poly1305 is used for message packets in BIP324.
+#[derive(Clone, Debug)]
+pub struct FSChaCha20Poly1305 {
+    key: [u8; 32],
+    message_counter: u32,
+}
+
+impl FSChaCha20Poly1305 {
+    pub fn new(key: [u8; 32]) -> Self {
+        FSChaCha20Poly1305 {
+            key,
+            message_counter: 0,
+        }
+    }
+
+    fn crypt(
+        &mut self,
+        aad: Vec<u8>,
+        mut contents: Vec<u8>,
+        crypt_type: CryptType,
+    ) -> Result<Vec<u8>, Error> {
+        let mut counter_div = (self.message_counter / REKEY_INTERVAL)
+            .to_le_bytes()
+            .to_vec();
+        counter_div.extend([0u8; 4]); // ok? invalid for 4 billion messages
+        let counter_mod = (self.message_counter % REKEY_INTERVAL).to_le_bytes();
+        let mut nonce = counter_mod.to_vec();
+        nonce.extend(counter_div); // mod slice then div slice
+        let cipher =
+            ChaCha20Poly1305::new(self.key, nonce.try_into().expect("Nonce is malformed."));
+        let converted_ciphertext: Vec<u8> = match crypt_type {
+            CryptType::Encrypt => {
+                let mut buffer = contents.clone();
+                buffer.extend([0u8; 16]);
+                cipher
+                    .encrypt(&mut contents, Some(&aad), &mut buffer)
+                    .map_err(|_| Error::Encryption)?;
+                buffer.to_vec()
+            }
+            CryptType::Decrypt => {
+                let mut ciphertext = contents.clone();
+                cipher
+                    .decrypt(&mut ciphertext, Some(&aad))
+                    .map_err(|_| Error::Decryption)?;
+                ciphertext[..ciphertext.len() - 16].to_vec()
+            }
+        };
+        if (self.message_counter + 1) % REKEY_INTERVAL == 0 {
+            let mut rekey_nonce = REKEY_INITIAL_NONCE.to_vec();
+            let mut counter_div = (self.message_counter / REKEY_INTERVAL)
+                .to_le_bytes()
+                .to_vec();
+            counter_div.extend([0u8; 4]);
+            let counter_mod = (self.message_counter % REKEY_INTERVAL).to_le_bytes();
+            let mut nonce = counter_mod.to_vec();
+            nonce.extend(counter_div);
+            rekey_nonce.extend(nonce[4..].to_vec());
+            let mut buffer = [0u8; 48];
+            let mut plaintext = [0u8; 32];
+            let cipher = ChaCha20Poly1305::new(
+                self.key,
+                rekey_nonce.try_into().expect("Nonce is malformed."),
+            );
+            cipher
+                .encrypt(&mut plaintext, Some(&aad), &mut buffer)
+                .map_err(|_| Error::Encryption)?;
+            self.key = buffer[0..32]
+                .try_into()
+                .expect("Cipher should be at least 32 bytes.");
+        }
+        self.message_counter += 1;
+        Ok(converted_ciphertext)
+    }
+
+    pub fn encrypt(&mut self, aad: Vec<u8>, contents: Vec<u8>) -> Result<Vec<u8>, Error> {
+        self.crypt(aad, contents, CryptType::Encrypt)
+    }
+
+    pub fn decrypt(&mut self, aad: Vec<u8>, contents: Vec<u8>) -> Result<Vec<u8>, Error> {
+        self.crypt(aad, contents, CryptType::Decrypt)
+    }
+}
+
+/// A wrapper over ChaCha20 (unauthenticated) stream cipher which handles automatically changing
+/// nonces and re-keying.
+///
+/// FSChaCha20 is used for lengths in BIP324. Should be noted that the lengths are still
+/// implicitly authenticated by the message packets.
+#[derive(Clone, Debug)]
+pub struct FSChaCha20 {
+    key: [u8; 32],
+    block_counter: u32,
+    chunk_counter: u32,
+}
+
+impl FSChaCha20 {
+    pub fn new(key: [u8; 32]) -> Self {
+        FSChaCha20 {
+            key,
+            block_counter: 0,
+            chunk_counter: 0,
+        }
+    }
+
+    pub fn crypt(&mut self, chunk: Vec<u8>) -> Vec<u8> {
+        let zeroes = (0u32).to_le_bytes().to_vec();
+        let counter_mod = (self.chunk_counter / REKEY_INTERVAL).to_le_bytes();
+        let mut nonce = zeroes.clone();
+        nonce.extend(counter_mod);
+        nonce.extend(zeroes);
+        let mut cipher = ChaCha20::new(self.key, nonce.try_into().expect("Nonce is malformed."), 0);
+        let mut buffer = chunk.clone();
+        cipher.seek(self.block_counter);
+        cipher.apply_keystream(&mut buffer);
+        self.block_counter += CHACHA_BLOCKS_USED;
+        if (self.chunk_counter + 1) % REKEY_INTERVAL == 0 {
+            let mut key_buffer = [0u8; 32];
+            cipher.seek(self.block_counter);
+            cipher.apply_keystream(&mut key_buffer);
+            self.block_counter = 0;
+            self.key = key_buffer;
+        }
+        self.chunk_counter += 1;
+        buffer
+    }
+}