harden decoders

Cargo.toml

@@ -187,7 +187,7 @@ url = { version = "2" }
 uuid = { version = "1", default-features = false }
 wasi = { version = "0.14", default-features = false }
 wasi-preview1-component-adapter-provider = { version = "49", default-features = false }
-wasm-tokio = { version = "0.6.1", default-features = false }
+wasm-tokio = { version = "0.7.0", default-features = false }
 wasm-wave = { version = "0.258", default-features = false }
 wasmparser = { version = "0.261", default-features = false }
 wasmtime = { version = "49", default-features = false }

crates/transport/src/frame/codec.rs

@@ -2,9 +2,9 @@ use std::sync::Arc;
 
 use bytes::{Bytes, BytesMut};
 use tracing::{instrument, trace};
-use wasm_tokio::{Leb128DecoderU32, Leb128DecoderU64, Leb128Encoder};
+use wasm_tokio::{DEFAULT_MAX_INITIAL_CAPACITY, Leb128DecoderU32, Leb128DecoderU64, Leb128Encoder};
 
-use super::{Frame, FrameRef};
+use crate::frame::{Frame, FrameRef, MAX_INITIAL_PATH_CAPACITY};
 
 /// [Frame] decoder
 pub struct Decoder {
@@ -31,7 +31,7 @@ impl Decoder {
 
 impl Default for Decoder {
     fn default() -> Self {
-        Self::new(32, u32::MAX.into())
+        Self::new(u32::MAX, u32::MAX.into())
     }
 }
 
@@ -59,27 +59,19 @@ impl tokio_util::codec::Decoder for Decoder {
                     ),
                 ));
             }
-            let n = n
-                .try_into()
+            let n = usize::try_from(n)
                 .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
             self.path_cap = n;
-            Vec::with_capacity(n)
+            Vec::with_capacity(n.min(MAX_INITIAL_PATH_CAPACITY))
         };
-        let n = self.path_cap.saturating_sub(src.len());
-        if n > 0 {
-            src.reserve(n);
-            self.path = Some(path);
-            return Ok(None);
-        }
         while self.path_cap > 0 {
             trace!(self.path_cap, "decoding path element");
             let Some(i) = Leb128DecoderU32.decode(src)? else {
                 self.path = Some(path);
                 return Ok(None);
             };
             trace!(i, "decoded path element");
-            let i = i
-                .try_into()
+            let i = usize::try_from(i)
                 .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
             path.push(i);
             self.path_cap -= 1;
@@ -100,8 +92,7 @@ impl tokio_util::codec::Decoder for Decoder {
                     ),
                 ));
             }
-            let n = n
-                .try_into()
+            let n = usize::try_from(n)
                 .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
             self.data_len = n;
             if n == 0 {
@@ -113,7 +104,7 @@ impl tokio_util::codec::Decoder for Decoder {
         }
         let n = self.data_len.saturating_sub(src.len());
         if n > 0 {
-            src.reserve(n);
+            src.reserve(n.min(DEFAULT_MAX_INITIAL_CAPACITY));
             self.path = Some(path);
             return Ok(None);
         }

crates/transport/src/frame/conn/mod.rs

@@ -5,7 +5,7 @@ use core::task::{Context, Poll, ready};
 
 use std::sync::Arc;
 
-use bytes::{Buf as _, BufMut as _, Bytes, BytesMut};
+use bytes::{Buf as _, Bytes, BytesMut};
 use futures::Sink as _;
 use pin_project_lite::pin_project;
 use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt as _};
@@ -16,7 +16,11 @@ use tokio_util::codec::Encoder;
 use tokio_util::io::StreamReader;
 use tokio_util::sync::PollSender;
 use tracing::{Instrument as _, Span, debug, error, instrument, trace};
