Merge pull request #27 from ssp-rs/fixup/encryption

fixup: encryption mode

Cargo.toml

@@ -37,7 +37,7 @@ features = ["alloc"]
 optional = true
 
 [dependencies.smol-jsonrpc]
-version = "0.1"
+version = "0.2"
 optional = true
 
 [features]

src/crc.rs

@@ -23,14 +23,24 @@ pub fn crc16(data: &[u8]) -> u16 {
             if crc & 0x8000 != 0 {
                 crc = (crc << 1) ^ CRC_POLY;
             } else {
-                crc = crc.overflowing_shl(1).0;
+                crc = saturating_shl(crc, 1);
             }
         }
     }
 
     crc
 }
 
+fn saturating_shl(b: u16, s: u32) -> u16 {
+    let (res, o) = b.overflowing_shl(s);
+
+    if o {
+        u16::MAX
+    } else {
+        res
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -98,4 +108,46 @@ mod tests {
         );
         assert_eq!(crc16(enc_data.as_ref()), 0x0000);
     }
+
+    #[test]
+    fn test_sample() {
+        let data_crc = [0xac, 0x1a];
+        let enc_data = [
+            0x00,
+            0x11,
+            0x7e,
+            0xe5,
+            0x65,
+            0x07,
+            0x0e,
+            0x2a,
+            0x8f,
+            0xab,
+            0xf7,
+            0xdd,
+            0xb3,
+            0x87,
+            0xe6,
+            0x45,
+            0xea,
+            0xb2,
+            0xbb,
+            data_crc[0],
+            data_crc[1],
+        ];
+        assert_eq!(
+            crc16(enc_data[..enc_data.len() - 2].as_ref()),
+            u16::from_be_bytes(data_crc)
+        );
+
+        let enc_crc = [0x95, 0xfb];
+        let enc_data = [
+            0x01, 0x05, 0x00, 0x00, 0x00, 0x07, 0x9c, 0x10, 0x69, 0xd9, 0x37, 0xf4, 0x18, 0x7f,
+            enc_crc[0], enc_crc[1],
+        ];
+        assert_eq!(
+            crc16(enc_data[..enc_data.len() - 2].as_ref()),
+            u16::from_le_bytes(enc_crc)
+        );
+    }
 }

src/encrypted.rs

@@ -31,6 +31,11 @@ pub fn sequence_count() -> SequenceCount {
     SEQUENCE_COUNT.load(Ordering::Relaxed).into()
 }
 
+/// Sets the current [SequenceCount].
+pub fn set_sequence_count(count: u32) {
+    SEQUENCE_COUNT.store(count, Ordering::SeqCst);
+}
+
 /// Increments the [SequenceCount].
 ///
 /// Returns the new [SequenceCount].
@@ -134,9 +139,9 @@ pub fn unstuff(buf: &mut [u8], mut end: usize) -> Result<usize> {
     Ok(end)
 }
 
