Revision control
Copy as Markdown
Other Tools
use crate::Marker;
use std::num::{NonZeroU32, NonZeroUsize};
/// Incremental MessagePack parser that can parse incomplete messages,
/// and report their estimated total length.
pub struct MessageLen {
/// The last operation interrupted
wip: Option<WIP>,
/// Max size estimate
max_position: NonZeroUsize,
/// Bytes read so far
position: usize,
/// Stack of open arrays and maps
/// It is not a complete stack. Used only when resumption is needed.
sequences_wip: Vec<Seq>,
/// Nesting of arrays and maps
current_depth: u16,
/// Configured limit
max_depth: u16,
/// Configured limit
max_len: u32,
}
/// [`MessageLen`] result
#[derive(Debug)]
pub enum LenError {
/// The message is truncated, and needs at least this many bytes to parse
Truncated(NonZeroUsize),
/// The message is invalid or exceeded size limits
ParseError,
}
impl LenError {
/// Get expected min length or 0 on error
#[must_use]
pub fn len(&self) -> usize {
match *self {
Self::ParseError => 0,
Self::Truncated(l) => l.get(),
}
}
}
impl MessageLen {
/// New parser with default limits
///
/// If you have all MessagePack data in memory already, you can use [`MessageLen::len_of`].
/// If you're reading data in a streaming fashion, you can feed chunks of data
/// to [`MessageLen::incremental_len`].
#[must_use]
pub fn new() -> Self {
Self::with_limits(1024, (u32::MAX as usize).min(isize::MAX as usize / 2))
}
/// * `max_depth` limits nesting of arrays and maps
///
/// * `max_len` is maximum size of any string, byte string, map, or array.
/// For maps and arrays this is the number of items, not bytes.
///
/// Messages can be both deep and wide, being `max_depth` * `max_len` in size.
/// You should also limit the maximum byte size of the message (outside of this parser).
#[must_use]
pub fn with_limits(max_depth: usize, max_len: usize) -> Self {
Self {
max_position: NonZeroUsize::new(1).unwrap(),
position: 0,
current_depth: 0,
max_depth: max_depth.min(u16::MAX as _) as u16,
max_len: max_len.min(u32::MAX as _) as u32,
sequences_wip: Vec::new(),
wip: Some(WIP::NextMarker),
}
}
/// Parse the entire message to find if it's complete, and what is its serialized size in bytes.
///
/// If it returns `Ok(len)`, then the first `len` bytes of the given slice
/// parse as a single MessagePack object.
/// The length may be shorter than the slice given (extra data is gracefully ignored).
///
/// `Err(LenError::Truncated(len))` means that the the object is incomplete, the slice is truncated,
/// and it would need *at least* this many bytes to parse.
/// The `len` is always the lower bound, and never exceeds actual message length.
///
/// `Err(LenError::ParseError)` — the end of the message is unknown.
///
/// Don't call this function in a loop. Use [`MessageLen::incremental_len`] instead.
pub fn len_of(complete_message: &[u8]) -> Result<usize, LenError> {
Self::with_limits(1024, 1 << 30).incremental_len(complete_message.as_ref())
}
/// Parse more bytes, and re-evaluate required message length.
///
/// This function is stateful and keeps "appending" the data to its evaluation.
///
/// * `Ok(len)` — size of the whole MessagePack object, in bytes, starting at the beginning
/// of all data given to this function, including previous calls (not just this slice).
/// The object's data may end before the end of this slice. In such case the extra bytes
/// are gracefully ignored, and have not been parsed.
///
/// * `Err(LenError::Truncated(len))` — all bytes of this slice have been consumed,
/// and that was still not enough. The object needs at least `len` bytes in total
/// (counting from the start of all data given to this function, not just this slice).
/// The `len` is always the lower bound, and never exceeds actual message length,
/// so it's safe to read the additional bytes without overshooting the end of the message.
///
/// * `Err(LenError::ParseError)` — the end of the message cannot be determined, and this
/// is a non-recoverable error. Any further calls to this function may return nonsense.
pub fn incremental_len(&mut self, mut next_message_fragment: &[u8]) -> Result<usize, LenError> {
let data = &mut next_message_fragment;
let Some(wip) = self.wip.take() else {
return Ok(self.position); // must have succeded already
};
match wip {
WIP::Data(Data { bytes_left }) => self.skip_data(data, bytes_left.get()),
WIP::MarkerLen(wip) => self.read_marker_with_len(data, wip),
WIP::NextMarker => self.read_one_item(data),
WIP::LimitExceeded => {
self.wip = Some(WIP::LimitExceeded); // put it back!
return Err(LenError::ParseError);
},
}.ok_or(LenError::Truncated(self.max_position))?;
while let Some(seq) = self.sequences_wip.pop() {
self.current_depth = seq.depth;
debug_assert!(self.wip.is_none());
self.read_sequence(data, seq.items_left.get() - 1).ok_or(LenError::Truncated(self.max_position))?;
}
debug_assert!(self.wip.is_none());
debug_assert!(self.max_position.get() <= self.position);
Ok(self.position)
}
/// Forget all the state. The next call to `incremental_len` will assume it's the start of a new message.
pub fn reset(&mut self) {
self.max_position = NonZeroUsize::new(1).unwrap();
self.position = 0;
self.current_depth = 0;
self.sequences_wip.clear();
self.wip = Some(WIP::NextMarker);
}
fn read_one_item(&mut self, data: &mut &[u8]) -> Option<()> {
debug_assert!(self.wip.is_none());
let marker = self.read_marker(data)?;
match marker {
Marker::FixPos(_) => Some(()),
Marker::FixMap(len) => self.read_sequence(data, u32::from(len) * 2),
Marker::FixArray(len) => self.read_sequence(data, u32::from(len)),
Marker::FixStr(len) => self.skip_data(data, len.into()),
Marker::Null |
Marker::Reserved |
Marker::False |
Marker::True => Some(()),
Marker::Str8 |
Marker::Str16 |
Marker::Str32 |
Marker::Bin8 |
Marker::Bin16 |
Marker::Bin32 |
Marker::Array16 |
Marker::Array32 |
Marker::Map16 |
Marker::Map32 => self.read_marker_with_len(data, MarkerLen { marker, buf: [0; 4], has: 0 }),
Marker::Ext8 |
Marker::Ext16 |
Marker::Ext32 => todo!(),
Marker::F32 => self.skip_data(data, 4),
Marker::F64 => self.skip_data(data, 8),
Marker::U8 => self.skip_data(data, 1),
Marker::U16 => self.skip_data(data, 2),
Marker::U32 => self.skip_data(data, 4),
Marker::U64 => self.skip_data(data, 8),
Marker::I8 => self.skip_data(data, 1),
Marker::I16 => self.skip_data(data, 2),
Marker::I32 => self.skip_data(data, 4),
Marker::I64 => self.skip_data(data, 8),
Marker::FixExt1 |
Marker::FixExt2 |
Marker::FixExt4 |
Marker::FixExt8 |
Marker::FixExt16 => todo!(),
Marker::FixNeg(_) => Some(()),
}
}
fn read_marker_with_len(&mut self, data: &mut &[u8], mut wip: MarkerLen) -> Option<()> {
let size = wip.size();
debug_assert!(wip.has < size && size > 0 && size <= 4);
let dest = &mut wip.buf[0..size as usize];
let wanted = dest.len().checked_sub(wip.has as _)?;
let taken = self.take_bytes(data, wanted as u32);
dest[wip.has as usize..][..taken.len()].copy_from_slice(taken);
wip.has += taken.len() as u8;
if wip.has < size {
return self.fail(WIP::MarkerLen(wip));
}
let len = match dest.len() {
1 => dest[0].into(),
2 => u16::from_be_bytes(dest.try_into().unwrap()).into(),
4 => u32::from_be_bytes(dest.try_into().unwrap()),
_ => {
debug_assert!(false);
return None;
},
};
if len >= self.max_len {
return self.fail(WIP::LimitExceeded);
}
match wip.marker {
Marker::Bin8 |
Marker::Bin16 |
Marker::Bin32 |
Marker::Str8 |
Marker::Str16 |
Marker::Str32 => self.skip_data(data, len),
Marker::Ext8 |
Marker::Ext16 |
Marker::Ext32 => todo!(),
Marker::Array16 |
Marker::Array32 => self.read_sequence(data, len),
Marker::Map16 |
Marker::Map32 => {
if let Some(len) = len.checked_mul(2).filter(|&l| l < self.max_len) {
self.read_sequence(data, len)
} else {
self.fail(WIP::LimitExceeded)
}
},
_ => {
debug_assert!(false);
None
},
}
}
fn read_sequence(&mut self, data: &mut &[u8], mut items_left: u32) -> Option<()> {
self.current_depth += 1;
if self.current_depth > self.max_depth {
return self.fail(WIP::LimitExceeded);
}
while let Some(non_zero) = NonZeroU32::new(items_left) {
let position_before_item = self.position;
self.read_one_item(data).or_else(|| {
self.set_max_position(position_before_item + items_left as usize);
// -1, because it will increase depth again when resumed
self.sequences_wip.push(Seq {
items_left: non_zero,
depth: self.current_depth - 1,
});
None
})?;
items_left -= 1;
}
debug_assert!(self.current_depth > 0);
self.current_depth -= 1;
Some(())
}
fn skip_data(&mut self, data: &mut &[u8], wanted: u32) -> Option<()> {
let taken = self.take_bytes(data, wanted);
if let Some(bytes_left) = NonZeroU32::new(wanted - taken.len() as u32) {
debug_assert!(data.is_empty());
self.fail(WIP::Data(Data { bytes_left }))
} else {
Some(())
}
}
fn read_marker(&mut self, data: &mut &[u8]) -> Option<Marker> {
let Some((&b, rest)) = data.split_first() else {
debug_assert!(data.is_empty());
return self.fail(WIP::NextMarker);
};
self.position += 1;
*data = rest;
Some(Marker::from_u8(b))
}
fn set_max_position(&mut self, position: usize) {
self.max_position = NonZeroUsize::new(self.max_position.get().max(position)).unwrap();
}
/// May return less than requested
fn take_bytes<'data>(&mut self, data: &mut &'data [u8], wanted: u32) -> &'data [u8] {
let (taken, rest) = data.split_at(data.len().min(wanted as usize));
self.position += taken.len();
*data = rest;
taken
}
#[inline(always)]
fn fail<T>(&mut self, wip: WIP) -> Option<T> {
debug_assert!(self.wip.is_none());
let pos = match self.wip.insert(wip) {
WIP::NextMarker => self.position + 1,
WIP::Data(Data { bytes_left }) => self.position + bytes_left.get() as usize,
WIP::MarkerLen(m) => self.position + (m.size() - m.has) as usize,
WIP::LimitExceeded => 0,
};
self.set_max_position(pos);
None
}
}
enum WIP {
NextMarker,
Data(Data),
MarkerLen(MarkerLen),
LimitExceeded,
}
struct Seq { items_left: NonZeroU32, depth: u16 }
struct Data { bytes_left: NonZeroU32 }
struct MarkerLen { marker: Marker, buf: [u8; 4], has: u8 }
impl MarkerLen {
fn size(&self) -> u8 {
match self.marker {
Marker::Bin8 => 1,
Marker::Bin16 => 2,
Marker::Bin32 => 4,
Marker::Ext8 => 1,
Marker::Ext16 => 2,
Marker::Ext32 => 4,
Marker::Str8 => 1,
Marker::Str16 => 2,
Marker::Str32 => 4,
Marker::Array16 => 2,
Marker::Array32 => 4,
Marker::Map16 => 2,
Marker::Map32 => 4,
_ => unimplemented!(),
}
}
}