Skip to content

Commit

Permalink
Implement the same OOM fix for reading Bytes
Browse files Browse the repository at this point in the history
Nobody noticed likely because `with-bytes` feature is rarely used.

CC: #411
  • Loading branch information
stepancheg committed May 20, 2019
1 parent 3fa92bb commit ef888e0
Showing 1 changed file with 11 additions and 5 deletions.
16 changes: 11 additions & 5 deletions protobuf/src/buf_read_iter.rs
Expand Up @@ -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())
}
}

Expand Down

0 comments on commit ef888e0

Please sign in to comment.