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
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
use super::error::CapacityExceeded;
use bytes::buf::UninitSlice;
use bytes::BufMut;
use core::mem;
use core::mem::MaybeUninit;

/// Non panic wrapper for `&mut [u8]` that implement [`BufMut`](bytes::buf::BufMut) trait.
///
/// # Example
/// ### work with uninitialized memory `&mut [MaybeUninit<u8>]`
/// ```rust
/// use bytes::buf::BufMut;
/// use static_bytes::SafeBytesSlice;
/// use core::mem::MaybeUninit;
/// //are you sure that's random?
/// fn fill_with_random(buf: &mut dyn BufMut, amount: usize) {
///     for _ in 0..amount {
///         buf.put_u8(9);
///     }
/// }
/// // This is optimized way of working with slices
/// // This is safe. See way in rust doc:
/// // https://doc.rust-lang.org/std/mem/union.MaybeUninit.html#initializing-an-array-element-by-element
/// let mut fixed_storage: [MaybeUninit<u8>; 128] = unsafe {
///        MaybeUninit::uninit().assume_init()
///    };
/// let mut slice = SafeBytesSlice::from(&mut fixed_storage[..]);
/// // your function that accept `&mut dyn BufMut`
/// fill_with_random(&mut slice, 32);
/// let output = slice.try_into_bytes().unwrap();
/// assert_eq!(output.len(), 32);
/// assert_eq!(output[31], 9);
/// ```
///
/// ### work with standard slice `&mut [u8]`
/// ```rust
/// use bytes::buf::BufMut;
/// use static_bytes::SafeBytesSlice;
/// # //are you sure that's random?
/// # fn fill_with_random(buf: &mut dyn BufMut, amount: usize) {
/// #    for _ in 0..amount {
/// #        buf.put_u8(9);
/// #    }
/// # }
///
/// let mut fixed_storage = [0u8; 64];
/// let mut slice = SafeBytesSlice::from(&mut fixed_storage[..]);
/// // your function that accept `&mut dyn BufMut`
/// // see function impl in `&mut [MaybeUninit<u8>]` example
/// fill_with_random(&mut slice, 32);
/// let output = slice.try_into_bytes().unwrap();
/// assert_eq!(output.len(), 32);
/// assert_eq!(output[31], 9);
/// ```
///
/// Is `fill_with_random()` random?
/// ![](https://starecat.com/content/wp-content/uploads/tour-of-accounting-over-here-we-have-our-random-number-generator-nine-nine-are-you-sure-thats-random-thats-the-problem-with-randomness-you-can-never-be-sure-gilbert-comic.jpg)
pub struct SafeBytesSlice<'a> {
    slice: &'a mut [MaybeUninit<u8>],
    bytes_written: usize,
    // capacity exceeded. User tried put more bytes than available in slice
    cap_exceeded: bool,
}

impl<'a> From<&'a mut [u8]> for SafeBytesSlice<'a> {
    fn from(slice: &'a mut [u8]) -> Self {
        let maybe_uninit_slice =
            unsafe { &mut *(&mut *slice as *mut [u8] as *mut [mem::MaybeUninit<u8>]) };

        Self::from(maybe_uninit_slice)
    }
}

impl<'a> From<&'a mut [MaybeUninit<u8>]> for SafeBytesSlice<'a> {
    fn from(slice: &'a mut [MaybeUninit<u8>]) -> Self {
        Self {
            slice,
            bytes_written: 0,
            cap_exceeded: false,
        }
    }
}

impl<'a> SafeBytesSlice<'a> {
    /// Returns the number of bytes wrote into inner slice.
    ///
    /// Use [`BufMut::remaining_mut()`](bytes::buf::BufMut::remaining_mut) to check the number of bytes that
    /// can be written from the current position until the end of the buffer is reached.
    ///
    /// # Example
    ///
    /// This is example of BAD usage. Showing why this method is private and probably should never
    /// be public.
    ///
    /// [Tracking Issue 1](https://github.com/xoac/static-bytes/issues/1)
    ///
    /// ```compile_fail
    /// use static_bytes::SafeBytesSlice;
    /// use bytes::buf::BufMut;
    /// # fn fill_with_random(buf: &mut dyn BufMut, amount: usize) {
    /// #    for _ in 0..amount {
    /// #        buf.put_u8(0xFF);
    /// #    }
    /// # }
    /// let mut static_data = [0u8; 32];
    /// let mut safe_slice = SafeBytesSlice::from(&mut static_data[..]);
    ///
    /// fill_with_random(&mut safe_slice, 33);
    /// assert_eq!(safe_slice.is_exceed(), true);
    /// let output_len = safe_slice.bytes_written();
    /// // ⚠️ This  allow use  output_len in improper way!
    /// drop(safe_slice);
    ///
    /// // output is incorrect contains 32 bytes but should 33.
    /// let _output = &static_data[..output_len];
    /// ```
    #[allow(dead_code)]
    fn bytes_written(&self) -> usize {
        debug_assert_eq!(self.cap_exceeded, false);
        self.bytes_written
    }

