Skip to content

Commit

Permalink
Merge pull request #34 from nyonson/some-error-returns
Browse files Browse the repository at this point in the history
Clean up chacha20 expects into errors
  • Loading branch information
nyonson authored Apr 18, 2024
2 parents 5dd56ae + a3472e5 commit 14958fb
Show file tree
Hide file tree
Showing 3 changed files with 87 additions and 75 deletions.
17 changes: 13 additions & 4 deletions protocol/src/chacha20poly1305.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,12 +14,14 @@ const ZEROES: [u8; 16] = [0u8; 16];
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum Error {
UnauthenticatedAdditionalData,
Cipher(chacha20::Error),
}

impl fmt::Display for Error {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Error::UnauthenticatedAdditionalData => write!(f, "Unauthenticated aad."),
Error::Cipher(e) => write!(f, "Cipher encryption/decrytion error {}", e),
}
}
}
Expand All @@ -29,10 +31,17 @@ impl std::error::Error for Error {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Error::UnauthenticatedAdditionalData => None,
Error::Cipher(e) => Some(e),
}
}
}

impl From<chacha20::Error> for Error {
fn from(e: chacha20::Error) -> Self {
Error::Cipher(e)
}
}

/// Encrypt and decrypt content along with a authentication tag.
#[derive(Debug)]
pub struct ChaCha20Poly1305 {
Expand All @@ -57,8 +66,8 @@ impl ChaCha20Poly1305 {
/// The 16-byte authentication tag.
pub fn encrypt(self, content: &mut [u8], aad: Option<&[u8]>) -> Result<[u8; 16], Error> {
let mut chacha = ChaCha20::new_from_block(self.key, self.nonce, 1);
chacha.apply_keystream(content);
let keystream = chacha.get_keystream(0);
chacha.apply_keystream(content)?;
let keystream = chacha.get_keystream(0)?;
let mut poly = Poly1305::new(
keystream[..32]
.try_into()
Expand Down Expand Up @@ -105,7 +114,7 @@ impl ChaCha20Poly1305 {
aad: Option<&[u8]>,
) -> Result<(), Error> {
let mut chacha = ChaCha20::new_from_block(self.key, self.nonce, 0);
let keystream = chacha.get_keystream(0);
let keystream = chacha.get_keystream(0)?;
let mut poly = Poly1305::new(
keystream[..32]
.try_into()
Expand Down Expand Up @@ -135,7 +144,7 @@ impl ChaCha20Poly1305 {
let derived_tag = poly.tag();
if derived_tag.eq(&tag) {
let mut chacha = ChaCha20::new_from_block(self.key, self.nonce, 1);
chacha.apply_keystream(content);
chacha.apply_keystream(content)?;
Ok(())
} else {
Err(Error::UnauthenticatedAdditionalData)
Expand Down
132 changes: 63 additions & 69 deletions protocol/src/chacha20poly1305/chacha20.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,30 @@
//! The ChaCha20 stream cipher based on RFC7539.
use alloc::fmt;

/// Errors encrypting and decrypting messages with ChaCha20 and Poly1305 authentication tags.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum Error {
InvalidKey,
}

impl fmt::Display for Error {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Error::InvalidKey => write!(f, "Invalid 32-byte key."),
}
}
}

#[cfg(feature = "std")]
impl std::error::Error for Error {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Error::InvalidKey => None,
}
}
}

/// The first four words (32-bit) of the ChaCha stream cipher state are constants.
const WORD_1: u32 = 0x61707865;
const WORD_2: u32 = 0x3320646e;
Expand Down Expand Up @@ -65,7 +90,7 @@ impl ChaCha20 {
}

/// Apply the keystream to a message.
pub fn apply_keystream<'a>(&'a mut self, to: &'a mut [u8]) -> &[u8] {
pub fn apply_keystream(&mut self, to: &mut [u8]) -> Result<(), Error> {
let num_full_blocks = to.len() / CHACHA_BLOCKSIZE;
let mut j = 0;
while j < num_full_blocks {
Expand All @@ -74,7 +99,7 @@ impl ChaCha20 {
self.nonce,
self.block_count,
self.seek_offset_bytes,
);
)?;
for (c, k) in to[j * CHACHA_BLOCKSIZE..(j + 1) * CHACHA_BLOCKSIZE]
.iter_mut()
.zip(kstream.iter())
Expand All @@ -90,17 +115,17 @@ impl ChaCha20 {
self.nonce,
self.block_count,
self.seek_offset_bytes,
);
)?;
for (c, k) in to[j * CHACHA_BLOCKSIZE..].iter_mut().zip(kstream.iter()) {
*c ^= *k
}
self.block_count += 1;
}
to
Ok(())
}

/// Get the keystream block at a specified block.
pub(crate) fn get_keystream(&mut self, block: u32) -> [u8; 64] {
pub(crate) fn get_keystream(&mut self, block: u32) -> Result<[u8; 64], Error> {
self.block(block);
keystream_at_slice(
self.key,
Expand Down Expand Up @@ -150,61 +175,25 @@ fn chacha_block(state: &mut [u32; 16]) {
}
}

fn prepare_state(key: [u8; 32], nonce: [u8; 12], count: u32) -> [u32; 16] {
fn prepare_state(key: [u8; 32], nonce: [u8; 12], count: u32) -> Result<[u32; 16], Error> {
let mut state: [u32; 16] = [0; 16];
state[0] = WORD_1;
state[1] = WORD_2;
state[2] = WORD_3;
state[3] = WORD_4;
state[4] = u32::from_le_bytes(key[0..4].try_into().expect("Valid slice of 32 byte array."));
state[5] = u32::from_le_bytes(key[4..8].try_into().expect("Valid slice of 32 byte array."));
state[6] = u32::from_le_bytes(
key[8..12]
.try_into()
.expect("Valid slice of 32 byte array."),
);
state[7] = u32::from_le_bytes(
key[12..16]
.try_into()
.expect("Valid slice of 32 byte array."),
);
state[8] = u32::from_le_bytes(
key[16..20]
.try_into()
.expect("Valid slice of 32 byte array."),
);
state[9] = u32::from_le_bytes(
key[20..24]
.try_into()
.expect("Valid slice of 32 byte array."),
);
state[10] = u32::from_le_bytes(
key[24..28]
.try_into()
.expect("Valid slice of 32 byte array."),
);
state[11] = u32::from_le_bytes(
key[28..32]
.try_into()
.expect("Valid slice of 32 byte array."),
);
state[4] = u32::from_le_bytes(key[0..4].try_into().map_err(|_| Error::InvalidKey)?);
state[5] = u32::from_le_bytes(key[4..8].try_into().map_err(|_| Error::InvalidKey)?);
state[6] = u32::from_le_bytes(key[8..12].try_into().map_err(|_| Error::InvalidKey)?);
state[7] = u32::from_le_bytes(key[12..16].try_into().map_err(|_| Error::InvalidKey)?);
state[8] = u32::from_le_bytes(key[16..20].try_into().map_err(|_| Error::InvalidKey)?);
state[9] = u32::from_le_bytes(key[20..24].try_into().map_err(|_| Error::InvalidKey)?);
state[10] = u32::from_le_bytes(key[24..28].try_into().map_err(|_| Error::InvalidKey)?);
state[11] = u32::from_le_bytes(key[28..32].try_into().map_err(|_| Error::InvalidKey)?);
state[12] = count;
state[13] = u32::from_le_bytes(
nonce[0..4]
.try_into()
.expect("Valid slice of 32 byte array."),
);
state[14] = u32::from_le_bytes(
nonce[4..8]
.try_into()
.expect("Valid slice of 32 byte array."),
);
state[15] = u32::from_le_bytes(
nonce[8..12]
.try_into()
.expect("Valid slice of 32 byte array."),
);
state
state[13] = u32::from_le_bytes(nonce[0..4].try_into().map_err(|_| Error::InvalidKey)?);
state[14] = u32::from_le_bytes(nonce[4..8].try_into().map_err(|_| Error::InvalidKey)?);
state[15] = u32::from_le_bytes(nonce[8..12].try_into().map_err(|_| Error::InvalidKey)?);
Ok(state)
}

fn keystream_from_state(state: &mut [u32; 16]) -> [u8; 64] {
Expand All @@ -218,18 +207,23 @@ fn keystream_from_state(state: &mut [u32; 16]) -> [u8; 64] {
keystream
}

fn keystream_at_slice(key: [u8; 32], nonce: [u8; 12], count: u32, seek: usize) -> [u8; 64] {
fn keystream_at_slice(
key: [u8; 32],
nonce: [u8; 12],
count: u32,
seek: usize,
) -> Result<[u8; 64], Error> {
let mut keystream: [u8; 128] = [0; 128];
let mut state = prepare_state(key, nonce, count);
let mut state = prepare_state(key, nonce, count)?;
chacha_block(&mut state);
let first_half = keystream_from_state(&mut state);
let mut state = prepare_state(key, nonce, count + 1);
let mut state = prepare_state(key, nonce, count + 1)?;
chacha_block(&mut state);
let second_half = keystream_from_state(&mut state);
keystream[..64].copy_from_slice(&first_half);
keystream[64..].copy_from_slice(&second_half);
let kstream: [u8; 64] = keystream[seek..seek + 64].try_into().expect("msg");
kstream
Ok(kstream)
}

#[cfg(test)]
Expand Down Expand Up @@ -346,7 +340,7 @@ mod tests {
let nonce = Vec::from_hex("000000090000004a00000000").unwrap();
let nonce: [u8; 12] = nonce.try_into().unwrap();
let count = 1;
let state = prepare_state(key, nonce, count);
let state = prepare_state(key, nonce, count).unwrap();
assert_eq!(state[4].to_be_bytes().to_lower_hex_string(), "03020100");
assert_eq!(state[10].to_be_bytes().to_lower_hex_string(), "1b1a1918");
assert_eq!(state[14].to_be_bytes().to_lower_hex_string(), "4a000000");
Expand All @@ -365,9 +359,9 @@ mod tests {
let mut chacha = ChaCha20::new(key, nonce, count);
let mut binding = [8; 3];
let to = binding.as_mut_slice();
chacha.apply_keystream(to);
chacha.apply_keystream(to).unwrap();
let mut chacha = ChaCha20::new(key, nonce, count);
chacha.apply_keystream(to);
chacha.apply_keystream(to).unwrap();
assert_eq!([8; 3], to);
}

Expand All @@ -382,9 +376,9 @@ mod tests {
let mut chacha = ChaCha20::new(key, nonce, count);
let mut binding = [8; 64];
let to = binding.as_mut_slice();
chacha.apply_keystream(to);
chacha.apply_keystream(to).unwrap();
let mut chacha = ChaCha20::new(key, nonce, count);
chacha.apply_keystream(to);
chacha.apply_keystream(to).unwrap();
assert_eq!([8; 64], to);
}

Expand All @@ -399,10 +393,10 @@ mod tests {
let mut chacha = ChaCha20::new(key, nonce, count);
let mut binding = *b"Ladies and Gentlemen of the class of '99: If I could offer you only one tip for the future, sunscreen would be it.";
let to = binding.as_mut_slice();
chacha.apply_keystream(to);
chacha.apply_keystream(to).unwrap();
assert_eq!(to, Vec::from_hex("6e2e359a2568f98041ba0728dd0d6981e97e7aec1d4360c20a27afccfd9fae0bf91b65c5524733ab8f593dabcd62b3571639d624e65152ab8f530c359f0861d807ca0dbf500d6a6156a38e088a22b65e52bc514d16ccf806818ce91ab77937365af90bbf74a35be6b40b8eedf2785e42874d").unwrap());
let mut chacha = ChaCha20::new(key, nonce, count);
chacha.apply_keystream(to);
chacha.apply_keystream(to).unwrap();
let binding = *b"Ladies and Gentlemen of the class of '99: If I could offer you only one tip for the future, sunscreen would be it.";
assert_eq!(binding, to);
}
Expand All @@ -418,10 +412,10 @@ mod tests {
let mut chacha = ChaCha20::new_from_block(key, nonce, block);
let mut binding = *b"Ladies and Gentlemen of the class of '99: If I could offer you only one tip for the future, sunscreen would be it.";
let to = binding.as_mut_slice();
chacha.apply_keystream(to);
chacha.apply_keystream(to).unwrap();
assert_eq!(to, Vec::from_hex("6e2e359a2568f98041ba0728dd0d6981e97e7aec1d4360c20a27afccfd9fae0bf91b65c5524733ab8f593dabcd62b3571639d624e65152ab8f530c359f0861d807ca0dbf500d6a6156a38e088a22b65e52bc514d16ccf806818ce91ab77937365af90bbf74a35be6b40b8eedf2785e42874d").unwrap());
chacha.block(block);
chacha.apply_keystream(to);
chacha.apply_keystream(to).unwrap();
let binding = *b"Ladies and Gentlemen of the class of '99: If I could offer you only one tip for the future, sunscreen would be it.";
assert_eq!(binding, to);
}
Expand All @@ -447,11 +441,11 @@ mod tests {
let message = gen_garbage(129);
let mut message2 = message.clone();
let msg = message2.as_mut_slice();
chacha.apply_keystream(msg);
chacha.apply_keystream(msg).unwrap();
let mut cipher = ChaCha20::new(key, nonce, 0);
let mut buffer = message;
cipher.seek(count);
cipher.apply_keystream(&mut buffer);
cipher.apply_keystream(&mut buffer).unwrap();
assert_eq!(buffer.as_slice(), msg);
}
}
Expand Down
13 changes: 11 additions & 2 deletions protocol/src/fschacha20poly1305.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,13 +12,15 @@ const REKEY_INITIAL_NONCE: [u8; 4] = [0xFF, 0xFF, 0xFF, 0xFF];
pub enum Error {
Encryption,
Decryption,
Cipher(crate::chacha20poly1305::chacha20::Error),
}

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"),
Error::Cipher(e) => write!(f, "Cipher encryption/decrytion error {}", e),
}
}
}
Expand All @@ -29,10 +31,17 @@ impl std::error::Error for Error {
match self {
Error::Encryption => None,
Error::Decryption => None,
Error::Cipher(e) => Some(e),
}
}
}

impl From<crate::chacha20poly1305::chacha20::Error> for Error {
fn from(e: crate::chacha20poly1305::chacha20::Error) -> Self {
Error::Cipher(e)
}
}

/// A wrapper over ChaCha20Poly1305 AEAD stream cipher which handles automatically changing
/// nonces and re-keying.
///
Expand Down Expand Up @@ -151,12 +160,12 @@ impl FSChaCha20 {
nonce[4..8].copy_from_slice(&counter_mod);
let mut cipher = ChaCha20::new(self.key, nonce, 0);
cipher.seek(self.block_counter);
cipher.apply_keystream(chunk);
cipher.apply_keystream(chunk)?;
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);
cipher.apply_keystream(&mut key_buffer)?;
self.block_counter = 0;
self.key = key_buffer;
}
Expand Down

0 comments on commit 14958fb

Please sign in to comment.