Skip to content

Commit

Permalink
fix(chacha20): byte conversions are infallible
Browse files Browse the repository at this point in the history
  • Loading branch information
rustaceanrob committed May 7, 2024
1 parent a82a180 commit 71b6ae1
Show file tree
Hide file tree
Showing 4 changed files with 110 additions and 154 deletions.
41 changes: 6 additions & 35 deletions protocol/src/chacha20poly1305.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,18 +14,12 @@ const ZEROES: [u8; 16] = [0u8; 16];
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum Error {
UnauthenticatedAdditionalData,
KeystreamMismatch,
Cipher(chacha20::Error),
AuthenticationTag(poly1305::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),
Error::AuthenticationTag(e) => write!(f, "Authentication tag error {}", e),
Error::KeystreamMismatch => write!(f, "Keystream is out of whack."),
}
}
}
Expand All @@ -35,25 +29,10 @@ 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),
Error::AuthenticationTag(e) => Some(e),
Error::KeystreamMismatch => None,
}
}
}

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

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

/// Encrypt and decrypt content along with a authentication tag.
#[derive(Debug)]
pub struct ChaCha20Poly1305 {
Expand All @@ -78,13 +57,9 @@ 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)?;
let mut poly = Poly1305::new(
keystream[..32]
.try_into()
.map_err(|_| Error::KeystreamMismatch)?,
)?;
chacha.apply_keystream(content);
let keystream = chacha.get_keystream(0);
let mut poly = Poly1305::new(keystream[..32].try_into().expect("infallible conversion"));
let aad = aad.unwrap_or(&[]);
// AAD and ciphertext are padded if not 16-byte aligned.
poly.add(aad);
Expand Down Expand Up @@ -126,12 +101,8 @@ 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 mut poly = Poly1305::new(
keystream[..32]
.try_into()
.map_err(|_| Error::KeystreamMismatch)?,
)?;
let keystream = chacha.get_keystream(0);
let mut poly = Poly1305::new(keystream[..32].try_into().expect("infallible conversion"));
let aad = aad.unwrap_or(&[]);
poly.add(aad);
// AAD and ciphertext are padded if not 16-byte aligned.
Expand All @@ -156,7 +127,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
159 changes: 89 additions & 70 deletions protocol/src/chacha20poly1305/chacha20.rs
Original file line number Diff line number Diff line change
@@ -1,30 +1,5 @@
//! The ChaCha20 stream cipher based on RFC7539.

use core::fmt;

/// Possible errors using the ChaCha20 cipher.
#[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 @@ -90,7 +65,7 @@ impl ChaCha20 {
}

