ssh-cipher: add Decryptor::peek

As discussed in RustCrypto/block-modes#91

Author: tarcieri

Cargo.lock

@@ -26,7 +26,7 @@ encoding = { package = "ssh-encoding", version = "0.3.0-rc.9" }
 aead = { version = "0.6.0-rc.10", optional = true, default-features = false }
 aes = { version = "0.9", optional = true, default-features = false }
 aes-gcm = { version = "0.11.0-rc.3", optional = true, default-features = false, features = ["aes"] }
-cbc = { version = "0.2", optional = true }
+cbc = { version = "0.2.1", optional = true }
 ctr = { version = "0.10", optional = true, default-features = false }
 ctutils = { version = "0.4", optional = true, default-features = false }
 chacha20 = { version = "0.10", optional = true, default-features = false, features = ["cipher", "legacy"] }

ssh-cipher/Cargo.toml

@@ -1,20 +1,20 @@
 //! Stateful decryptor object.
 
 use crate::{Cipher, Error, Result};
-use cipher::KeyIvInit;
+use cipher::{KeyIvInit, SetIvState};
 use core::fmt::{self, Debug};
 
-#[cfg(feature = "aes-ctr")]
-use crate::{Ctr128BE, encryptor::ctr_encrypt as ctr_decrypt};
 #[cfg(any(feature = "aes-cbc", feature = "aes-ctr"))]
 use aes::{Aes128, Aes192, Aes256};
 #[cfg(any(feature = "aes-cbc", feature = "tdes"))]
 use cipher::{
-    Block,
+    Block, IvState, StreamCipherSeek,
     block::{BlockCipherDecrypt, BlockModeDecrypt},
 };
 #[cfg(feature = "tdes")]
 use des::TdesEde3;
+#[cfg(feature = "aes-ctr")]
+use {crate::encryptor::ctr_encrypt as ctr_decrypt, ctr::Ctr128BE};
 
 /// Stateful decryptor object for unauthenticated SSH symmetric ciphers.
 ///
@@ -43,6 +43,16 @@ enum Inner {
     TDesCbc(cbc::Decryptor<TdesEde3>),
 }
 