    /// Returns true if the inner slice contains 0 bytes.
    #[inline]
    #[allow(dead_code)]
    fn is_empty(&self) -> bool {
        self.bytes_written() == 0
    }

    /// Returns filled bytes (`&[u8]`) or error if capacity exceeded.
    pub fn try_into_bytes(self) -> Result<&'a [u8], CapacityExceeded> {
        if self.is_exceed() {
            Err(CapacityExceeded {})
        } else {
            // TODO in the future there will be function for this
            // https://doc.rust-lang.org/std/mem/union.MaybeUninit.html#method.slice_get_ref
            //
            // Safty this is safe because memory is properly initialized up to self.bytes_written
            // MaybeUninit<T> is #[repr(transparent)]
            Ok(unsafe {
                &*(&self.slice[..self.bytes_written] as *const [core::mem::MaybeUninit<u8>]
                    as *const [u8])
            })
        }
    }

    /// Returns true if capacity was exceeded.
    ///
    /// `SafeBytesSlice` is not usable anymore - there is no access to inner bytes because they are
    /// in improper state.
    pub fn is_exceed(&self) -> bool {
        self.cap_exceeded
    }
}

// Implement required methods
unsafe impl<'a> BufMut for SafeBytesSlice<'a> {
    fn remaining_mut(&self) -> usize {
        debug_assert!(self.bytes_written <= self.slice.len());
        self.slice.len() - self.bytes_written
    }

    unsafe fn advance_mut(&mut self, cnt: usize) {
        let new_bytes_written = self.bytes_written + cnt;
        if new_bytes_written > self.slice.len() {
            self.bytes_written = self.slice.len(); // make `remaining_mut()` return 0
            self.cap_exceeded = true;
        } else {
            self.bytes_written = new_bytes_written;
        }
    }

    fn chunk_mut(&mut self) -> &mut UninitSlice {
        let bytes = &mut self.slice[self.bytes_written..];
        let len = bytes.len();
        let ptr = bytes.as_mut_ptr() as *mut _;
        unsafe { UninitSlice::from_raw_parts_mut(ptr, len) }
    }

    fn put_slice(&mut self, src: &[u8]) {
        use core::ptr;
        // check if we have enough data to put slice. If no set flag instead of panic!
        let src_len = src.len();
        if self.remaining_mut() < src_len {
            self.bytes_written = self.slice.len(); // make `remaining_mut()` return 0
            self.cap_exceeded = true;
            return;
        }

        // enough inner capacity to execute safe copy
        unsafe {
            let dst = self.chunk_mut();
            ptr::copy_nonoverlapping(src[..].as_ptr(), dst.as_mut_ptr() as *mut u8, src_len);
            self.advance_mut(src_len);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn fill_with_random(buf: &mut dyn BufMut, amount: usize) {
        for _ in 0..amount {
            buf.put_u8(0xFF);
        }
    }

    #[test]
    fn usefullness() {
        // standard way
        let mut data = [0u8; 32];
        let mut slice = &mut data[..];
        let slice_len = slice.len();
        fill_with_random(&mut slice, 27);
        // how many data was wrote (?)
        let n = slice_len - slice.len();
        assert_eq!(n, 27);
        let _wrote_data = &data[..n];

        // new way
        let mut raw = [0u8; 32];
        let mut slice = SafeBytesSlice::from(&mut raw[..]);
        fill_with_random(&mut slice, 27);
        // slice already contains bytes_written();
        let _wrote_data = match slice.try_into_bytes() {
            Ok(bytes) => bytes,
            Err(_err) => unimplemented!(),
        };
    }

    #[test]
    fn naive_test() {
        let mut static_data = [0u8; 32];
        let mut safe_slice = SafeBytesSlice::from(&mut static_data[..]);

        fill_with_random(&mut safe_slice, 32);
        assert_eq!(safe_slice.is_exceed(), false);

        for v in safe_slice
            .try_into_bytes()
            .expect("not expected capacity")
            .iter()
        {
            assert_eq!(*v, 0xFF);
        }

        // reuse data
        let mut safe_slice = SafeBytesSlice::from(&mut static_data[..]);
        fill_with_random(&mut safe_slice, 33);
        assert_eq!(safe_slice.is_exceed(), true);
    }
}