fix(rust): bound decoder allocations and pre-size encoded frames

rust/loro-protocol/examples/encoding_bench.rs

@@ -0,0 +1,31 @@
+use loro_protocol::{encode, BatchId, CrdtType, ProtocolMessage};
+use std::{hint::black_box, time::Instant};
+
+fn main() {
+    for &(count, bytes) in &[(1, 32), (1, 4096), (64, 1024), (1, 240 * 1024)] {
+        let message = ProtocolMessage::DocUpdate {
+            crdt: CrdtType::Loro,
+            room_id: "benchmark-room".into(),
+            batch_id: BatchId([7; 8]),
+            updates: vec![vec![42; bytes]; count],
+        };
+        let expected = encode(&message).unwrap();
+        assert_eq!(loro_protocol::decode(&expected).unwrap(), message);
+        let iterations = if bytes > 100_000 { 2000 } else { 10_000 };
+        let mut samples = Vec::new();
+        for _ in 0..15 {
+            let start = Instant::now();
+            for _ in 0..iterations {
+                black_box(encode(black_box(&message)).unwrap());
+            }
+            samples.push(start.elapsed().as_nanos() / iterations);
+        }
+        samples.sort_unstable();
+        println!(
+            "count={count} payload={bytes} encoded={} median_ns={} p95_ns={}",
+            expected.len(),
+            samples[7],
+            samples[14]
+        );
+    }
+}

rust/loro-protocol/src/bytes.rs