-use wasm_tokio::{AsyncReadCore as _, AsyncReadLeb128 as _, Leb128Encoder};
+use wasm_tokio::{
+    AsyncReadCore as _, AsyncReadLeb128 as _, DEFAULT_MAX_INITIAL_CAPACITY, Leb128Encoder,
+};
+
+use crate::frame::MAX_INITIAL_PATH_CAPACITY;
 
 mod client;
 mod server;
@@ -581,14 +585,13 @@ async fn ingress(
         let n = AsyncReadExt::chain([b].as_slice(), &mut rx)
             .read_u32_leb128()
             .await?;
-        let n = n
-            .try_into()
+        let n = usize::try_from(n)
             .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
         trace!(n, "read path length");
         let tx = if n == 0 {
             param_tx
         } else {
-            let mut path = Vec::with_capacity(n);
+            let mut path = Vec::with_capacity(n.min(MAX_INITIAL_PATH_CAPACITY));
             for i in 0..n {
                 trace!(i, "reading path element");
                 let p = rx.read_u32_leb128().await?;
@@ -611,18 +614,24 @@ async fn ingress(
         };
         trace!("reading data length");
         let n = rx.read_u32_leb128().await?;
-        let n = n
-            .try_into()
-            .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
         trace!(n, "read data length");
-        let mut buf = BytesMut::with_capacity(n);
-        buf.put_bytes(0, n);
-        trace!("reading data");
-        rx.read_exact(&mut buf).await?;
-        trace!(?buf, "read data");
-        tx.send(Ok(buf.freeze())).await.map_err(|_| {
-            std::io::Error::new(std::io::ErrorKind::BrokenPipe, "stream receiver closed")
-        })?;
+        let mut data = (&mut rx).take(n.into());
+        while data.limit() > 0 {
+            let k = usize::try_from(data.limit())
+                .unwrap_or(usize::MAX)
+                .min(DEFAULT_MAX_INITIAL_CAPACITY);
+            let mut buf = BytesMut::with_capacity(k);
+            trace!(len = k, "reading data chunk");
+            while buf.len() < k {
+                if data.read_buf(&mut buf).await? == 0 {
+                    return Err(std::io::ErrorKind::UnexpectedEof.into());
+                }
+            }
+            trace!(?buf, "read data chunk");
+            tx.send(Ok(buf.freeze())).await.map_err(|_| {
+                std::io::Error::new(std::io::ErrorKind::BrokenPipe, "stream receiver closed")
+            })?;
+        }
     }
 }
 
@@ -676,3 +685,48 @@ pub trait ConnHandler<Rx, Tx> {
 }
 
 impl<Rx, Tx> ConnHandler<Rx, Tx> for () {}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test_log::test(tokio::test)]
+    async fn ingress_truncated_path() {
+        let index = std::sync::Mutex::new(IndexTrie::Empty);
+        let (tx, _rx) = mpsc::channel(1);
+        let err = ingress([0x02, 0x00].as_slice(), &index, &tx)
+            .await
+            .unwrap_err();
+        assert_eq!(err.kind(), std::io::ErrorKind::UnexpectedEof);
+    }
+
+    #[test_log::test(tokio::test)]
+    async fn ingress_chunked_data() {
+        let index = std::sync::Mutex::new(IndexTrie::Empty);
+        let (tx, mut rx) = mpsc::channel(4);
+        let n = DEFAULT_MAX_INITIAL_CAPACITY + 1;
+        let mut frame = vec![0x00, 0x81, 0x80, 0x40];
+        frame.resize(frame.len() + n, 0x42);
+        ingress(frame.as_slice(), &index, &tx).await.unwrap();
+        drop(tx);
+        let mut chunks = vec![];
+        while let Some(chunk) = rx.recv().await {
+            chunks.push(chunk.unwrap().len());
+        }
+        assert_eq!(chunks, [DEFAULT_MAX_INITIAL_CAPACITY, 1]);
+    }
+
+    #[test_log::test(tokio::test)]
+    async fn ingress_truncated_data() {
+        let index = std::sync::Mutex::new(IndexTrie::Empty);
+        let (tx, _rx) = mpsc::channel(1);
+        let err = ingress(
+            [0x00, 0xff, 0xff, 0xff, 0xff, 0x0f, 0x42].as_slice(),
+            &index,
+            &tx,
+        )
+        .await
+        .unwrap_err();
+        assert_eq!(err.kind(), std::io::ErrorKind::UnexpectedEof);
+    }
+}