/// Apply the keystream to a message.
pub fn apply_keystream(&mut self, to: &mut [u8]) -> Result<(), Error> {
pub fn apply_keystream(&mut self, to: &mut [u8]) {
let num_full_blocks = to.len() / CHACHA_BLOCKSIZE;
let mut j = 0;
while j < num_full_blocks {
Expand All @@ -99,7 +74,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 @@ -115,17 +90,16 @@ 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;
}
Ok(())
}

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

fn prepare_state(key: [u8; 32], nonce: [u8; 12], count: u32) -> Result<[u32; 16], Error> {
fn prepare_state(key: [u8; 32], nonce: [u8; 12], count: u32) -> [u32; 16] {
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().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[4] = u32::from_le_bytes(
key[0..4]
.try_into()
.expect("infalliable conversion to 4 byte array"),
);
state[5] = u32::from_le_bytes(
key[4..8]
.try_into()
.expect("infalliable conversion to 4 byte array"),
);
state[6] = u32::from_le_bytes(
key[8..12]
.try_into()
.expect("infalliable conversion to 4 byte array"),
);
state[7] = u32::from_le_bytes(
key[12..16]
.try_into()
.expect("infalliable conversion to 4 byte array"),
);
state[8] = u32::from_le_bytes(
key[16..20]
.try_into()
.expect("infalliable conversion to 4 byte array"),
);
state[9] = u32::from_le_bytes(
key[20..24]
.try_into()
.expect("infalliable conversion to 4 byte array"),
);
state[10] = u32::from_le_bytes(
key[24..28]
.try_into()
.expect("infalliable conversion to 4 byte array"),
);
state[11] = u32::from_le_bytes(
key[28..32]
.try_into()
.expect("infalliable conversion to 4 byte array"),
);
state[12] = count;
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)
state[13] = u32::from_le_bytes(
nonce[0..4]
.try_into()
.expect("infalliable conversion to 4 byte array"),
);
state[14] = u32::from_le_bytes(
nonce[4..8]
.try_into()
.expect("infalliable conversion to 4 byte array"),
);
state[15] = u32::from_le_bytes(
nonce[8..12]
.try_into()
.expect("infalliable conversion to 4 byte array"),
);
state
}

fn keystream_from_state(state: &mut [u32; 16]) -> [u8; 64] {
Expand All @@ -207,23 +225,18 @@ 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,
) -> Result<[u8; 64], Error> {
fn keystream_at_slice(key: [u8; 32], nonce: [u8; 12], count: u32, seek: usize) -> [u8; 64] {
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");
Ok(kstream)
let kstream: [u8; 64] = keystream[seek..seek + 64].try_into().expect("infallable");
kstream
}

#[cfg(test)]
Expand Down Expand Up @@ -341,7 +354,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).unwrap();
let state = prepare_state(key, nonce, count);
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 @@ -359,10 +372,11 @@ mod tests {
let count = 1;
let mut chacha = ChaCha20::new(key, nonce, count);
let mut binding = [8; 3];
chacha.apply_keystream(&mut binding[..]).unwrap();
let to = binding.as_mut_slice();
chacha.apply_keystream(to);
let mut chacha = ChaCha20::new(key, nonce, count);
chacha.apply_keystream(&mut binding[..]).unwrap();
assert_eq!([8; 3], binding);
chacha.apply_keystream(to);
assert_eq!([8; 3], to);
}

#[test]
Expand All @@ -375,10 +389,11 @@ mod tests {
let count = 1;
let mut chacha = ChaCha20::new(key, nonce, count);
let mut binding = [8; 64];
chacha.apply_keystream(&mut binding[..]).unwrap();
let to = binding.as_mut_slice();
chacha.apply_keystream(to);
let mut chacha = ChaCha20::new(key, nonce, count);
chacha.apply_keystream(&mut binding[..]).unwrap();
assert_eq!([8; 64], binding);
chacha.apply_keystream(to);
assert_eq!([8; 64], to);
}

#[test]
Expand All @@ -391,11 +406,13 @@ mod tests {
let count = 64;
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.";
chacha.apply_keystream(&mut binding[..]).unwrap();
assert_eq!(binding[..], Vec::from_hex("6e2e359a2568f98041ba0728dd0d6981e97e7aec1d4360c20a27afccfd9fae0bf91b65c5524733ab8f593dabcd62b3571639d624e65152ab8f530c359f0861d807ca0dbf500d6a6156a38e088a22b65e52bc514d16ccf806818ce91ab77937365af90bbf74a35be6b40b8eedf2785e42874d").unwrap());
let to = binding.as_mut_slice();
chacha.apply_keystream(to);
assert_eq!(to, Vec::from_hex("6e2e359a2568f98041ba0728dd0d6981e97e7aec1d4360c20a27afccfd9fae0bf91b65c5524733ab8f593dabcd62b3571639d624e65152ab8f530c359f0861d807ca0dbf500d6a6156a38e088a22b65e52bc514d16ccf806818ce91ab77937365af90bbf74a35be6b40b8eedf2785e42874d").unwrap());
let mut chacha = ChaCha20::new(key, nonce, count);
chacha.apply_keystream(&mut binding[..]).unwrap();
assert_eq!(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.");
chacha.apply_keystream(to);
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);
}

#[test]
Expand All @@ -408,11 +425,13 @@ mod tests {
let block: u32 = 1;
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.";
chacha.apply_keystream(&mut binding[..]).unwrap();
assert_eq!(binding[..], Vec::from_hex("6e2e359a2568f98041ba0728dd0d6981e97e7aec1d4360c20a27afccfd9fae0bf91b65c5524733ab8f593dabcd62b3571639d624e65152ab8f530c359f0861d807ca0dbf500d6a6156a38e088a22b65e52bc514d16ccf806818ce91ab77937365af90bbf74a35be6b40b8eedf2785e42874d").unwrap());
let to = binding.as_mut_slice();
chacha.apply_keystream(to);
assert_eq!(to, Vec::from_hex("6e2e359a2568f98041ba0728dd0d6981e97e7aec1d4360c20a27afccfd9fae0bf91b65c5524733ab8f593dabcd62b3571639d624e65152ab8f530c359f0861d807ca0dbf500d6a6156a38e088a22b65e52bc514d16ccf806818ce91ab77937365af90bbf74a35be6b40b8eedf2785e42874d").unwrap());
chacha.block(block);
chacha.apply_keystream(&mut binding[..]).unwrap();
assert_eq!(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.");
chacha.apply_keystream(to);
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);
}

#[cfg(feature = "std")]
Expand All @@ -436,11 +455,11 @@ mod tests {
let message = gen_garbage(129);
let mut message2 = message.clone();
let msg = message2.as_mut_slice();
chacha.apply_keystream(msg).unwrap();
chacha.apply_keystream(msg);
let mut cipher = ChaCha20::new(key, nonce, 0);
let mut buffer = message;
cipher.seek(count);
cipher.apply_keystream(&mut buffer).unwrap();
cipher.apply_keystream(&mut buffer);
assert_eq!(buffer.as_slice(), msg);
}
}
Expand Down
Loading

0 comments on commit 71b6ae1

Please sign in to comment.