Merge remote-tracking branch 'origin/main' into pr/702

Author: Eugeny

cryptovec/Cargo.toml

@@ -6,7 +6,7 @@ edition = "2024"
 license = "Apache-2.0"
 name = "russh-cryptovec"
 repository = "https://github.com/warp-tech/russh"
-version = "0.59.0"
+version = "0.60.3"
 rust-version = "1.85"
 
 [dependencies]

cryptovec/src/cryptovec.rs

@@ -153,6 +153,57 @@ impl Default for CryptoVec {
     }
 }
 
+const MAX_CAPACITY: usize = 1usize << (usize::BITS - 2);
+
+#[cold]
+#[inline(never)]
+#[allow(clippy::panic)]
+fn capacity_overflow(len: usize) -> ! {
+    panic!("CryptoVec capacity overflow: {len}")
+}
+
+#[cold]
+#[inline(never)]
+#[allow(clippy::panic)]
+fn length_overflow(lhs: usize, rhs: usize) -> ! {
+    panic!("CryptoVec length overflow: {lhs} + {rhs}")
+}
+
+#[cold]
+#[inline(never)]
+fn alloc_failed(layout: std::alloc::Layout) -> ! {
+    std::alloc::handle_alloc_error(layout)
+}
+
+#[inline]
+fn checked_capacity(len: usize) -> usize {
+    if len > MAX_CAPACITY {
+        capacity_overflow(len);
+    }
+    len.next_power_of_two()
+}
+
+#[inline]
+unsafe fn alloc_zeroed(capacity: usize) -> *mut u8 {
+    debug_assert!(capacity > 0);
+    let layout = unsafe { std::alloc::Layout::from_size_align_unchecked(capacity, 1) };
+    let p = unsafe { std::alloc::alloc_zeroed(layout) };
+    if p.is_null() {
+        alloc_failed(layout);
+    }
+    let _ = mlock(p, capacity);
+    p
+}
+
+#[inline]
+fn checked_len_sum(lhs: usize, rhs: usize) -> usize {
+    let sum = lhs.wrapping_add(rhs);
+    if sum < lhs {
+        length_overflow(lhs, rhs);
+    }
+    sum
+}
+
 impl CryptoVec {
     /// Creates a new `CryptoVec`.
     pub fn new() -> CryptoVec {
@@ -161,27 +212,27 @@ impl CryptoVec {
 
     /// Creates a new `CryptoVec` with `n` zeros.
     pub fn new_zeroed(size: usize) -> CryptoVec {
-        unsafe {
-            let capacity = size.next_power_of_two();
-            let layout = std::alloc::Layout::from_size_align_unchecked(capacity, 1);
-            let p = std::alloc::alloc_zeroed(layout);
-            let _ = mlock(p, capacity);
-            CryptoVec { p, capacity, size }
+        if size == 0 {
+            return CryptoVec::default();
         }
+
+        let capacity = checked_capacity(size);
+        let p = unsafe { alloc_zeroed(capacity) };
+        CryptoVec { p, capacity, size }
     }
 
     /// Creates a new `CryptoVec` with capacity `capacity`.
     pub fn with_capacity(capacity: usize) -> CryptoVec {
-        unsafe {
-            let capacity = capacity.next_power_of_two();
-            let layout = std::alloc::Layout::from_size_align_unchecked(capacity, 1);
-            let p = std::alloc::alloc_zeroed(layout);
-            let _ = mlock(p, capacity);
-            CryptoVec {
-                p,
-                capacity,
-                size: 0,
-            }
+        if capacity == 0 {
+            return CryptoVec::default();
+        }
+
+        let capacity = checked_capacity(capacity);
+        let p = unsafe { alloc_zeroed(capacity) };
+        CryptoVec {
+            p,
+            capacity,
+            size: 0,
         }
     }
 
@@ -220,29 +271,21 @@ impl CryptoVec {
         } else {
             // realloc ! and erase the previous memory.
             unsafe {
-                let next_capacity = size.next_power_of_two();
+                let next_capacity = checked_capacity(size);
                 let old_ptr = self.p;
-                let next_layout = std::alloc::Layout::from_size_align_unchecked(next_capacity, 1);
-                self.p = std::alloc::alloc_zeroed(next_layout);
-                let _ = mlock(self.p, next_capacity);
+                let next_ptr = alloc_zeroed(next_capacity);
 
                 if self.capacity > 0 {
-                    std::ptr::copy_nonoverlapping(old_ptr, self.p, self.size);
+                    std::ptr::copy_nonoverlapping(old_ptr, next_ptr, self.size);
                     zeroize(old_ptr, self.size);
                     let _ = munlock(old_ptr, self.capacity);
                     let layout = std::alloc::Layout::from_size_align_unchecked(self.capacity, 1);
                     std::alloc::dealloc(old_ptr, layout);
                 }
 
-                if self.p.is_null() {
-                    #[allow(clippy::panic)]
-                    {
-                        panic!("Realloc failed, pointer = {self:?} {size:?}")
-                    }
-                } else {
-                    self.capacity = next_capacity;
-                    self.size = size;
-                }
+                self.p = next_ptr;
+                self.capacity = next_capacity;
+                self.size = size;
             }
         }
     }
@@ -262,7 +305,7 @@ impl CryptoVec {
     /// Append a new byte at the end of this CryptoVec.
     pub fn push(&mut self, s: u8) {
         let size = self.size;
-        self.resize(size + 1);
+        self.resize(checked_len_sum(size, 1));
         unsafe { *self.p.add(size) = s }
     }
 
@@ -274,7 +317,8 @@ impl CryptoVec {
         mut r: R,
     ) -> Result<usize, std::io::Error> {
         let cur_size = self.size;
-        self.resize(cur_size + n_bytes);
+        let target_size = checked_len_sum(cur_size, n_bytes);
+        self.resize(target_size);
         let s = unsafe { std::slice::from_raw_parts_mut(self.p.add(cur_size), n_bytes) };
         // Resize the buffer to its appropriate size.
         match r.read(s) {
@@ -319,7 +363,7 @@ impl CryptoVec {
     /// ```
     pub fn resize_mut(&mut self, n: usize) -> &mut [u8] {
         let size = self.size;
-        self.resize(size + n);
+        self.resize(checked_len_sum(size, n));
         unsafe { std::slice::from_raw_parts_mut(self.p.add(size), n) }
     }
 
@@ -331,7 +375,8 @@ impl CryptoVec {
     /// ```
     pub fn extend(&mut self, s: &[u8]) {
         let size = self.size;
-        self.resize(size + s.len());
+        let added = s.len();
+        self.resize(checked_len_sum(size, added));
         unsafe {
             std::ptr::copy_nonoverlapping(s.as_ptr(), self.p.add(size), s.len());
         }
@@ -438,7 +483,7 @@ fn optimization_barrier(dst: *mut u8, size: usize) {
 
 #[cfg(test)]
 mod test {
-    use super::CryptoVec;
+    use super::{CryptoVec, checked_capacity};
 
     #[test]
     fn test_new() {
@@ -569,13 +614,39 @@ mod test {
         assert!(crypto_vec.is_empty());
     }
 
+    #[test]
+    fn test_with_capacity_zero() {
+        let crypto_vec = CryptoVec::with_capacity(0);
+        assert_eq!(crypto_vec.size, 0);
+        assert_eq!(crypto_vec.capacity, 0);
+    }
+
+    #[test]
+    fn test_new_zeroed_zero() {
+        let crypto_vec = CryptoVec::new_zeroed(0);
+        assert_eq!(crypto_vec.size, 0);
+        assert_eq!(crypto_vec.capacity, 0);
+    }
+
     #[test]
     fn test_extend() {
         let mut crypto_vec = CryptoVec::new();
         crypto_vec.extend(b"test");
         assert_eq!(crypto_vec.as_ref(), b"test");
     }
 
+    #[test]
+    #[should_panic(expected = "CryptoVec capacity overflow")]
+    fn test_checked_capacity_overflow_panics() {
+        let _ = checked_capacity(usize::MAX);
+    }
+
+    #[test]
+    #[should_panic(expected = "CryptoVec capacity overflow")]
+    fn test_checked_capacity_rejects_values_above_max_capacity() {
+        let _ = checked_capacity(super::MAX_CAPACITY + 1);
+    }
+
     #[test]
     fn test_write_all_from() {
         let mut crypto_vec = CryptoVec::new();

cryptovec/src/platform/unix.rs

@@ -7,9 +7,14 @@ use super::MemoryLockError;
 
 /// Unlock memory on drop for Unix-based systems.
 pub fn munlock(ptr: *const u8, len: usize) -> Result<(), MemoryLockError> {
+    if len == 0 {
+        return Ok(());
+    }
+    let Some(ptr) = NonNull::new(ptr as *mut c_void) else {
+        return Err(MemoryLockError::new("munlock: null pointer".into()));
+    };
     unsafe {
         Errno::clear();
-        let ptr = NonNull::new_unchecked(ptr as *mut c_void);
         nix::sys::mman::munlock(ptr, len).map_err(|e| {
             MemoryLockError::new(format!("munlock: {} (0x{:x})", e.desc(), e as i32))
         })?;
@@ -18,9 +23,14 @@ pub fn munlock(ptr: *const u8, len: usize) -> Result<(), MemoryLockError> {
 }
 
 pub fn mlock(ptr: *const u8, len: usize) -> Result<(), MemoryLockError> {
+    if len == 0 {
+        return Ok(());
+    }
+    let Some(ptr) = NonNull::new(ptr as *mut c_void) else {
+        return Err(MemoryLockError::new("mlock: null pointer".into()));
+    };
     unsafe {
         Errno::clear();
-        let ptr = NonNull::new_unchecked(ptr as *mut c_void);
         nix::sys::mman::mlock(ptr, len)
             .map_err(|e| MemoryLockError::new(format!("mlock: {} (0x{:x})", e.desc(), e as i32)))?;
     }

russh/Cargo.toml

@@ -9,7 +9,7 @@ license = "Apache-2.0"
 name = "russh"
 readme = "../README.md"
 repository = "https://github.com/warp-tech/russh"
-version = "0.60.2"
+version = "0.60.3"
 rust-version = "1.85"
 
 [features]
@@ -91,7 +91,7 @@ rand_core = { version = "0.10.0" }
 rand.workspace = true
 ring = { version = "0.17.14", optional = true }
 rsa = { version = "0.10.0-rc.18", optional = true }
-russh-cryptovec = { version = "0.59.0", path = "../cryptovec", features = [
+russh-cryptovec = { version = "0.60.3", path = "../cryptovec", features = [
   "ssh-encoding",
 ] }
 russh-util = { version = "0.52.0", path = "../russh-util" }

russh/src/keys/agent/client.rs

@@ -17,6 +17,8 @@ pub trait AgentStream: AsyncRead + AsyncWrite {}
 
 impl<S: AsyncRead + AsyncWrite> AgentStream for S {}
 
+const MAX_AGENT_FRAME_LEN: usize = 256 * 1024;
+
 /// SSH agent client.
 pub struct AgentClient<S: AgentStream> {
     stream: S,
@@ -112,25 +114,29 @@ impl AgentClient<tokio::net::windows::named_pipe::NamedPipeClient> {
 }
 
 impl<S: AgentStream + Unpin> AgentClient<S> {
-    async fn read_response(&mut self) -> Result<(), Error> {
-        // Writing the message
-        self.stream.write_all(&self.buf).await?;
-        self.stream.flush().await?;
-
-        // Reading the length
+    async fn read_frame(&mut self) -> Result<(), Error> {
         self.buf.clear();
         self.buf.resize(4, 0);
         self.stream.read_exact(&mut self.buf).await?;
 
-        // Reading the rest of the buffer
         let len = BigEndian::read_u32(&self.buf) as usize;
+        if len > MAX_AGENT_FRAME_LEN {
+            return Err(Error::AgentProtocolError);
+        }
+
         self.buf.clear();
         self.buf.resize(len, 0);
         self.stream.read_exact(&mut self.buf).await?;
-
         Ok(())
     }
 
+    async fn read_response(&mut self) -> Result<(), Error> {
+        // Writing the message
+        self.stream.write_all(&self.buf).await?;
+        self.stream.flush().await?;
+        self.read_frame().await
+    }
+
     async fn read_success(&mut self) -> Result<(), Error> {
         self.read_response().await?;
         if self.buf.first() == Some(&msg::SUCCESS) {
@@ -570,3 +576,42 @@ impl<S: AgentStream + Unpin> AgentClient<S> {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use byteorder::{BigEndian, ByteOrder};
+    use tokio::io::{AsyncReadExt, AsyncWriteExt};
+
+    use super::{AgentClient, MAX_AGENT_FRAME_LEN};
+    use crate::keys::Error;
+
+    #[test]
+    fn oversized_agent_response_is_rejected_before_allocation() -> std::io::Result<()> {
+        let runtime = tokio::runtime::Builder::new_current_thread()
+            .enable_all()
+            .build()?;
+
+        runtime.block_on(async {
+            let (mut writer, reader) = tokio::io::duplex(64);
+            let server = tokio::spawn(async move {
+                let mut frame = [0u8; 4];
+                writer.read_exact(&mut frame).await?;
+                let len = BigEndian::read_u32(&frame) as usize;
+                let mut body = vec![0; len];
+                writer.read_exact(&mut body).await?;
+
+                BigEndian::write_u32(&mut frame, (MAX_AGENT_FRAME_LEN + 1) as u32);
+                writer.write_all(&frame).await?;
+                Ok::<(), std::io::Error>(())
+            });
+
+            let mut client = AgentClient::connect(reader);
+            let err = client.request_identities().await.unwrap_err();
+            assert!(matches!(err, Error::AgentProtocolError));
+            server.await.expect("server task")?;
+            Ok::<(), std::io::Error>(())
+        })?;
+
+        Ok(())
+    }
+}