crates/transport/src/frame/mod.rs

@@ -21,6 +21,8 @@ pub use oneshot::*;
 /// Framing protocol version
 pub const PROTOCOL: u8 = 0;
 
+const MAX_INITIAL_PATH_CAPACITY: usize = 32;
+
 /// Owned wRPC frame
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub struct Frame {

crates/transport/src/value.rs

@@ -26,9 +26,9 @@ use wasm_tokio::cm::{
 };
 use wasm_tokio::{
     CoreNameDecoder, CoreNameEncoder, CoreVecDecoder, CoreVecDecoderBytes, CoreVecEncoderBytes,
-    Leb128DecoderI8, Leb128DecoderI16, Leb128DecoderI32, Leb128DecoderI64, Leb128DecoderI128,
-    Leb128DecoderU8, Leb128DecoderU16, Leb128DecoderU32, Leb128DecoderU64, Leb128DecoderU128,
-    Leb128Encoder, Utf8Codec,
+    DEFAULT_MAX_INITIAL_CAPACITY, Leb128DecoderI8, Leb128DecoderI16, Leb128DecoderI32,
+    Leb128DecoderI64, Leb128DecoderI128, Leb128DecoderU8, Leb128DecoderU16, Leb128DecoderU32,
+    Leb128DecoderU64, Leb128DecoderU128, Leb128Encoder, Utf8Codec,
 };
 
 use crate::BufferedIncoming;
@@ -801,16 +801,27 @@ where
             if len == 0 {
                 return Ok(Some(Vec::default()));
             }
-            let len = len
-                .try_into()
+            let len = usize::try_from(len)
                 .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
-            self.ret = Vec::with_capacity(len);
-            self.deferred = Vec::with_capacity(len);
+            self.ret = Vec::with_capacity(
+                len.min(DEFAULT_MAX_INITIAL_CAPACITY / mem::size_of::<T::Item>().max(1)),
+            );
+            self.deferred = Vec::with_capacity(len.min(
+                DEFAULT_MAX_INITIAL_CAPACITY
+                    / mem::size_of::<Option<DeferredFn<BufferedIncoming>>>(),
+            ));
             self.cap = len;
         }
         while self.cap > 0 {
-            let Some(v) = self.dec.decode(src)? else {
-                return Ok(None);
+            let v = match self.dec.decode(src) {
+                Ok(Some(v)) => v,
+                Ok(None) => return Ok(None),
+                Err(err) => {
+                    self.cap = 0;
+                    self.ret.clear();
+                    self.deferred.clear();
+                    return Err(err);
+                }
             };
             self.ret.push(v);
             self.deferred.push(self.dec.take_deferred());
@@ -1804,8 +1815,7 @@ where
                             for (i, item) in zip(i.., chunk) {
                                 enc.encode(item, &mut buf)?;
                                 if let Some(f) = enc.take_deferred() {
-                                    let i = i
-                                        .try_into()
+                                    let i = usize::try_from(i)
                                         .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
                                     let w = w.index(&[i]).map_err(std::io::Error::other)?;
                                     trace!("spawning transmit task");
@@ -2265,9 +2275,37 @@ mod tests {
 
     use super::*;
 
+    #[test_log::test(tokio::test)]
+    async fn list_decoder_reset_on_error() -> anyhow::Result<()> {
+        let mut dec = <Vec<Vec<String>> as Decode>::Decoder::default();
+        let mut buf = BytesMut::from(&[0x02, 0x01, 0x01, b'a', 0x01, 0x01, 0xff][..]);
+        assert!(tokio_util::codec::Decoder::decode(&mut dec, &mut buf).is_err());
+        let mut buf = BytesMut::from(&[0x01, 0x01, 0x01, b'b'][..]);
+        assert_eq!(
+            tokio_util::codec::Decoder::decode(&mut dec, &mut buf)?,
+            Some(vec![vec!["b".to_string()]])
+        );
+        Ok(())
+    }
+
+    #[test_log::test(tokio::test)]
+    async fn list_decoder_reset_deferred_on_error() -> anyhow::Result<()> {
+        let mut dec =
+            <Vec<Pin<Box<dyn Stream<Item = Bytes> + Send>>> as Decode>::Decoder::default();
+        let mut buf = BytesMut::from(&[0x02, 0x00, 0xff, 0xff, 0xff, 0xff, 0xff][..]);
+        assert!(tokio_util::codec::Decoder::decode(&mut dec, &mut buf).is_err());
+        let mut buf = BytesMut::from(&[0x00][..]);
+        let vs = tokio_util::codec::Decoder::decode(&mut dec, &mut buf)?;
+        assert!(vs.is_some_and(|vs| vs.is_empty()));
+        if let Some(_f) = Deferred::<BufferedIncoming>::take_deferred(&mut dec) {
+            bail!("no deferred read should have been returned");
+        }
+        Ok(())
+    }
+
     #[test_log::test(tokio::test)]
     async fn codec() -> anyhow::Result<()> {
-        let mut buf = BytesMut::new();
+        let mut buf = BytesMut::default();
         let mut enc = <(u8, u32) as Encode>::Encoder::default();
         enc.encode((0x42, 0x42), &mut buf)?;
         if let Some(_f) = Deferred::<Outgoing>::take_deferred(&mut enc) {
@@ -2282,17 +2320,17 @@ mod tests {
         let mut enc = <f32 as Encode>::Encoder::default();
 
         // A non-canonical (e.g. signalling) `NaN` is canonicalized on encode.
-        let mut buf = BytesMut::new();
+        let mut buf = BytesMut::default();
         enc.encode(f32::from_bits(0x7f80_0001), &mut buf).unwrap();
         assert_eq!(buf.as_ref(), CANONICAL_NAN_F32.to_le_bytes());
 
         // A negative `NaN` is canonicalized to the (positive) canonical `NaN`.
-        let mut buf = BytesMut::new();
+        let mut buf = BytesMut::default();
         enc.encode(f32::from_bits(0xffc0_0000), &mut buf).unwrap();
         assert_eq!(buf.as_ref(), CANONICAL_NAN_F32.to_le_bytes());
 
         // Non-`NaN` values are encoded unchanged.
-        let mut buf = BytesMut::new();
+        let mut buf = BytesMut::default();
         enc.encode(1.5_f32, &mut buf).unwrap();
         assert_eq!(buf.as_ref(), 1.5_f32.to_bits().to_le_bytes());
     }
@@ -2301,17 +2339,17 @@ mod tests {
     fn canonical_nan_f64() {
         let mut enc = <f64 as Encode>::Encoder::default();
 
-        let mut buf = BytesMut::new();
+        let mut buf = BytesMut::default();
         enc.encode(f64::from_bits(0x7ff0_0000_0000_0001), &mut buf)
             .unwrap();
         assert_eq!(buf.as_ref(), CANONICAL_NAN_F64.to_le_bytes());
 
-        let mut buf = BytesMut::new();
+        let mut buf = BytesMut::default();
         enc.encode(f64::from_bits(0xfff8_0000_0000_0000), &mut buf)
             .unwrap();
         assert_eq!(buf.as_ref(), CANONICAL_NAN_F64.to_le_bytes());
 
-        let mut buf = BytesMut::new();
+        let mut buf = BytesMut::default();
         enc.encode(1.5_f64, &mut buf).unwrap();
         assert_eq!(buf.as_ref(), 1.5_f64.to_bits().to_le_bytes());
     }