      1
      2
      3
      4
      5
      6
      7
      8
      9
     10
     11
     12
     13
     14
     15
     16
     17
     18
     19
     20
     21
     22
     23
     24
     25
     26
     27
     28
     29
     30
     31
     32
     33
     34
     35
     36
     37
     38
     39
     40
     41
     42
     43
     44
     45
     46
     47
     48
     49
     50
     51
     52
     53
     54
     55
     56
     57
     58
     59
     60
     61
     62
     63
     64
     65
     66
     67
     68
     69
     70
     71
     72
     73
     74
     75
     76
     77
     78
     79
     80
     81
     82
     83
     84
     85
     86
     87
     88
     89
     90
     91
     92
     93
     94
     95
     96
     97
     98
     99
    100
    101
    102
    103
    104
    105
    106
    107
    108
    109
    110
    111
    112
    113
    114
    115
    116
    117
    118
    119
    120
    121
    122
    123
    124
    125
    126
    127
    128
    129
    130
    131
    132

    use crate::io::ReadBuf;
    use std::mem::MaybeUninit;

    mod private {
        pub trait Sealed {}

        impl Sealed for Vec<u8> {}
        impl Sealed for &mut Vec<u8> {}
    }

    /// A sealed trait that constrains the generic type parameter in `VecWithInitialized<V>`.  That struct's safety relies
    /// on certain invariants upheld by `Vec<u8>`.
    pub(crate) trait VecU8: AsMut<Vec<u8>> + private::Sealed {}

    impl VecU8 for Vec<u8> {}
    impl VecU8 for &mut Vec<u8> {}
    /// This struct wraps a `Vec<u8>` or `&mut Vec<u8>`, combining it with a
    /// `num_initialized`, which keeps track of the number of initialized bytes
    /// in the unused capacity.
    ///
    /// The purpose of this struct is to remember how many bytes were initialized
    /// through a `ReadBuf` from call to call.
    ///
    /// This struct has the safety invariant that the first `num_initialized` of the
    /// vector's allocation must be initialized at any time.
    #[derive(Debug)]
    pub(crate) struct VecWithInitialized<V> {
        vec: V,
        // The number of initialized bytes in the vector.
        // Always between `vec.len()` and `vec.capacity()`.
        num_initialized: usize,
    }

    impl VecWithInitialized<Vec<u8>> {
        #[cfg(feature = "io-util")]
        pub(crate) fn take(&mut self) -> Vec<u8> {
            self.num_initialized = 0;
            std::mem::take(&mut self.vec)
        }
    }

    impl<V> VecWithInitialized<V>
    where
        V: VecU8,
    {
        pub(crate) fn new(mut vec: V) -> Self {
            // SAFETY: The safety invariants of vector guarantee that the bytes up
            // to its length are initialized.
            Self {
                num_initialized: vec.as_mut().len(),
                vec,
            }
        }

        pub(crate) fn reserve(&mut self, num_bytes: usize) {
            let vec = self.vec.as_mut();
            if vec.capacity() - vec.len() >= num_bytes {
                return;
            }
            // SAFETY: Setting num_initialized to `vec.len()` is correct as
            // `reserve` does not change the length of the vector.
            self.num_initialized = vec.len();
            vec.reserve(num_bytes);
        }

        #[cfg(feature = "io-util")]
        pub(crate) fn is_empty(&mut self) -> bool {
            self.vec.as_mut().is_empty()
        }

        pub(crate) fn get_read_buf<'a>(&'a mut self) -> ReadBuf<'a> {
            let num_initialized = self.num_initialized;

            // SAFETY: Creating the slice is safe because of the safety invariants
            // on Vec<u8>. The safety invariants of `ReadBuf` will further guarantee
            // that no bytes in the slice are de-initialized.
            let vec = self.vec.as_mut();
            let len = vec.len();
            let cap = vec.capacity();
            let ptr = vec.as_mut_ptr().cast::<MaybeUninit<u8>>();
            let slice = unsafe { std::slice::from_raw_parts_mut::<'a, MaybeUninit<u8>>(ptr, cap) };

            // SAFETY: This is safe because the safety invariants of
            // VecWithInitialized say that the first num_initialized bytes must be
            // initialized.
            let mut read_buf = ReadBuf::uninit(slice);
            unsafe {
                read_buf.assume_init(num_initialized);
            }
            read_buf.set_filled(len);

            read_buf
        }

        pub(crate) fn apply_read_buf(&mut self, parts: ReadBufParts) {
            let vec = self.vec.as_mut();
            assert_eq!(vec.as_ptr(), parts.ptr);

            // SAFETY:
            // The ReadBufParts really does point inside `self.vec` due to the above
            // check, and the safety invariants of `ReadBuf` guarantee that the
            // first `parts.initialized` bytes of `self.vec` really have been
            // initialized. Additionally, `ReadBuf` guarantees that `parts.len` is
            // at most `parts.initialized`, so the first `parts.len` bytes are also
            // initialized.
            //
            // Note that this relies on the fact that `V` is either `Vec<u8>` or
            // `&mut Vec<u8>`, so the vector returned by `self.vec.as_mut()` cannot
            // change from call to call.
            unsafe {
                self.num_initialized = parts.initialized;
                vec.set_len(parts.len);
            }
        }
    }

    pub(crate) struct ReadBufParts {
        // Pointer is only used to check that the ReadBuf actually came from the
        // right VecWithInitialized.
        ptr: *const u8,
        len: usize,
        initialized: usize,
    }

    // This is needed to release the borrow on `VecWithInitialized<V>`.
    pub(crate) fn into_read_buf_parts(rb: ReadBuf<'_>) -> ReadBufParts {
        ReadBufParts {
            ptr: rb.filled().as_ptr(),
            len: rb.filled().len(),
            initialized: rb.initialized().len(),
        }
    }