-#[cfg(test)]
+#[cfg(all(test, feature = "std"))]
 pub mod tests {
-    use crate::{len, AesKey, MessageOps, Result};
+    use crate::{len, std::sync::Mutex, AesKey, MessageOps, Result};
 
     use super::*;
 
@@ -145,6 +150,9 @@ pub mod tests {
         b'p',
     ];
 
+    // Add a mutex to run tests one at a time
+    static T: Mutex<()> = Mutex::new(());
+
     fn test_key() -> AesKey {
         AesKey::clone_from_slice(TEST_KEY.as_ref())
     }
@@ -153,6 +161,8 @@ pub mod tests {
     fn test_command_encryption() -> Result<()> {
         use crate::PollCommand;
 
+        let _lock = T.lock();
+
         let key = test_key();
 
         let mut poll_msg = PollCommand::new();
@@ -175,6 +185,8 @@ pub mod tests {
     fn test_response_encryption() -> Result<()> {
         use crate::{PollResponse, ResponseOps, ResponseStatus};
 
+        let _lock = T.lock();
+
         let key = test_key();
 
         let mut poll_msg = PollResponse::new();
@@ -197,6 +209,8 @@ pub mod tests {
 
     #[test]
     fn test_byte_stuffing() -> Result<()> {
+        let _lock = T.lock();
+
         let mut buf = [0x7f, 0xaa, 0xbb, 0x00];
         let exp = [0x7f, 0x7f, 0xaa, 0xbb];
         let end = 2;
@@ -234,6 +248,8 @@ pub mod tests {
 
     #[test]
     fn test_byte_unstuffing() -> Result<()> {
+        let _lock = T.lock();
+
         let mut buf = [0x7f, 0x7f, 0xaa, 0xbb];
         let exp = [0x7f, 0xaa, 0xbb, 0x00];
         let end = buf.len() - 1;
@@ -268,4 +284,71 @@ pub mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_encrypt_decrypt_stuffing() -> Result<()> {
+        use crate::PollCommand;
+
+        let _lock = T.lock();
+
+        let mut msg = PollCommand::new();
+        let _clear_csum = msg.calculate_checksum();
+
+        let key = AesKey::from([
+            0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x10, 0x11, 0x12, 0x13, 0x14,
+            0x15, 0x16,
+        ]);
+
+        #[cfg(feature = "std")]
+        println!("Clear-text command buffer: {:x?}", msg.buf());
+
+        let mut enc_msg = EncryptedCommand::new();
+
+        enc_msg.set_message_data(&mut msg)?;
+        let wrap_msg = enc_msg.encrypt(&key);
+
+        let _wrap_csum = wrap_msg.checksum();
+
+        #[cfg(feature = "std")]
+        println!("Wrapped encrypted command buffer: {:x?}", wrap_msg.buf());
+
+        let dec_msg = EncryptedResponse::decrypt(&key, wrap_msg);
+
+        dec_msg.verify_checksum()?;
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_encrypt_known_keys() -> Result<()> {
+        use crate::{HostProtocolVersionCommand, ProtocolVersion};
+
+        let _lock = T.lock();
+
+        let key = AesKey::from([
+            0x67, 0x45, 0x23, 0x01, 0x67, 0x45, 0x23, 0x01, 0x5e, 0xfa, 0xe5, 0x0d, 0x00, 0x00,
+            0x00, 0x00,
+        ]);
+
+        let exp_enc_bytes = [
+            0x7f, 0x80, 0x11, 0x7e, 0x86, 0x01, 0xf3, 0xa7, 0x85, 0xec, 0x6f, 0x4b, 0x3d, 0x0f,
+            0xf0, 0xfe, 0xd4, 0x8d, 0x16, 0x64, 0x2a, 0x23,
+        ];
+
+        let msg = HostProtocolVersionCommand::new().with_version(ProtocolVersion::Six);
+
+        set_sequence_count(0);
+
+        let mut wrap_msg = EncryptedCommand::new()
+            .with_count(0u32.into())
+            .with_message_data(&msg)?
+            .encrypt(&key);
+
+        wrap_msg.set_sequence_id(128u8.into());
+        wrap_msg.calculate_checksum();
+
+        assert_eq!(wrap_msg.buf(), exp_enc_bytes.as_ref());
+
+        wrap_msg.verify_checksum()
+    }
 }

src/encrypted/command.rs

@@ -47,6 +47,17 @@ impl EncryptedCommand {
         self.buf[index::COUNT..index::COUNT_END].as_ref()
     }
 
+    fn set_count(&mut self, count: SequenceCount) {
+        self.buf[index::COUNT..index::COUNT_END]
+            .copy_from_slice(count.as_inner().to_le_bytes().as_ref());
+    }
+
+    /// Builder function that sets the [SequenceCount].
+    pub fn with_count(mut self, count: SequenceCount) -> Self {
+        self.set_count(count);
+        self
+    }
+
     /// Gets the message data.
     pub fn message_data(&self) -> &[u8] {
         let start = self.data_start();
@@ -75,29 +86,37 @@ impl EncryptedCommand {
     /// encrypted message is wrapped in an outer standard SSP message.
     ///
     /// Matryoshka dolls all the way down...
-    pub fn set_message_data(&mut self, message: &mut dyn CommandOps) -> Result<()> {
+    pub fn set_message_data(&mut self, message: &dyn CommandOps) -> Result<()> {
         let len = message.data_len();
 
         if message.data().len() != len {
             return Err(Error::InvalidDataLength((len, message.data().len())));
         }
 
-        message.calculate_checksum();
-
         if (0..=len::MAX_ENCRYPTED_DATA).contains(&len) {
-            self.buf[index::LEN] = len as u8;
+            self.set_data_len(len as u8);
 
             let start = self.data_start();
             let end = self.data_end();
 
-            self.buf[start..end].copy_from_slice(message.data());
+            let data = message.data();
+
+            log::trace!("Encrypted data: {data:x?}, length: {len}");
+
+            self.buf[start..end].copy_from_slice(data);
 
             Ok(())
         } else {
             Err(Error::InvalidDataLength((len, len::MAX_ENCRYPTED_DATA)))
         }
     }
 
+    /// Builder function that sets the message data.
+    pub fn with_message_data(mut self, message: &dyn CommandOps) -> Result<Self> {
+        self.set_message_data(message)?;
+        Ok(self)
+    }
+
     /// Gets random packing data used to make the encrypted packet a mutliple of the [AES block
     /// length](crate::len::AES).
     ///
@@ -133,7 +152,10 @@ impl EncryptedCommand {
         let start = self.packing_start();
         let end = self.packing_end();
 
+        #[cfg(not(test))]
         rng.fill_bytes(&mut self.buf[start..end]);
+        #[cfg(test)]
+        self.buf[start..end].copy_from_slice([0; 255][..end - start].as_ref());
     }
 
     /// Adds random packing data to make the encrypted packet a mutliple of the [AES block
@@ -155,7 +177,10 @@ impl EncryptedCommand {
         let start = self.packing_start();
         let end = self.packing_end();
 
+        #[cfg(not(test))]
         rng.fill_bytes(&mut self.buf[start..end]);
+        #[cfg(test)]
+        self.buf[start..end].copy_from_slice([0; 255][..end - start].as_ref());
     }
 
     fn packing_start(&self) -> usize {
@@ -185,30 +210,52 @@ impl EncryptedCommand {
     ///
     /// Converts the [EncryptedCommand] message into a standard [WrappedEncryptedMessage].
     pub fn encrypt(mut self, key: &AesKey) -> WrappedEncryptedMessage {
-        use crate::aes;
+        //use crate::aes;
+        use aes::cipher::{BlockEncrypt, KeyInit};
 
         self.set_packing();
+
+        if super::sequence_count().as_inner() != 0 {
+            self.set_count(super::sequence_count());
+        }
+
         self.calculate_checksum();
 
         let mut enc_msg = WrappedEncryptedMessage::new();
 
         let enc_len = self.len();
         enc_msg.set_data_len(enc_len as u8);
 
+        log::trace!("Encrypted message: {:x?}", self.buf());
+
         let plain_data = self.encrypt_data();
         let cipher_data = enc_msg.data_mut()[1..].as_mut();
 
-        if let Err(err) = aes::aes_encrypt_inplace(key.as_ref(), plain_data, cipher_data) {
-            log::error!("error encrypting message: {err}");
+        let ciph = aes::Aes128::new(key);
+
+        for (pchunk, cchunk) in plain_data
+            .chunks_exact(16)
+            .zip(cipher_data.chunks_exact_mut(16))
+        {
+            ciph.encrypt_block_b2b(pchunk.into(), cchunk.into());
         }
 
-        super::increment_sequence_count();
-        log::trace!("encryption sequence count: {}", super::sequence_count());
+        enc_msg.calculate_checksum();
+        if let Err(err) = enc_msg.verify_checksum() {
+            log::error!("error validating wrapped encrypted checksum: {err}");
+        }
 
         if let Err(err) = enc_msg.stuff_encrypted_data() {
             log::error!("error stuffing encrypted command message: {err}");
         }
 
+        log::trace!("encryption sequence count: {}", super::sequence_count());
+        #[cfg(any(not(test), feature = "test-crypto"))]
+        log::trace!(
+            "next encryption sequence count: {}",
+            super::increment_sequence_count()
+        );
+
         enc_msg
     }
 

src/encrypted/response.rs

@@ -233,8 +233,15 @@ impl EncryptedResponse {
             log::error!("error decrypting response message: {err}");
         }
 
-        super::increment_sequence_count();
-        log::trace!("decryption sequence count: {}", super::sequence_count());
+        log::trace!("decrypted data: {:x?}", plain_data);
+
+        let seq_count = super::sequence_count();
+        let dec_count = dec_msg.count();
+
+        if seq_count != dec_count {
+            log::error!("decryption sequence count is out of sync: have: {dec_count}, expected: {seq_count}");
+            super::set_sequence_count(dec_count.as_inner());
+        }
 
         dec_msg
     }