From e54863c770baa67ac46ad4cd8c55b6e4267d030e Mon Sep 17 00:00:00 2001 From: Louis Barthonet Date: Mon, 14 Sep 2026 14:41:20 -0700 Subject: [PATCH] THRIFT-6290: Rust: skip string/binary fields without heap-allocating Client: rs skip() used read_bytes() and dropped the Vec. Binary and compact now read the length prefix and discard the payload through a stack buffer. Co-Authored-By: Grok 4.7 --- lib/rs/src/protocol/binary.rs | 79 +++++++++---- lib/rs/src/protocol/compact.rs | 40 ++++--- lib/rs/src/protocol/mod.rs | 207 ++++++++++++++++++++++++++++++++- lib/rs/src/protocol/stored.rs | 4 + 4 files changed, 291 insertions(+), 39 deletions(-) diff --git a/lib/rs/src/protocol/binary.rs b/lib/rs/src/protocol/binary.rs index 7faf73cd09d..84c5f01c6fa 100644 --- a/lib/rs/src/protocol/binary.rs +++ b/lib/rs/src/protocol/binary.rs @@ -93,6 +93,31 @@ where } Ok(()) } + + fn read_binary_len(&mut self) -> crate::Result { + let num_bytes = self.transport.read_i32::()?; + + if num_bytes < 0 { + return Err(crate::Error::Protocol(ProtocolError::new( + ProtocolErrorKind::NegativeSize, + format!("Negative byte array size: {}", num_bytes), + ))); + } + + if let Some(max_size) = self.config.max_string_size() { + if num_bytes as usize > max_size { + return Err(crate::Error::Protocol(ProtocolError::new( + ProtocolErrorKind::SizeLimit, + format!( + "Byte array size {} exceeds maximum allowed size of {}", + num_bytes, max_size + ), + ))); + } + } + + Ok(num_bytes as usize) + } } impl TInputProtocol for TBinaryInputProtocol @@ -192,34 +217,19 @@ where } fn read_bytes(&mut self) -> crate::Result> { - let num_bytes = self.transport.read_i32::()?; - - if num_bytes < 0 { - return Err(crate::Error::Protocol(ProtocolError::new( - ProtocolErrorKind::NegativeSize, - format!("Negative byte array size: {}", num_bytes), - ))); - } - - if let Some(max_size) = self.config.max_string_size() { - if num_bytes as usize > max_size { - return Err(crate::Error::Protocol(ProtocolError::new( - ProtocolErrorKind::SizeLimit, - format!( - "Byte array size {} exceeds maximum allowed size of {}", - num_bytes, max_size - ), - ))); - } - } - - let mut buf = vec![0u8; num_bytes as usize]; + let num_bytes = self.read_binary_len()?; + let mut buf = vec![0u8; num_bytes]; self.transport .read_exact(&mut buf) .map(|_| buf) .map_err(From::from) } + fn skip_binary(&mut self) -> crate::Result<()> { + let num_bytes = self.read_binary_len()?; + super::discard_exact(&mut self.transport, num_bytes).map_err(From::from) + } + fn read_bool(&mut self) -> crate::Result { let b = self.read_i8()?; match b { @@ -1176,6 +1186,31 @@ mod tests { assert_eq!(i_prot.recursion_depth, 0); } + #[test] + fn must_reject_negative_binary_size() { + let mem = TBufferChannel::with_capacity(16, 16); + let (r_mem, mut w_mem) = mem.split().unwrap(); + let mut i_prot = TBinaryInputProtocol::new(r_mem, true); + w_mem.set_readable_bytes(&[0xFF, 0xFF, 0xFF, 0xFF]); + match i_prot.read_bytes() { + Err(crate::Error::Protocol(e)) => { + assert_eq!(e.kind, ProtocolErrorKind::NegativeSize); + } + other => panic!("Expected NegativeSize, got {:?}", other), + } + + let mem = TBufferChannel::with_capacity(16, 16); + let (r_mem, mut w_mem) = mem.split().unwrap(); + let mut i_prot = TBinaryInputProtocol::new(r_mem, true); + w_mem.set_readable_bytes(&[0xFF, 0xFF, 0xFF, 0xFF]); + match i_prot.skip_binary() { + Err(crate::Error::Protocol(e)) => { + assert_eq!(e.kind, ProtocolErrorKind::NegativeSize); + } + other => panic!("Expected NegativeSize, got {:?}", other), + } + } + #[test] fn must_reject_negative_container_sizes() { let mem = TBufferChannel::with_capacity(40, 40); diff --git a/lib/rs/src/protocol/compact.rs b/lib/rs/src/protocol/compact.rs index 0809e907470..e8f35c3a17f 100644 --- a/lib/rs/src/protocol/compact.rs +++ b/lib/rs/src/protocol/compact.rs @@ -158,6 +158,24 @@ where "Variable-length int over 10 bytes.", ))) } + + fn read_binary_len(&mut self) -> crate::Result { + let len = self.read_varint32()?; + + if let Some(max_size) = self.config.max_string_size() { + if len as usize > max_size { + return Err(crate::Error::Protocol(ProtocolError::new( + ProtocolErrorKind::SizeLimit, + format!( + "Byte array size {} exceeds maximum allowed size of {}", + len, max_size + ), + ))); + } + } + + Ok(len as usize) + } } impl TInputProtocol for TCompactInputProtocol @@ -302,27 +320,19 @@ where } fn read_bytes(&mut self) -> crate::Result> { - let len = self.read_varint32()?; - - if let Some(max_size) = self.config.max_string_size() { - if len as usize > max_size { - return Err(crate::Error::Protocol(ProtocolError::new( - ProtocolErrorKind::SizeLimit, - format!( - "Byte array size {} exceeds maximum allowed size of {}", - len, max_size - ), - ))); - } - } - - let mut buf = vec![0u8; len as usize]; + let len = self.read_binary_len()?; + let mut buf = vec![0u8; len]; self.transport .read_exact(&mut buf) .map_err(From::from) .map(|_| buf) } + fn skip_binary(&mut self) -> crate::Result<()> { + let len = self.read_binary_len()?; + super::discard_exact(&mut self.transport, len).map_err(From::from) + } + fn read_i8(&mut self) -> crate::Result { self.read_byte().map(|i| i as i8) } diff --git a/lib/rs/src/protocol/mod.rs b/lib/rs/src/protocol/mod.rs index 0d952af34a0..6ce63e16e3c 100644 --- a/lib/rs/src/protocol/mod.rs +++ b/lib/rs/src/protocol/mod.rs @@ -60,6 +60,7 @@ use std::convert::{From, TryFrom}; use std::fmt; use std::fmt::{Display, Formatter}; +use std::io::Read; use crate::transport::{TReadTransport, TWriteTransport}; use crate::{ProtocolError, ProtocolErrorKind, TConfiguration}; @@ -208,7 +209,7 @@ pub trait TInputProtocol { TType::I32 => self.read_i32().map(|_| ()), TType::I64 => self.read_i64().map(|_| ()), TType::Double => self.read_double().map(|_| ()), - TType::String => self.read_bytes().map(|_| ()), + TType::String => self.skip_binary(), TType::Uuid => self.read_uuid().map(|_| ()), TType::Struct => { self.read_struct_begin()?; @@ -256,6 +257,11 @@ pub trait TInputProtocol { } } + /// Skip a binary or string field payload. + fn skip_binary(&mut self) -> crate::Result<()> { + self.read_bytes().map(|_| ()) + } + // utility (DO NOT USE IN GENERATED CODE!!!!) // @@ -399,6 +405,10 @@ where (**self).read_bytes() } + fn skip_binary(&mut self) -> crate::Result<()> { + (**self).skip_binary() + } + fn read_i8(&mut self) -> crate::Result { (**self).read_i8() } @@ -997,6 +1007,17 @@ pub(crate) fn check_container_size( } } +pub(crate) fn discard_exact(reader: &mut R, mut count: usize) -> std::io::Result<()> { + const CHUNK: usize = 256; + let mut buf = [0u8; CHUNK]; + while count > 0 { + let n = count.min(CHUNK); + reader.read_exact(&mut buf[..n])?; + count -= n; + } + Ok(()) +} + /// Extract the field id from a Thrift field identifier. /// /// `field_ident` must *not* have `TFieldIdentifier.field_type` of type `TType::Stop`. @@ -1014,7 +1035,9 @@ pub fn field_id(field_ident: &TFieldIdentifier) -> crate::Result { #[cfg(test)] mod tests { - use std::io::Cursor; + use std::cell::Cell; + use std::io::{Cursor, Read}; + use std::rc::Rc; use super::*; use crate::transport::{TReadTransport, TWriteTransport}; @@ -1133,4 +1156,184 @@ mod tests { let data = build_struct_with_unknown_binary_field(&[]); assert_eq!(read_struct_skipping_unknown(&data).unwrap(), 42); } + + fn build_struct_with_unknown_binary_then_i64(payload: &[u8], after: i64) -> Vec { + let mut buf = build_struct_with_unknown_binary_field(payload); + assert_eq!(buf.pop(), Some(0x00)); + buf.push(0x0A); // field 2: TType::I64 + buf.extend_from_slice(&2_i16.to_be_bytes()); + buf.extend_from_slice(&after.to_be_bytes()); + buf.push(0x00); // stop + buf + } + + fn skip_unknown_and_read_i64_fields( + proto: &mut P, + ) -> crate::Result<(i64, Option)> { + proto.read_struct_begin()?; + let mut first = None; + let mut second = None; + loop { + let field = proto.read_field_begin()?; + if field.field_type == TType::Stop { + break; + } + match field.id { + Some(1) if field.field_type == TType::I64 => { + first = Some(proto.read_i64()?); + } + Some(2) if field.field_type == TType::I64 => { + second = Some(proto.read_i64()?); + } + _ => { + proto.skip(field.field_type)?; + } + } + proto.read_field_end()?; + } + proto.read_struct_end()?; + Ok(( + first.ok_or_else(|| { + crate::Error::Protocol(crate::ProtocolError { + kind: crate::ProtocolErrorKind::InvalidData, + message: "missing known field".to_string(), + }) + })?, + second, + )) + } + + /// Records the largest slice `read` is asked to fill. `read_exact` passes + /// the full requested buffer, so this is the size the protocol asked for, + /// not the number of bytes returned. + struct PeakRead { + inner: R, + peak: Rc>, + } + + impl Read for PeakRead { + fn read(&mut self, buf: &mut [u8]) -> std::io::Result { + self.peak.set(self.peak.get().max(buf.len())); + self.inner.read(buf) + } + } + + fn peak_reader(data: Vec) -> (PeakRead>>, Rc>) { + let peak = Rc::new(Cell::new(0)); + ( + PeakRead { + inner: Cursor::new(data), + peak: Rc::clone(&peak), + }, + peak, + ) + } + + fn assert_skip_fill_at_most_256(peak: &Cell) { + assert!( + peak.get() <= 256, + "skip asked the transport to fill {} bytes", + peak.get() + ); + } + + #[test] + fn must_skip_large_binary_field_and_read_following_field() { + let payload = vec![0xABu8; 1024]; + let data = build_struct_with_unknown_binary_then_i64(&payload, 7); + let (reader, peak) = peak_reader(data); + let mut proto = TBinaryInputProtocol::new(reader, true); + let (first, second) = skip_unknown_and_read_i64_fields(&mut proto).unwrap(); + assert_eq!(first, 42); + assert_eq!(second, Some(7)); + assert_skip_fill_at_most_256(&peak); + } + + #[test] + fn must_skip_large_binary_field_through_boxed_protocol() { + let payload = vec![0xCDu8; 1024]; + let data = build_struct_with_unknown_binary_then_i64(&payload, 9); + let (reader, peak) = peak_reader(data); + let mut proto: Box = Box::new(TBinaryInputProtocol::new(reader, true)); + let (first, second) = skip_unknown_and_read_i64_fields(&mut proto).unwrap(); + assert_eq!(first, 42); + assert_eq!(second, Some(9)); + assert_skip_fill_at_most_256(&peak); + } + + #[test] + fn must_skip_large_binary_field_through_stored_protocol() { + let payload = vec![0x11u8; 1024]; + let data = build_struct_with_unknown_binary_then_i64(&payload, 13); + let (reader, peak) = peak_reader(data); + let mut inner = TBinaryInputProtocol::new(reader, true); + let mut proto = TStoredInputProtocol::new( + &mut inner, + TMessageIdentifier::new("unused", TMessageType::Call, 0), + ); + let (first, second) = skip_unknown_and_read_i64_fields(&mut proto).unwrap(); + assert_eq!(first, 42); + assert_eq!(second, Some(13)); + assert_skip_fill_at_most_256(&peak); + } + + #[test] + fn must_skip_binary_respects_max_string_size() { + let payload = vec![0u8; 32]; + let data = build_struct_with_unknown_binary_field(&payload); + let config = crate::TConfiguration::builder() + .max_string_size(Some(8)) + .build() + .unwrap(); + let mut proto = TBinaryInputProtocol::with_config(Cursor::new(data), true, config); + proto.read_struct_begin().unwrap(); + let field = proto.read_field_begin().unwrap(); + assert_eq!(field.id, Some(1)); + let _ = proto.read_i64().unwrap(); + proto.read_field_end().unwrap(); + let field = proto.read_field_begin().unwrap(); + assert_eq!(field.id, Some(99)); + let err = proto.skip(field.field_type).unwrap_err(); + match err { + crate::Error::Protocol(p) => { + assert_eq!(p.kind, crate::ProtocolErrorKind::SizeLimit); + } + other => panic!("expected SizeLimit, got {:?}", other), + } + } + + fn compact_struct_with_unknown_binary_then_i64(payload: &[u8], after: i64) -> Vec { + let mut buf = Vec::new(); + { + let mut o = TCompactOutputProtocol::new(&mut buf); + o.write_struct_begin(&TStructIdentifier::new("S")).unwrap(); + o.write_field_begin(&TFieldIdentifier::new("known", TType::I64, 1)) + .unwrap(); + o.write_i64(42).unwrap(); + o.write_field_end().unwrap(); + o.write_field_begin(&TFieldIdentifier::new("bin", TType::String, 99)) + .unwrap(); + o.write_bytes(payload).unwrap(); + o.write_field_end().unwrap(); + o.write_field_begin(&TFieldIdentifier::new("after", TType::I64, 2)) + .unwrap(); + o.write_i64(after).unwrap(); + o.write_field_end().unwrap(); + o.write_field_stop().unwrap(); + o.write_struct_end().unwrap(); + } + buf + } + + #[test] + fn must_skip_large_binary_field_compact_and_read_following_field() { + let payload = vec![0xEFu8; 1024]; + let data = compact_struct_with_unknown_binary_then_i64(&payload, 11); + let (reader, peak) = peak_reader(data); + let mut proto = TCompactInputProtocol::new(reader); + let (first, second) = skip_unknown_and_read_i64_fields(&mut proto).unwrap(); + assert_eq!(first, 42); + assert_eq!(second, Some(11)); + assert_skip_fill_at_most_256(&peak); + } } diff --git a/lib/rs/src/protocol/stored.rs b/lib/rs/src/protocol/stored.rs index 04d3277faf3..2274505d49a 100644 --- a/lib/rs/src/protocol/stored.rs +++ b/lib/rs/src/protocol/stored.rs @@ -138,6 +138,10 @@ impl<'a> TInputProtocol for TStoredInputProtocol<'a> { self.inner.read_bytes() } + fn skip_binary(&mut self) -> crate::Result<()> { + self.inner.skip_binary() + } + fn read_bool(&mut self) -> crate::Result { self.inner.read_bool() }