@@ -9,8 +9,14 @@ pub struct BytesWriter {
 impl BytesWriter {
     #[must_use]
     pub fn new() -> Self {
+        Self::with_capacity(32)
+    }
+
+    /// Create a writer with space for a known encoded message size.
+    #[must_use]
+    pub fn with_capacity(capacity: usize) -> Self {
         Self {
-            buf: Vec::with_capacity(32),
+            buf: Vec::with_capacity(capacity),
         }
     }
 
@@ -119,6 +125,9 @@ impl<'a> BytesReader<'a> {
         let mut shift: u32 = 0;
         loop {
             let byte = self.read_byte()?;
+            if shift == 63 && byte & 0x7f > 1 {
+                return Err("uleb128 too large".into());
+            }
             result |= u64::from(byte & 0x7f) << shift;
             if (byte & 0x80) == 0 {
                 break;

rust/loro-protocol/src/elo.rs

@@ -49,9 +49,12 @@ pub struct ParsedEloRecord<'a> {
 
 /// Decode an ELO container into a list of record byte slices.
 /// The returned slices borrow from `data`.
-pub fn decode_elo_container<'a>(data: &'a [u8]) -> Result<Vec<&'a [u8]>, String> {
+pub fn decode_elo_container(data: &[u8]) -> Result<Vec<&[u8]>, String> {
     let mut r = BytesReader::new(data);
     let n = usize::try_from(r.read_uleb128()?).map_err(|_| "length too large".to_string())?;
+    if n > r.remaining() {
+        return Err("ELO record count exceeds remaining payload".into());
+    }
     let mut out: Vec<&[u8]> = Vec::with_capacity(n);
     for _ in 0..n {
         out.push(r.read_var_bytes()?);
@@ -99,7 +102,7 @@ fn parse_delta<'a>(
     if peer_id.len() > 64 {
         return Err("Invalid ELO delta span: peerId too long".into());
     }
-    if key_id.as_bytes().len() > 64 {
+    if key_id.len() > 64 {
         return Err("Invalid ELO delta span: keyId too long".into());
     }
 
@@ -126,6 +129,10 @@ fn parse_snapshot<'a>(
 ) -> Result<ParsedEloRecord<'a>, String> {
     let count =
         usize::try_from(r.read_uleb128()?).map_err(|_| "vv length too large".to_string())?;
+    // At least a peer-ID length byte and a counter byte are needed per entry.
+    if count > r.remaining() / 2 {
+        return Err("ELO version vector count exceeds remaining payload".into());
+    }
     let mut vv: Vec<(Vec<u8>, u64)> = Vec::with_capacity(count);
     for _ in 0..count {
         let pid = r.read_var_bytes()?.to_vec();
@@ -143,7 +150,7 @@ fn parse_snapshot<'a>(
     if iv_bytes.len() != 12 {
         return Err("Invalid ELO snapshot: IV must be 12 bytes".into());
     }
-    if key_id.as_bytes().len() > 64 {
+    if key_id.len() > 64 {
         return Err("Invalid ELO snapshot: keyId too long".into());
     }
     // vv should be sorted by peer_id bytes ascending

rust/loro-protocol/src/encoding.rs

@@ -7,6 +7,75 @@ use crate::protocol::*;
 
 const MAX_ROOM_ID_LENGTH: usize = 128;
 
+fn uleb128_len(n: u64) -> usize {
+    (64 - n.leading_zeros()).max(1).div_ceil(7) as usize
+}
+
+fn var_bytes_len(len: usize) -> usize {
+    len.saturating_add(uleb128_len(len as u64))
+}
+
+/// Size only the fields that the encoder actually writes. Optional JoinError
+/// fields depend on the error code, not merely on whether they are present.
+fn payload_len(message: &ProtocolMessage) -> usize {
+    match message {
+        ProtocolMessage::JoinRequest { auth, version, .. } => {
+            var_bytes_len(auth.len()).saturating_add(var_bytes_len(version.len()))
+        }
+        ProtocolMessage::JoinResponseOk {
+            permission,
+            version,
+            extra,
+            ..
+        } => {
+            let permission_len = match permission {
+                Permission::Read => 4,
+                Permission::Write => 5,
+            };
+            var_bytes_len(permission_len)
+                .saturating_add(var_bytes_len(version.len()))
+                .saturating_add(var_bytes_len(extra.as_ref().map_or(0, Vec::len)))
+        }
+        ProtocolMessage::JoinError {
+            code,
+            message,
+            receiver_version,
+            app_code,
+            ..
+        } => {
+            let mut len = 1usize.saturating_add(var_bytes_len(message.len()));
+            if matches!(code, JoinErrorCode::VersionUnknown) {
+                if let Some(version) = receiver_version {
+                    len = len.saturating_add(var_bytes_len(version.len()));
+                }
+            }
+            if matches!(code, JoinErrorCode::AppError) {
+                if let Some(app) = app_code {
+                    len = len.saturating_add(var_bytes_len(app.len()));
+                }
+            }
+            len
+        }
+        ProtocolMessage::DocUpdate { updates, .. } => updates.iter().fold(
+            uleb128_len(updates.len() as u64).saturating_add(8),
+            |len, update| len.saturating_add(var_bytes_len(update.len())),
+        ),
+        ProtocolMessage::DocUpdateFragmentHeader {
+            fragment_count,
+            total_size_bytes,
+            ..
+        } => 8 + uleb128_len(*fragment_count) + uleb128_len(*total_size_bytes),
+        ProtocolMessage::DocUpdateFragment {
+            index, fragment, ..
+        } => (8 + uleb128_len(*index)).saturating_add(var_bytes_len(fragment.len())),
+        ProtocolMessage::RoomError { message, .. } => {
+            1usize.saturating_add(var_bytes_len(message.len()))
+        }
+        ProtocolMessage::Ack { .. } => 9,
+        ProtocolMessage::Leave { .. } => 0,
+    }
+}
+
 fn encode_crdt(w: &mut BytesWriter, crdt: CrdtType) {
     w.push_bytes(&crdt.magic_bytes());
 }
@@ -48,10 +117,8 @@ fn decode_permission(r: &mut BytesReader) -> Result<Permission, String> {
 ///
 /// Returns an owned `Vec<u8>` containing the encoded message.
 pub fn encode(message: &ProtocolMessage) -> Result<Vec<u8>, String> {
-    let mut w = BytesWriter::new();
-
     // Common: crdt, room_id, type
-    match message {
+    let (crdt, room_id) = match message {
         ProtocolMessage::JoinRequest {
             crdt,
             room_id,
@@ -100,15 +167,21 @@ pub fn encode(message: &ProtocolMessage) -> Result<Vec<u8>, String> {
             ref_id: _,
             status: _,
         }
-        | ProtocolMessage::Leave { crdt, room_id } => {
-            let room_id_bytes = room_id.as_bytes();
-            if room_id_bytes.len() > MAX_ROOM_ID_LENGTH {
-                return Err("Room ID too long".into());
-            }
-            encode_crdt(&mut w, *crdt);
-            w.push_var_bytes(room_id_bytes);
-        }
+        | ProtocolMessage::Leave { crdt, room_id } => (*crdt, room_id),
+    };
+    if room_id.len() > MAX_ROOM_ID_LENGTH {
+        return Err("Room ID too long".into());
+    }
+    let size = (5 + var_bytes_len(room_id.len())).saturating_add(payload_len(message));
+    if size > MAX_MESSAGE_SIZE {
+        return Err(format!(
+            "Message size {} exceeds maximum {}",
+            size, MAX_MESSAGE_SIZE
+        ));
     }
+    let mut w = BytesWriter::with_capacity(size);
+    encode_crdt(&mut w, crdt);
+    w.push_var_bytes(room_id.as_bytes());
 
     // Type byte
     let ty: u8 = match message {
@@ -208,6 +281,7 @@ pub fn encode(message: &ProtocolMessage) -> Result<Vec<u8>, String> {
     }
 
     let out = w.finalize();
+    debug_assert_eq!(out.len(), size);
     if out.len() > MAX_MESSAGE_SIZE {
         return Err(format!(
             "Message size {} exceeds maximum {}",
@@ -286,7 +360,13 @@ pub fn decode(buf: &[u8]) -> Result<ProtocolMessage, String> {
             }
         }
         MessageType::DocUpdate => {
-            let count = r.read_uleb128()? as usize;
+            let count = usize::try_from(r.read_uleb128()?)
+                .map_err(|_| "Invalid DocUpdate: update count too large".to_string())?;
+            // Every update needs at least its one-byte length prefix, followed
+            // by the batch ID. Bound attacker-controlled capacity before allocation.
+            if count > r.remaining().saturating_sub(8) {
+                return Err("Invalid DocUpdate: update count exceeds remaining payload".into());
+            }
             let mut updates = Vec::with_capacity(count);
             for _ in 0..count {
                 updates.push(r.read_var_bytes()?.to_vec());

rust/loro-protocol/tests/decode_bounds.rs

@@ -0,0 +1,44 @@
+use loro_protocol::{
+    decode, decode_elo_container, parse_elo_record_header, BytesReader, BytesWriter, EloRecordKind,
+};
+
+#[test]
+fn overflowing_varint_is_not_silently_truncated() {
+    // The tenth byte can contain only bit 63. A payload of 2 overflows u64.
+    let mut bytes = vec![0x80; 9];
+    bytes.push(2);
+    assert!(BytesReader::new(&bytes).read_uleb128().is_err());
+    let mut writer = BytesWriter::new();
+    writer.push_uleb128(u64::MAX);
+    assert_eq!(
+        BytesReader::new(&writer.finalize()).read_uleb128().unwrap(),
+        u64::MAX
+    );
+}
+
+#[test]
+fn impossible_update_count_returns_error_without_panicking() {
+    let mut writer = BytesWriter::new();
+    writer.push_bytes(b"%LOR");
+    writer.push_var_string("x");
+    writer.push_byte(3); // DocUpdate
+    writer.push_uleb128(u64::MAX);
+    writer.push_bytes(&[0; 8]);
+    let bytes = writer.finalize();
+    assert!(decode(&bytes).is_err());
+}
+
+#[test]
+fn impossible_encrypted_record_count_returns_error_without_panicking() {
+    let mut writer = BytesWriter::new();
+    writer.push_uleb128(u64::MAX);
+    assert!(decode_elo_container(&writer.finalize()).is_err());
+}
+
+#[test]
+fn impossible_encrypted_snapshot_vector_returns_error_without_panicking() {
+    let mut writer = BytesWriter::new();
+    writer.push_byte(EloRecordKind::Snapshot as u8);
+    writer.push_uleb128(u64::MAX);
+    assert!(parse_elo_record_header(&writer.finalize()).is_err());
+}

rust/loro-protocol/tests/elo.rs

@@ -73,7 +73,7 @@ fn decode_container_and_parse_delta() {
         _ => panic!("expected delta header"),
     }
     assert_eq!(parsed.ct, &ct[..]);
-    assert!(parsed.aad.len() > 0);
+    assert!(!parsed.aad.is_empty());
 }
 
 #[test]

rust/loro-protocol/tests/encoding_limits.rs

@@ -0,0 +1,48 @@
+use loro_protocol::{decode, encode, BatchId, CrdtType, ProtocolMessage, MAX_MESSAGE_SIZE};
+
+fn update(payload: usize, room: &str) -> ProtocolMessage {
+    ProtocolMessage::DocUpdate {
+        crdt: CrdtType::Loro,
+        room_id: room.into(),
+        updates: vec![vec![17; payload]],
+        batch_id: BatchId([0; 8]),
+    }
+}
+
+#[test]
+fn encoding_preserves_exact_limit_and_room_error_precedence() {
+    // 4 magic + 2 room + 1 type + 1 count + 3 payload length + 8 batch ID.
+    let limit = MAX_MESSAGE_SIZE - 19;
+    let message = update(limit, "x");
+    let bytes = encode(&message).unwrap();
+    assert_eq!(bytes.len(), MAX_MESSAGE_SIZE);
+    assert_eq!(decode(&bytes).unwrap(), message);
+    assert_eq!(
+        encode(&update(limit + 1, "x")).unwrap_err(),
+        format!(
+            "Message size {} exceeds maximum {}",
+            MAX_MESSAGE_SIZE + 1,
+            MAX_MESSAGE_SIZE
+        )
+    );
+    assert_eq!(
+        encode(&update(MAX_MESSAGE_SIZE, &"x".repeat(129))).unwrap_err(),
+        "Room ID too long"
+    );
+}
+
+#[test]
+fn varint_boundaries_and_many_empty_updates_keep_wire_parity() {
+    for size in [0, 1, 127, 128, 16383, 16384] {
+        let message = update(size, "世界");
+        let bytes = encode(&message).unwrap();
+        assert_eq!(decode(&bytes).unwrap(), message);
+    }
+    let message = ProtocolMessage::DocUpdate {
+        crdt: CrdtType::Loro,
+        room_id: "x".into(),
+        updates: vec![vec![]; MAX_MESSAGE_SIZE],
+        batch_id: BatchId([0; 8]),
+    };
+    assert!(encode(&message).unwrap_err().starts_with("Message size "));
+}