1 use crate::io::ReadBuf;
2 use std::mem::MaybeUninit;
3
4 mod private {
5 pub trait Sealed {}
6
7 impl Sealed for Vec<u8> {}
8 impl Sealed for &mut Vec<u8> {}
9 }
10
11 /// A sealed trait that constrains the generic type parameter in `VecWithInitialized<V>`. That struct's safety relies
12 /// on certain invariants upheld by `Vec<u8>`.
13 pub(crate) trait VecU8: AsMut<Vec<u8>> + private::Sealed {}
14
15 impl VecU8 for Vec<u8> {}
16 impl VecU8 for &mut Vec<u8> {}
17 /// This struct wraps a `Vec<u8>` or `&mut Vec<u8>`, combining it with a
18 /// `num_initialized`, which keeps track of the number of initialized bytes
19 /// in the unused capacity.
20 ///
21 /// The purpose of this struct is to remember how many bytes were initialized
22 /// through a `ReadBuf` from call to call.
23 ///
24 /// This struct has the safety invariant that the first `num_initialized` of the
25 /// vector's allocation must be initialized at any time.
26 #[derive(Debug)]
27 pub(crate) struct VecWithInitialized<V> {
28 vec: V,
29 // The number of initialized bytes in the vector.
30 // Always between `vec.len()` and `vec.capacity()`.
31 num_initialized: usize,
32 }
33
34 impl VecWithInitialized<Vec<u8>> {
35 #[cfg(feature = "io-util")]
take(&mut self) -> Vec<u8>36 pub(crate) fn take(&mut self) -> Vec<u8> {
37 self.num_initialized = 0;
38 std::mem::take(&mut self.vec)
39 }
40 }
41
42 impl<V> VecWithInitialized<V>
43 where
44 V: VecU8,
45 {
new(mut vec: V) -> Self46 pub(crate) fn new(mut vec: V) -> Self {
47 // SAFETY: The safety invariants of vector guarantee that the bytes up
48 // to its length are initialized.
49 Self {
50 num_initialized: vec.as_mut().len(),
51 vec,
52 }
53 }
54
reserve(&mut self, num_bytes: usize)55 pub(crate) fn reserve(&mut self, num_bytes: usize) {
56 let vec = self.vec.as_mut();
57 if vec.capacity() - vec.len() >= num_bytes {
58 return;
59 }
60 // SAFETY: Setting num_initialized to `vec.len()` is correct as
61 // `reserve` does not change the length of the vector.
62 self.num_initialized = vec.len();
63 vec.reserve(num_bytes);
64 }
65
66 #[cfg(feature = "io-util")]
is_empty(&mut self) -> bool67 pub(crate) fn is_empty(&mut self) -> bool {
68 self.vec.as_mut().is_empty()
69 }
70
get_read_buf<'a>(&'a mut self) -> ReadBuf<'a>71 pub(crate) fn get_read_buf<'a>(&'a mut self) -> ReadBuf<'a> {
72 let num_initialized = self.num_initialized;
73
74 // SAFETY: Creating the slice is safe because of the safety invariants
75 // on Vec<u8>. The safety invariants of `ReadBuf` will further guarantee
76 // that no bytes in the slice are de-initialized.
77 let vec = self.vec.as_mut();
78 let len = vec.len();
79 let cap = vec.capacity();
80 let ptr = vec.as_mut_ptr().cast::<MaybeUninit<u8>>();
81 let slice = unsafe { std::slice::from_raw_parts_mut::<'a, MaybeUninit<u8>>(ptr, cap) };
82
83 // SAFETY: This is safe because the safety invariants of
84 // VecWithInitialized say that the first num_initialized bytes must be
85 // initialized.
86 let mut read_buf = ReadBuf::uninit(slice);
87 unsafe {
88 read_buf.assume_init(num_initialized);
89 }
90 read_buf.set_filled(len);
91
92 read_buf
93 }
94
apply_read_buf(&mut self, parts: ReadBufParts)95 pub(crate) fn apply_read_buf(&mut self, parts: ReadBufParts) {
96 let vec = self.vec.as_mut();
97 assert_eq!(vec.as_ptr(), parts.ptr);
98
99 // SAFETY:
100 // The ReadBufParts really does point inside `self.vec` due to the above
101 // check, and the safety invariants of `ReadBuf` guarantee that the
102 // first `parts.initialized` bytes of `self.vec` really have been
103 // initialized. Additionally, `ReadBuf` guarantees that `parts.len` is
104 // at most `parts.initialized`, so the first `parts.len` bytes are also
105 // initialized.
106 //
107 // Note that this relies on the fact that `V` is either `Vec<u8>` or
108 // `&mut Vec<u8>`, so the vector returned by `self.vec.as_mut()` cannot
109 // change from call to call.
110 unsafe {
111 self.num_initialized = parts.initialized;
112 vec.set_len(parts.len);
113 }
114 }
115 }
116
117 pub(crate) struct ReadBufParts {
118 // Pointer is only used to check that the ReadBuf actually came from the
119 // right VecWithInitialized.
120 ptr: *const u8,
121 len: usize,
122 initialized: usize,
123 }
124
125 // This is needed to release the borrow on `VecWithInitialized<V>`.
into_read_buf_parts(rb: ReadBuf<'_>) -> ReadBufParts126 pub(crate) fn into_read_buf_parts(rb: ReadBuf<'_>) -> ReadBufParts {
127 ReadBufParts {
128 ptr: rb.filled().as_ptr(),
129 len: rb.filled().len(),
130 initialized: rb.initialized().len(),
131 }
132 }
133