cu_aligner: fix a panic on payload-less ticks and bound snapshot decoding

components/tasks/cu_aligner/src/buffers.rs

@@ -1,4 +1,5 @@
 use circular_buffer::FixedCircularBuffer;
+use cu29::bincode::de::read::Reader;
 use cu29::bincode::de::{Decode, Decoder};
 use cu29::bincode::enc::{Encode, Encoder};
 use cu29::bincode::error::{DecodeError, EncodeError};
@@ -30,6 +31,37 @@ fn extract_tov_time_right(tov: &Tov) -> Option<CuTime> {
     }
 }
 
+/// Largest snapshot accepted for a single buffered message.
+///
+/// The runtime decodes keyframes with `bincode`'s `NoLimit` configuration, so
+/// without a bound here a corrupted snapshot can declare any length it likes and
+/// the process dies allocating it. 256 MiB is far above any realistic single
+/// message (the biggest in-tree user is `cu_image_aligner`, which buffers whole
+/// camera frames) and far below the point where the allocation itself is the
+/// problem.
+const MAX_MSG_SNAPSHOT_BYTES: usize = 256 * 1024 * 1024;
+
+/// Budget handed to the inner decode of a message snapshot.
+///
+/// Careful: `bincode`'s limit counts *claimed* bytes, not wire bytes. Decoding a
+/// container claims `len * size_of::<T>()` (see `Decoder::claim_container_read`),
+/// so a `Vec<u64>` of small varints claims about eight times what it occupies on
+/// the wire. The budget is therefore a multiple of the wire cap, sized for the
+/// widest primitives in practice.
+///
+/// This is deliberately asymmetric with [`MAX_MSG_SNAPSHOT_BYTES`]: `freeze`
+/// bounds wire bytes, `thaw` bounds claimed bytes, and no encoder-side limit
+/// exists in `bincode` to make the two agree. A payload holding an enormous
+/// collection of multi-byte elements could still be written and then refused on
+/// the way back in. Both caps sit far above any realistic buffered message, so
+/// the gap is a documented corner rather than a live concern.
+const MAX_MSG_CLAIM_BYTES: usize = MAX_MSG_SNAPSHOT_BYTES * 8;
+
+/// How much of a message snapshot is read before checking there is really more to
+/// come. A large but honest payload costs a few extra reads; a bogus declared
+/// length costs one buffer of this size and nothing more.
+const SNAPSHOT_READ_CHUNK: usize = 8 * 1024;
+
 fn encode_buffered_msg<P, E>(
     msg: &CuStampedData<P, CuMsgMetadata>,
     encoder: &mut E,
@@ -39,6 +71,13 @@ where
     E: Encoder,
 {
     let bytes = cu29::bincode::encode_to_vec(msg, cu29::bincode::config::standard())?;
+    // Bound what we write, so a snapshot cannot be larger than what
+    // `decode_buffered_msg` is willing to read back by length.
+    if bytes.len() > MAX_MSG_SNAPSHOT_BYTES {
+        return Err(EncodeError::Other(
+            "alignment buffer message is too large to snapshot",
+        ));
+    }
     Encode::encode(&bytes, encoder)
 }
 
@@ -49,9 +88,37 @@ where
     P: CuMsgPayload,
     D: Decoder,
 {
-    let bytes: Vec<u8> = Decode::decode(decoder)?;
+    // Same wire format as `Vec<u8>`: a u64 length prefix followed by the bytes.
+    // Read it by hand rather than through `Vec::<u8>::decode`, which allocates the
+    // whole declared length up front, before any bound is checked.
+    let declared_len: u64 = Decode::decode(decoder)?;
+    let declared_len =
+        usize::try_from(declared_len).map_err(|_| DecodeError::OutsideUsizeRange(declared_len))?;
+    if declared_len > MAX_MSG_SNAPSHOT_BYTES {
+        return Err(DecodeError::LimitExceeded);
+    }
+
+    // Read in chunks so a length that the stream cannot actually satisfy fails on
+    // end-of-input having allocated only what really arrived.
+    let mut bytes = Vec::new();
+    let mut remaining = declared_len;
+    let mut chunk = [0u8; SNAPSHOT_READ_CHUNK];
+    while remaining > 0 {
+        let take = remaining.min(SNAPSHOT_READ_CHUNK);
+        decoder.claim_bytes_read(take)?;
+        decoder.reader().read(&mut chunk[..take])?;
+        bytes.extend_from_slice(&chunk[..take]);
+        remaining -= take;
+    }
+
+    // The inner decode needs the same treatment: a field inside the snapshot can
+    // declare its own length, so it gets a limit rather than `NoLimit`. See
+    // `MAX_MSG_CLAIM_BYTES` for why this budget is not the wire cap above.
     let (msg, bytes_read): (CuStampedData<P, CuMsgMetadata>, usize) =
-        cu29::bincode::decode_from_slice(&bytes, cu29::bincode::config::standard())?;
+        cu29::bincode::decode_from_slice(
+            &bytes,
+            cu29::bincode::config::standard().with_limit::<MAX_MSG_CLAIM_BYTES>(),
+        )?;
     if bytes_read != bytes.len() {
         return Err(DecodeError::OtherString(
             "alignment buffer message snapshot had trailing bytes".to_string(),
@@ -219,10 +286,85 @@ pub use alignment_buffers;
 
 #[cfg(test)]
 mod tests {
+    use super::*;
     use cu29::clock::Tov;
-    use cu29::cutask::*;
     use std::time::Duration;
 
+    type TestBuffer = TimeboundCircularBuffer<4, u32, CuMsgMetadata>;
+
+    /// Drives `thaw` through the normal bincode entry point.
+    struct Thawed(TestBuffer);
+
+    impl Decode<()> for Thawed {
+        fn decode<D: Decoder<Context = ()>>(decoder: &mut D) -> Result<Self, DecodeError> {
+            let mut buffer = TestBuffer::new();
+            buffer.thaw(decoder)?;
+            Ok(Thawed(buffer))
+        }
+    }
+
+    /// Drives `freeze` through the normal bincode entry point.
+    struct Frozen<'a>(&'a TestBuffer);
+
+    impl Encode for Frozen<'_> {
+        fn encode<E: Encoder>(&self, encoder: &mut E) -> Result<(), EncodeError> {
+            self.0.freeze(encoder)
+        }
+    }
+
+    /// A corrupted snapshot must be rejected, not turned into a multi-gigabyte
+    /// allocation. The runtime decodes with `NoLimit`, so the bound has to live here.
+    #[test]
+    fn thaw_rejects_a_bogus_message_length() {
+        let config = cu29::bincode::config::standard();
+
+        // A length past the per-message cap is rejected before anything is read.
+        let bytes = cu29::bincode::encode_to_vec((1u64, 5_000_000_000u64), config).unwrap();
+        let Err(err) = cu29::bincode::decode_from_slice::<Thawed, _>(&bytes, config) else {
+            panic!("a snapshot claiming 5 GB must be rejected");
+        };
+        assert!(
+            matches!(err, DecodeError::LimitExceeded),
+            "expected the per-message cap to reject it, got {err:?}"
+        );
+
+        // A length under the cap that the stream cannot satisfy fails on
+        // end-of-input, having allocated only what actually arrived. The old code
+        // reported the same error kind here, so this half guards the chunked-read
+        // path rather than the cap: what changed is that 64 MB is no longer
+        // allocated up front before the failure.
+        let bytes = cu29::bincode::encode_to_vec((1u64, 64_000_000u64), config).unwrap();
+        let Err(err) = cu29::bincode::decode_from_slice::<Thawed, _>(&bytes, config) else {
+            panic!("a snapshot promising 64 MB of absent bytes must be rejected");
+        };
+        assert!(
+            matches!(err, DecodeError::UnexpectedEnd { .. }),
+            "expected an end-of-input error, got {err:?}"
+        );
+    }
+
+    /// The chunked read must still round-trip an honest snapshot byte for byte.
+    #[test]
+    fn freeze_thaw_round_trips() {
+        let mut buffer = TestBuffer::new();
+        for (i, payload) in [11u32, 22, 33].into_iter().enumerate() {
+            let mut msg = CuStampedData::<u32, CuMsgMetadata>::new(Some(payload));
+            msg.tov = Tov::Time(Duration::from_secs(i as u64 + 1).into());
+            buffer.push(msg);
+        }
+
+        let config = cu29::bincode::config::standard();
+        let bytes = cu29::bincode::encode_to_vec(Frozen(&buffer), config).unwrap();
+        let (Thawed(restored), _) =
+            cu29::bincode::decode_from_slice::<Thawed, _>(&bytes, config).unwrap();
+
+        assert_eq!(restored.inner.len(), buffer.inner.len());
+        for (a, b) in restored.inner.iter().zip(buffer.inner.iter()) {
+            assert_eq!(a.payload(), b.payload());
+            assert_eq!(a.tov, b.tov);
+        }
+    }
+
     #[test]
     fn simple_init_test() {
         alignment_buffers!(AlignmentBuffers, buffer1: TimeboundCircularBuffer<10, CuStampedData<u32, CuMsgMetadata>>, buffer2: TimeboundCircularBuffer<12, CuStampedData<u64, CuMsgMetadata>>);

components/tasks/cu_aligner/src/lib.rs

@@ -79,11 +79,16 @@ macro_rules! define_task {
                 input: &Self::Input<'_>,
                 output: &mut Self::Output<'_>,
             ) -> CuResult<()> {
-                // add the incoming data into the buffers
-                // input is a tuple of &CuMsg<T> for each T in the input
+                // Add the incoming data into the buffers.
+                // input is a tuple of &CuMsg<T> for each T in the input.
+                // A tick where the upstream task had nothing to emit carries a tov
+                // but no payload. It holds no data to align, so it must not take a
+                // slot in the fixed-size buffer nor advance the alignment window.
                 paste::paste! {
                     $(
-                        self.aligner.[<buffer $index>].push(input.$index.clone());
+                        if input.$index.payload().is_some() {
+                            self.aligner.[<buffer $index>].push(input.$index.clone());
+                        }
                     )*
                 }
 
@@ -93,10 +98,12 @@ macro_rules! define_task {
                     return Ok(());
                 };
 
-                // Populate the CuArray fields in the output message
+                // Populate the CuArray fields in the output message.
+                // `TimeboundCircularBuffer::push` is public, so skip payload-less
+                // messages here too rather than unwrapping them.
                 let output_payload = output.payload_mut().get_or_insert_with(Default::default);
                 $(
-                    output_payload.$index.fill_from_iter(tuple_of_iters.$index.map(|msg| msg.payload().unwrap().clone()));
+                    output_payload.$index.fill_from_iter(tuple_of_iters.$index.filter_map(|msg| msg.payload().cloned()));
                 )*
                 Ok(())
             }
@@ -125,6 +132,56 @@ mod tests {
         let result = aligner.process(&ctx, &input, &mut output);
         assert!(result.is_ok());
     }
+    /// A task that had nothing to emit still ticks, so a message can carry a tov
+    /// with no payload. Such a message used to reach `payload().unwrap()` and abort
+    /// the process.
+    #[test]
+    fn test_aligner_tolerates_payload_less_ticks() {
+        let mut config = ComponentConfig::default();
+        config.set("target_alignment_window_ms", 100);
+        config.set("stale_data_horizon_ms", 1000);
+        let mut aligner = AlignerTask::new(Some(&config), ()).unwrap();
+        let ctx = CuContext::new_with_clock();
+
+        let tov = Tov::Time(CuTime::from_millis(100));
+        let mut empty = CuStampedData::<f32, CuMsgMetadata>::new(None);
+        empty.tov = tov;
+        let mut present = CuStampedData::<i32, CuMsgMetadata>::new(Some(7));
+        present.tov = tov;
+
+        let mut output =
+            CuStampedData::<(CuArray<f32, 5>, CuArray<i32, 10>), CuMsgMetadata>::default();
+        aligner
+            .process(&ctx, &(&empty, &present), &mut output)
+            .unwrap();
+
+        // The payload-less tick contributes nothing and takes no buffer slot, so
+        // that stream's array comes back empty. An empty array is already a normal
+        // outcome here: `iter_window` selects on time, so a stream whose data all
+        // falls outside the window yields nothing even when every message it sent
+        // carried a payload. Consumers have to handle that either way.
+        let payload = output.payload().unwrap();
+        assert_eq!(payload.0.len(), 0);
+        assert_eq!(payload.1.as_slice(), &[7]);
+
+        // Once real data arrives on that stream, both align as usual.
+        let tov = Tov::Time(CuTime::from_millis(150));
+        let mut left = CuStampedData::<f32, CuMsgMetadata>::new(Some(1.5));
+        left.tov = tov;
+        let mut right = CuStampedData::<i32, CuMsgMetadata>::new(Some(9));
+        right.tov = tov;
+
+        let mut output =
+            CuStampedData::<(CuArray<f32, 5>, CuArray<i32, 10>), CuMsgMetadata>::default();
+        aligner
+            .process(&ctx, &(&left, &right), &mut output)
+            .unwrap();
+
+        let payload = output.payload().unwrap();
+        assert_eq!(payload.0.as_slice(), &[1.5]);
+        assert_eq!(payload.1.as_slice(), &[7, 9]);
+    }
+
     mod string_payload {
         use super::*;