diff options
author | Android Build Coastguard Worker <android-build-coastguard-worker@google.com> | 2022-03-08 02:05:24 +0000 |
---|---|---|
committer | Android Build Coastguard Worker <android-build-coastguard-worker@google.com> | 2022-03-08 02:05:24 +0000 |
commit | 654557256191b7445462d18abfc80ae9459379e4 (patch) | |
tree | f0d3f05c18afd04b9824b485585518fe06641ee1 /src/coded_output_stream.rs | |
parent | c6e248b58bc0927b1a69c56d43a497361ae2e095 (diff) | |
parent | 0af506eab5a6d52635e775e0ad9b6439da826a3a (diff) | |
download | protobuf-android-security-13.0.0_r10.tar.gz |
Snap for 8270536 from 0af506eab5a6d52635e775e0ad9b6439da826a3a to tm-releaseandroid-vts-13.0_r8android-vts-13.0_r7android-vts-13.0_r6android-vts-13.0_r5android-vts-13.0_r4android-vts-13.0_r3android-vts-13.0_r2android-vts-13.0_r1android-security-13.0.0_r9android-security-13.0.0_r8android-security-13.0.0_r7android-security-13.0.0_r6android-security-13.0.0_r5android-security-13.0.0_r4android-security-13.0.0_r3android-security-13.0.0_r2android-security-13.0.0_r19android-security-13.0.0_r18android-security-13.0.0_r17android-security-13.0.0_r16android-security-13.0.0_r15android-security-13.0.0_r14android-security-13.0.0_r13android-security-13.0.0_r12android-security-13.0.0_r11android-security-13.0.0_r10android-security-13.0.0_r1android-platform-13.0.0_r9android-platform-13.0.0_r8android-platform-13.0.0_r7android-platform-13.0.0_r6android-platform-13.0.0_r5android-platform-13.0.0_r4android-platform-13.0.0_r3android-platform-13.0.0_r21android-platform-13.0.0_r20android-platform-13.0.0_r2android-platform-13.0.0_r19android-platform-13.0.0_r18android-platform-13.0.0_r17android-platform-13.0.0_r16android-platform-13.0.0_r15android-platform-13.0.0_r14android-platform-13.0.0_r13android-platform-13.0.0_r12android-platform-13.0.0_r11android-platform-13.0.0_r10android-platform-13.0.0_r1android-cts-13.0_r8android-cts-13.0_r7android-cts-13.0_r6android-cts-13.0_r5android-cts-13.0_r4android-cts-13.0_r3android-cts-13.0_r2android-cts-13.0_r1android-13.0.0_r8android-13.0.0_r7android-13.0.0_r6android-13.0.0_r5android-13.0.0_r4android-13.0.0_r31android-13.0.0_r3android-13.0.0_r2android-13.0.0_r12android-13.0.0_r1android13-tests-releaseandroid13-security-releaseandroid13-s3-releaseandroid13-s2-releaseandroid13-s1-releaseandroid13-releaseandroid13-platform-releaseandroid13-gsi
Change-Id: I1afc3e5e65c4576b5c9584ef20cb5bdbf6211d5a
Diffstat (limited to 'src/coded_output_stream.rs')
-rw-r--r-- | src/coded_output_stream.rs | 121 |
1 files changed, 86 insertions, 35 deletions
diff --git a/src/coded_output_stream.rs b/src/coded_output_stream.rs index bbfe228..2bbe0a3 100644 --- a/src/coded_output_stream.rs +++ b/src/coded_output_stream.rs @@ -1,17 +1,23 @@ -use crate::misc::remaining_capacity_as_slice_mut; -use crate::misc::remove_lifetime_mut; +use std::io; +use std::io::Write; +use std::mem; +use std::mem::MaybeUninit; +use std::ptr; +use std::slice; + +use crate::misc::maybe_uninit_write; +use crate::misc::maybe_uninit_write_slice; +use crate::misc::vec_spare_capacity_mut; use crate::varint; use crate::wire_format; use crate::zigzag::encode_zig_zag_32; use crate::zigzag::encode_zig_zag_64; use crate::Message; use crate::ProtobufEnum; +use crate::ProtobufError; use crate::ProtobufResult; use crate::UnknownFields; use crate::UnknownValueRef; -use std::io; -use std::io::Write; -use std::mem; /// Equal to the default buffer size of `BufWriter`, so when /// `CodedOutputStream` wraps `BufWriter`, it often skips double buffering. @@ -58,17 +64,27 @@ where Ok(v) } +/// Output buffer/writer for `CodedOutputStream`. enum OutputTarget<'a> { Write(&'a mut dyn Write, Vec<u8>), Vec(&'a mut Vec<u8>), + /// The buffer is passed as `&[u8]` to `CodedOutputStream` constructor + /// and immediately converted to `buffer` field of `CodedOutputStream`, + /// it is not needed to be stored here. + /// Lifetime parameter of `CodedOutputStream` guarantees the buffer is valid + /// during the lifetime of `CodedOutputStream`. Bytes, } /// Buffered write with handy utilities pub struct CodedOutputStream<'a> { target: OutputTarget<'a>, - // alias to buf from target - buffer: &'a mut [u8], + // Actual buffer is owned by `OutputTarget`, + // and here we alias the buffer so access to the buffer is branchless: + // access does not require switch by actual target type: `&[], `Vec`, `Write` etc. + // We don't access the actual buffer in `OutputTarget` except when + // we initialize `buffer` field here. + buffer: *mut [MaybeUninit<u8>], // within buffer position: usize, } @@ -81,15 +97,16 @@ impl<'a> CodedOutputStream<'a> { let buffer_len = OUTPUT_STREAM_BUFFER_SIZE; let mut buffer_storage = Vec::with_capacity(buffer_len); - unsafe { - buffer_storage.set_len(buffer_len); - } - let buffer = unsafe { remove_lifetime_mut(&mut buffer_storage as &mut [u8]) }; + // SAFETY: we are not using the `buffer_storage` + // except for initializing the `buffer` field. + // See `buffer` field documentation. + let buffer = vec_spare_capacity_mut(&mut buffer_storage); + let buffer: *mut [MaybeUninit<u8>] = buffer; CodedOutputStream { target: OutputTarget::Write(writer, buffer_storage), - buffer: buffer, + buffer, position: 0, } } @@ -98,9 +115,12 @@ impl<'a> CodedOutputStream<'a> { /// /// Attempt to write more than bytes capacity results in error. pub fn bytes(bytes: &'a mut [u8]) -> CodedOutputStream<'a> { + // SAFETY: it is safe to cast from &mut [u8] to &mut [MaybeUninit<u8>]. + let buffer = + ptr::slice_from_raw_parts_mut(bytes.as_mut_ptr() as *mut MaybeUninit<u8>, bytes.len()); CodedOutputStream { target: OutputTarget::Bytes, - buffer: bytes, + buffer, position: 0, } } @@ -110,9 +130,10 @@ impl<'a> CodedOutputStream<'a> { /// Caller should call `flush` at the end to guarantee vec contains /// all written data. pub fn vec(vec: &'a mut Vec<u8>) -> CodedOutputStream<'a> { + let buffer: *mut [MaybeUninit<u8>] = &mut []; CodedOutputStream { target: OutputTarget::Vec(vec), - buffer: &mut [], + buffer, position: 0, } } @@ -125,7 +146,7 @@ impl<'a> CodedOutputStream<'a> { pub fn check_eof(&self) { match self.target { OutputTarget::Bytes => { - assert_eq!(self.buffer.len() as u64, self.position as u64); + assert_eq!(self.buffer().len() as u64, self.position as u64); } OutputTarget::Write(..) | OutputTarget::Vec(..) => { panic!("must not be called with Writer or Vec"); @@ -133,10 +154,25 @@ impl<'a> CodedOutputStream<'a> { } } + #[inline(always)] + fn buffer(&self) -> &[MaybeUninit<u8>] { + // SAFETY: see the `buffer` field documentation about invariants. + unsafe { &*(self.buffer as *mut [MaybeUninit<u8>]) } + } + + #[inline(always)] + fn filled_buffer_impl<'s>(buffer: *mut [MaybeUninit<u8>], position: usize) -> &'s [u8] { + // SAFETY: this function is safe assuming `buffer` and `position` + // are `self.buffer` and `safe.position`: + // * `CodedOutputStream` has invariant that `position <= buffer.len()`. + // * `buffer` is filled up to `position`. + unsafe { slice::from_raw_parts_mut(buffer as *mut u8, position) } + } + fn refresh_buffer(&mut self) -> ProtobufResult<()> { match self.target { OutputTarget::Write(ref mut write, _) => { - write.write_all(&self.buffer[0..self.position as usize])?; + write.write_all(Self::filled_buffer_impl(self.buffer, self.position))?; self.position = 0; } OutputTarget::Vec(ref mut vec) => unsafe { @@ -144,11 +180,14 @@ impl<'a> CodedOutputStream<'a> { assert!(vec_len + self.position <= vec.capacity()); vec.set_len(vec_len + self.position); vec.reserve(1); - self.buffer = remove_lifetime_mut(remaining_capacity_as_slice_mut(vec)); + self.buffer = vec_spare_capacity_mut(vec); self.position = 0; }, OutputTarget::Bytes => { - panic!("refresh_buffer must not be called on CodedOutputStream create from slice"); + return Err(ProtobufError::IoError(io::Error::new( + io::ErrorKind::Other, + "given slice is too small to serialize the message", + ))); } } Ok(()) @@ -167,20 +206,22 @@ impl<'a> CodedOutputStream<'a> { /// Write a byte pub fn write_raw_byte(&mut self, byte: u8) -> ProtobufResult<()> { - if self.position as usize == self.buffer.len() { + if self.position as usize == self.buffer().len() { self.refresh_buffer()?; } - self.buffer[self.position as usize] = byte; + unsafe { maybe_uninit_write(&mut (&mut *self.buffer)[self.position as usize], byte) }; self.position += 1; Ok(()) } /// Write bytes pub fn write_raw_bytes(&mut self, bytes: &[u8]) -> ProtobufResult<()> { - if bytes.len() <= self.buffer.len() - self.position { + if bytes.len() <= self.buffer().len() - self.position { let bottom = self.position as usize; let top = bottom + (bytes.len() as usize); - self.buffer[bottom..top].copy_from_slice(bytes); + // SAFETY: see the `buffer` field documentation about invariants. + let buffer = unsafe { &mut (&mut *self.buffer)[bottom..top] }; + maybe_uninit_write_slice(buffer, bytes); self.position += bytes.len(); return Ok(()); } @@ -189,8 +230,11 @@ impl<'a> CodedOutputStream<'a> { assert!(self.position == 0); - if self.position + bytes.len() < self.buffer.len() { - self.buffer[self.position..self.position + bytes.len()].copy_from_slice(bytes); + if self.position + bytes.len() < self.buffer().len() { + // SAFETY: see the `buffer` field documentation about invariants. + let buffer = + unsafe { &mut (&mut *self.buffer)[self.position..self.position + bytes.len()] }; + maybe_uninit_write_slice(buffer, bytes); self.position += bytes.len(); return Ok(()); } @@ -204,9 +248,7 @@ impl<'a> CodedOutputStream<'a> { } OutputTarget::Vec(ref mut vec) => { vec.extend(bytes); - unsafe { - self.buffer = remove_lifetime_mut(remaining_capacity_as_slice_mut(vec)); - } + self.buffer = vec_spare_capacity_mut(vec) } } Ok(()) @@ -223,30 +265,38 @@ impl<'a> CodedOutputStream<'a> { /// Write varint pub fn write_raw_varint32(&mut self, value: u32) -> ProtobufResult<()> { - if self.buffer.len() - self.position >= 5 { + if self.buffer().len() - self.position >= 5 { // fast path - let len = varint::encode_varint32(value, &mut self.buffer[self.position..]); + let len = unsafe { + varint::encode_varint32(value, &mut (&mut *self.buffer)[self.position..]) + }; self.position += len; Ok(()) } else { // slow path let buf = &mut [0u8; 5]; - let len = varint::encode_varint32(value, buf); + let len = varint::encode_varint32(value, unsafe { + slice::from_raw_parts_mut(buf.as_mut_ptr() as *mut MaybeUninit<u8>, buf.len()) + }); self.write_raw_bytes(&buf[..len]) } } /// Write varint pub fn write_raw_varint64(&mut self, value: u64) -> ProtobufResult<()> { - if self.buffer.len() - self.position >= 10 { + if self.buffer().len() - self.position >= 10 { // fast path - let len = varint::encode_varint64(value, &mut self.buffer[self.position..]); + let len = unsafe { + varint::encode_varint64(value, &mut (&mut *self.buffer)[self.position..]) + }; self.position += len; Ok(()) } else { // slow path let buf = &mut [0u8; 10]; - let len = varint::encode_varint64(value, buf); + let len = varint::encode_varint64(value, unsafe { + slice::from_raw_parts_mut(buf.as_mut_ptr() as *mut MaybeUninit<u8>, buf.len()) + }); self.write_raw_bytes(&buf[..len]) } } @@ -532,13 +582,14 @@ impl<'a> Write for CodedOutputStream<'a> { #[cfg(test)] mod test { + use std::io::Write; + use std::iter; + use crate::coded_output_stream::CodedOutputStream; use crate::hex::decode_hex; use crate::hex::encode_hex; use crate::wire_format; use crate::ProtobufResult; - use std::io::Write; - use std::iter; fn test_write<F>(expected: &str, mut gen: F) where |