From ef888e06b8d5324449dfdd6d2d818cad872670cc Mon Sep 17 00:00:00 2001 From: Stepan Koltsov Date: Mon, 20 May 2019 04:07:28 +0100 Subject: [PATCH] Implement the same OOM fix for reading Bytes Nobody noticed likely because `with-bytes` feature is rarely used. CC: https://github.com/stepancheg/rust-protobuf/issues/411 --- protobuf/src/buf_read_iter.rs | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/protobuf/src/buf_read_iter.rs b/protobuf/src/buf_read_iter.rs index 4752f482f..c623a68f7 100644 --- a/protobuf/src/buf_read_iter.rs +++ b/protobuf/src/buf_read_iter.rs @@ -233,15 +233,21 @@ impl<'ignore> BufReadIter<'ignore> { self.pos_within_buf += len; Ok(r) } else { - let mut r = BytesMut::with_capacity(len); - unsafe { - { + if len >= READ_RAW_BYTES_MAX_ALLOC { + // We cannot trust `len` because protobuf message could be malformed. + // Reading should not result in OOM when allocating a buffer. + let mut v = Vec::new(); + self.read_exact_to_vec(len, &mut v)?; + Ok(Bytes::from(v)) + } else { + let mut r = BytesMut::with_capacity(len); + unsafe { let buf = &mut r.bytes_mut()[..len]; self.read_exact(buf)?; + r.advance_mut(len); } - r.advance_mut(len); + Ok(r.freeze()) } - Ok(r.freeze()) } }