From f9efdbaf56625b42bbdeeb86576bf9ce813269b5 Mon Sep 17 00:00:00 2001 From: Petra Date: Sat, 10 Oct 2026 10:50:32 -0400 Subject: [PATCH] fix: bound protocol allocations and pre-size encoded frames --- rust/loro-protocol/examples/encoding_bench.rs | 31 ++++++ rust/loro-protocol/src/bytes.rs | 11 +- rust/loro-protocol/src/elo.rs | 13 ++- rust/loro-protocol/src/encoding.rs | 104 ++++++++++++++++-- rust/loro-protocol/tests/decode_bounds.rs | 44 ++++++++ rust/loro-protocol/tests/elo.rs | 2 +- rust/loro-protocol/tests/encoding_limits.rs | 48 ++++++++ 7 files changed, 236 insertions(+), 17 deletions(-) create mode 100644 rust/loro-protocol/examples/encoding_bench.rs create mode 100644 rust/loro-protocol/tests/decode_bounds.rs create mode 100644 rust/loro-protocol/tests/encoding_limits.rs diff --git a/rust/loro-protocol/examples/encoding_bench.rs b/rust/loro-protocol/examples/encoding_bench.rs new file mode 100644 index 0000000..ac1273c --- /dev/null +++ b/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] + ); + } +} diff --git a/rust/loro-protocol/src/bytes.rs b/rust/loro-protocol/src/bytes.rs index 552e3c7..917cdd7 100644 --- a/rust/loro-protocol/src/bytes.rs +++ b/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; diff --git a/rust/loro-protocol/src/elo.rs b/rust/loro-protocol/src/elo.rs index 1ace443..1aa9872 100644 --- a/rust/loro-protocol/src/elo.rs +++ b/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, String> { +pub fn decode_elo_container(data: &[u8]) -> Result, 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, 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, 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 diff --git a/rust/loro-protocol/src/encoding.rs b/rust/loro-protocol/src/encoding.rs index a23c022..1c54910 100644 --- a/rust/loro-protocol/src/encoding.rs +++ b/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 { /// /// Returns an owned `Vec` containing the encoded message. pub fn encode(message: &ProtocolMessage) -> Result, 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, 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, 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 { } } 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()); diff --git a/rust/loro-protocol/tests/decode_bounds.rs b/rust/loro-protocol/tests/decode_bounds.rs new file mode 100644 index 0000000..cd4b16c --- /dev/null +++ b/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()); +} diff --git a/rust/loro-protocol/tests/elo.rs b/rust/loro-protocol/tests/elo.rs index c2237d7..103f3d3 100644 --- a/rust/loro-protocol/tests/elo.rs +++ b/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] diff --git a/rust/loro-protocol/tests/encoding_limits.rs b/rust/loro-protocol/tests/encoding_limits.rs new file mode 100644 index 0000000..4e9e4b0 --- /dev/null +++ b/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 ")); +}