From 294c75b32b26659f8fe0cc819abfd68231f59214 Mon Sep 17 00:00:00 2001 From: Parker Timmerman Date: Fri, 26 Jan 2024 12:14:33 -0500 Subject: [PATCH] start, make allocations in the decoding path fallible --- src/encoding.rs | 45 +++++++++++++++++++++++++++++++++++++-------- src/error.rs | 7 +++++++ src/types.rs | 2 +- 3 files changed, 45 insertions(+), 9 deletions(-) diff --git a/src/encoding.rs b/src/encoding.rs index e4d2aa274..ad2b605cd 100644 --- a/src/encoding.rs +++ b/src/encoding.rs @@ -4,7 +4,7 @@ #![allow(clippy::implicit_hasher, clippy::ptr_arg)] -use alloc::collections::BTreeMap; +use alloc::collections::{BTreeMap, TryReserveError}; use alloc::format; use alloc::string::String; use alloc::vec::Vec; @@ -452,6 +452,7 @@ macro_rules! merge_repeated_numeric { merge_loop(values, buf, ctx, |values, buf, ctx| { let mut value = Default::default(); $merge($wire_type, &mut value, buf, ctx)?; + values.try_reserve(1)?; values.push(value); Ok(()) }) @@ -460,6 +461,7 @@ macro_rules! merge_repeated_numeric { check_wire_type($wire_type, wire_type)?; let mut value = Default::default(); $merge(wire_type, &mut value, buf, ctx)?; + values.try_reserve(1)?; values.push(value); Ok(()) } @@ -771,7 +773,15 @@ macro_rules! length_delimited { check_wire_type(WireType::LengthDelimited, wire_type)?; let mut value = Default::default(); merge(wire_type, &mut value, buf, ctx)?; + + // Explicitly reserve before pushing so we can return an error instead of panicking from + // and out-of-memory exception. + // + // Note: Reserving a single element will still cause ammortizied growth of the Vec like + // pushing does, so there is no performance impact of reserving early. + values.try_reserve(1)?; values.push(value); + Ok(()) } @@ -877,12 +887,13 @@ pub trait BytesAdapter: sealed::BytesAdapter {} mod sealed { use super::{Buf, BufMut}; + use alloc::collections::TryReserveError; pub trait BytesAdapter: Default + Sized + 'static { fn len(&self) -> usize; /// Replace contents of this buffer with the contents of another buffer. - fn replace_with(&mut self, buf: B) + fn replace_with(&mut self, buf: B) -> Result<(), TryReserveError> where B: Buf; @@ -904,11 +915,12 @@ impl sealed::BytesAdapter for Bytes { Buf::remaining(self) } - fn replace_with(&mut self, mut buf: B) + fn replace_with(&mut self, mut buf: B) -> Result<(), TryReserveError> where B: Buf, { *self = buf.copy_to_bytes(buf.remaining()); + Ok(()) } fn append_to(&self, buf: &mut B) @@ -926,13 +938,14 @@ impl sealed::BytesAdapter for Vec { Vec::len(self) } - fn replace_with(&mut self, buf: B) + fn replace_with(&mut self, buf: B) -> Result<(), TryReserveError> where B: Buf, { self.clear(); - self.reserve(buf.remaining()); + self.try_reserve(buf.remaining())?; self.put(buf); + Ok(()) } fn append_to(&self, buf: &mut B) @@ -985,11 +998,11 @@ pub mod bytes { // This is intended for A and B both being Bytes so it is zero-copy. // Some combinations of A and B types may cause a double-copy, // in which case merge_one_copy() should be used instead. - value.replace_with(buf.copy_to_bytes(len)); + value.replace_with(buf.copy_to_bytes(len))?; Ok(()) } - pub(super) fn merge_one_copy( + pub fn merge_one_copy( wire_type: WireType, value: &mut A, buf: &mut B, @@ -1007,7 +1020,7 @@ pub mod bytes { let len = len as usize; // If we must copy, make sure to copy only once. - value.replace_with(buf.take(len)); + value.replace_with(buf.take(len))?; Ok(()) } @@ -1111,7 +1124,15 @@ pub mod message { check_wire_type(WireType::LengthDelimited, wire_type)?; let mut msg = M::default(); merge(WireType::LengthDelimited, &mut msg, buf, ctx)?; + + // Explicitly reserve before pushing so we can return an error instead of panicking from + // and out-of-memory exception. + // + // Note: Reserving a single element will still cause ammortizied growth of the Vec like + // pushing does, so there is no performance impact of reserving early. + messages.try_reserve(1)?; messages.push(msg); + Ok(()) } @@ -1202,7 +1223,15 @@ pub mod group { check_wire_type(WireType::StartGroup, wire_type)?; let mut msg = M::default(); merge(tag, WireType::StartGroup, &mut msg, buf, ctx)?; + + // Explicitly reserve before pushing so we can return an error instead of panicking from + // and out-of-memory exception. + // + // Note: Reserving a single element will still cause ammortizied growth of the Vec like + // pushing does, so there is no performance impact of reserving early. + messages.try_reserve(1)?; messages.push(msg); + Ok(()) } diff --git a/src/error.rs b/src/error.rs index 756ee8172..d59569944 100644 --- a/src/error.rs +++ b/src/error.rs @@ -2,6 +2,7 @@ use alloc::borrow::Cow; use alloc::boxed::Box; +use alloc::collections::TryReserveError; use alloc::vec::Vec; use core::fmt; @@ -69,6 +70,12 @@ impl fmt::Display for DecodeError { } } +impl From for DecodeError { + fn from(error: TryReserveError) -> Self { + DecodeError::new(error.to_string()) + } +} + #[cfg(feature = "std")] impl std::error::Error for DecodeError {} diff --git a/src/types.rs b/src/types.rs index 864a2adda..ce0369967 100644 --- a/src/types.rs +++ b/src/types.rs @@ -343,7 +343,7 @@ impl Message for Vec { B: Buf, { if tag == 1 { - bytes::merge(wire_type, self, buf, ctx) + bytes::merge_one_copy(wire_type, self, buf, ctx) } else { skip_field(wire_type, tag, buf, ctx) }