+/// Current IV state or position within the cipher.
+enum State {
+    #[cfg(feature = "aes-cbc")]
+    AesCbc(aes::Block),
+    #[cfg(feature = "aes-ctr")]
+    AesCtr(u64),
+    #[cfg(feature = "tdes")]
+    TDesCbc(Block<TdesEde3>),
+}
+
 impl Decryptor {
     /// Create a new decryptor object with the given [`Cipher`], `key`, and `iv` (i.e.
     /// initialization vector).
@@ -124,6 +134,77 @@ impl Decryptor {
 
         Ok(())
     }
+
+    /// Call the provided function with an ephemeral [`Decryptor`] state which will be reset upon
+    /// completion, returning the result of the function.
+    pub fn peek<T, F>(&mut self, mut f: F) -> Result<T>
+    where
+        F: FnMut(&mut Self) -> Result<T>,
+    {
+        let state = self.state();
+        let ret = f(self);
+        self.set_state(state)?;
+        ret
+    }
+
+    /// Get the current cipher state, i.e. IV or position within the stream cipher.
+    fn state(&self) -> State {
+        match &self.inner {
+            #[cfg(feature = "aes-cbc")]
+            Inner::Aes128Cbc(cipher) => State::AesCbc(cipher.iv_state()),
+            #[cfg(feature = "aes-cbc")]
+            Inner::Aes192Cbc(cipher) => State::AesCbc(cipher.iv_state()),
+            #[cfg(feature = "aes-cbc")]
+            Inner::Aes256Cbc(cipher) => State::AesCbc(cipher.iv_state()),
+            #[cfg(feature = "aes-ctr")]
+            Inner::Aes128Ctr(cipher) => State::AesCtr(cipher.current_pos()),
+            #[cfg(feature = "aes-ctr")]
+            Inner::Aes192Ctr(cipher) => State::AesCtr(cipher.current_pos()),
+            #[cfg(feature = "aes-ctr")]
+            Inner::Aes256Ctr(cipher) => State::AesCtr(cipher.current_pos()),
+            #[cfg(feature = "tdes")]
+            Inner::TDesCbc(cipher) => State::TDesCbc(cipher.iv_state()),
+        }
+    }
+
+    /// Set the current cipher state.
+    fn set_state(&mut self, state: State) -> Result<()> {
+        match (&mut self.inner, state) {
+            #[cfg(feature = "aes-cbc")]
+            (Inner::Aes128Cbc(cipher), State::AesCbc(iv)) => {
+                cipher.set_iv(&iv);
+                Ok(())
+            }
+            #[cfg(feature = "aes-cbc")]
+            (Inner::Aes192Cbc(cipher), State::AesCbc(iv)) => {
+                cipher.set_iv(&iv);
+                Ok(())
+            }
+            #[cfg(feature = "aes-cbc")]
+            (Inner::Aes256Cbc(cipher), State::AesCbc(iv)) => {
+                cipher.set_iv(&iv);
+                Ok(())
+            }
+            #[cfg(feature = "aes-ctr")]
+            (Inner::Aes128Ctr(cipher), State::AesCtr(pos)) => {
+                cipher.try_seek(pos).map_err(|_| Error::Crypto)
+            }
+            #[cfg(feature = "aes-ctr")]
+            (Inner::Aes192Ctr(cipher), State::AesCtr(pos)) => {
+                cipher.try_seek(pos).map_err(|_| Error::Crypto)
+            }
+            #[cfg(feature = "aes-ctr")]
+            (Inner::Aes256Ctr(cipher), State::AesCtr(pos)) => {
+                cipher.try_seek(pos).map_err(|_| Error::Crypto)
+            }
+            #[cfg(feature = "tdes")]
+            (Inner::TDesCbc(cipher), State::TDesCbc(iv)) => {
+                cipher.set_iv(&iv);
+                Ok(())
+            }
+            _ => Err(Error::Crypto), // should be unreachable
+        }
+    }
 }
 
 impl Debug for Decryptor {

ssh-cipher/src/decryptor.rs

@@ -1,7 +1,7 @@
 //! Stateful encryptor object.
 
 use crate::{Cipher, Error, Result};
-use cipher::{Block, BlockCipherEncrypt, KeyIvInit};
+use cipher::{Block, BlockCipherEncrypt, KeyIvInit, StreamCipher};
 use core::fmt::{self, Debug};
 
 #[cfg(any(feature = "aes-cbc", feature = "aes-ctr"))]
@@ -12,8 +12,8 @@ use cipher::block::BlockModeEncrypt;
 use des::TdesEde3;
 #[cfg(feature = "aes-ctr")]
 use {
-    crate::Ctr128BE,
-    cipher::{BlockSizeUser, StreamCipherCore, array::sizes::U16},
+    cipher::{BlockSizeUser, array::sizes::U16},
+    ctr::Ctr128BE,
 };
 
 /// Stateful encryptor object for unauthenticated SSH symmetric ciphers.
@@ -155,13 +155,7 @@ pub(crate) fn ctr_encrypt<C>(encryptor: &mut Ctr128BE<C>, buffer: &mut [u8]) ->
 where
     C: BlockCipherEncrypt + BlockSizeUser<BlockSize = U16>,
 {
-    let (blocks, remaining) = Block::<C>::slice_as_chunks_mut(buffer);
-
-    // Ensure input is block-aligned.
-    if !remaining.is_empty() {
-        return Err(Error::Length);
-    }
-
-    encryptor.apply_keystream_blocks(blocks);
-    Ok(())
+    encryptor
+        .try_apply_keystream(buffer)
+        .map_err(|_| Error::Crypto)
 }

ssh-cipher/src/encryptor.rs

@@ -77,10 +77,6 @@ pub type AesGcmNonce = Array<u8, U12>;
 /// `chacha20-poly1305@openssh.com`.
 pub type Tag = Array<u8, U16>;
 
-/// Counter mode with a 128-bit big endian counter.
-#[cfg(feature = "aes-ctr")]
-type Ctr128BE<Cipher> = ctr::CtrCore<Cipher, ctr::flavors::Ctr128BE>;
-
 /// Cipher algorithms.
 #[derive(Copy, Clone, Debug, Eq, Hash, PartialEq, PartialOrd, Ord)]
 #[non_exhaustive]