172 lines
3.6 KiB
Rust
172 lines
3.6 KiB
Rust
#[must_use]
|
|
pub struct ReadBuffer<'a> {
|
|
data: &'a [u8],
|
|
consumed: usize,
|
|
used: bool,
|
|
|
|
decrypted_consumed: &'a mut usize,
|
|
}
|
|
|
|
impl<'a> ReadBuffer<'a> {
|
|
#[inline]
|
|
pub(crate) fn new(buffer: &'a [u8], decrypted_consumed: &'a mut usize) -> Self {
|
|
Self {
|
|
data: buffer,
|
|
consumed: 0,
|
|
used: false,
|
|
decrypted_consumed,
|
|
}
|
|
}
|
|
|
|
#[inline]
|
|
#[must_use]
|
|
pub fn len(&self) -> usize {
|
|
self.data.len() - self.consumed
|
|
}
|
|
|
|
#[inline]
|
|
#[must_use]
|
|
pub fn is_empty(&self) -> bool {
|
|
self.len() == 0
|
|
}
|
|
|
|
#[inline]
|
|
pub fn peek(&mut self, count: usize) -> &'a [u8] {
|
|
let count = self.len().min(count);
|
|
let start = self.consumed;
|
|
|
|
self.used = true;
|
|
|
|
&self.data[start..start + count]
|
|
}
|
|
|
|
#[inline]
|
|
pub fn peek_all(&mut self) -> &'a [u8] {
|
|
self.peek(self.len())
|
|
}
|
|
|
|
#[inline]
|
|
pub fn pop(&mut self, count: usize) -> &'a [u8] {
|
|
let count = self.len().min(count);
|
|
let start = self.consumed;
|
|
self.consumed += count;
|
|
self.used = true;
|
|
|
|
&self.data[start..start + count]
|
|
}
|
|
|
|
#[inline]
|
|
pub fn pop_all(&mut self) -> &'a [u8] {
|
|
self.pop(self.len())
|
|
}
|
|
|
|
#[inline]
|
|
pub fn revert(self) {
|
|
core::mem::forget(self);
|
|
}
|
|
|
|
#[inline]
|
|
pub fn pop_into(&mut self, buf: &mut [u8]) -> usize {
|
|
let to_copy = self.pop(buf.len());
|
|
|
|
buf[..to_copy.len()].copy_from_slice(to_copy);
|
|
|
|
to_copy.len()
|
|
}
|
|
}
|
|
|
|
impl Drop for ReadBuffer<'_> {
|
|
#[inline]
|
|
fn drop(&mut self) {
|
|
*self.decrypted_consumed += if self.used {
|
|
self.consumed
|
|
} else {
|
|
self.data.len()
|
|
};
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod test {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn dropping_unused_buffer_consumes_all() {
|
|
let mut consumed = 1000;
|
|
let buffer = [0, 1, 2, 3];
|
|
|
|
_ = ReadBuffer::new(&buffer, &mut consumed);
|
|
|
|
assert_eq!(consumed, 1004);
|
|
}
|
|
|
|
#[test]
|
|
fn pop_moves_internal_cursor() {
|
|
let mut consumed = 0;
|
|
|
|
let mut buffer = ReadBuffer::new(&[0, 1, 2, 3], &mut consumed);
|
|
|
|
assert_eq!(buffer.pop(1), &[0]);
|
|
assert_eq!(buffer.pop(1), &[1]);
|
|
assert_eq!(buffer.pop(1), &[2]);
|
|
}
|
|
|
|
#[test]
|
|
fn dropping_consumes_as_many_bytes_as_used() {
|
|
let mut consumed = 0;
|
|
|
|
let mut buffer = ReadBuffer::new(&[0, 1, 2, 3], &mut consumed);
|
|
|
|
assert_eq!(buffer.pop(1), &[0]);
|
|
assert_eq!(buffer.pop(1), &[1]);
|
|
assert_eq!(buffer.pop(1), &[2]);
|
|
|
|
core::mem::drop(buffer);
|
|
|
|
assert_eq!(consumed, 3);
|
|
}
|
|
|
|
#[test]
|
|
fn pop_returns_fewer_bytes_if_requested_more_than_what_it_has() {
|
|
let mut consumed = 0;
|
|
|
|
let mut buffer = ReadBuffer::new(&[0, 1, 2, 3], &mut consumed);
|
|
|
|
assert_eq!(buffer.pop(1), &[0]);
|
|
assert_eq!(buffer.pop(1), &[1]);
|
|
assert_eq!(buffer.pop(4), &[2, 3]);
|
|
assert_eq!(buffer.pop(1), &[]);
|
|
|
|
core::mem::drop(buffer);
|
|
|
|
assert_eq!(consumed, 4);
|
|
}
|
|
|
|
#[test]
|
|
fn peek_does_not_consume() {
|
|
let mut consumed = 0;
|
|
|
|
let mut buffer = ReadBuffer::new(&[0, 1, 2, 3], &mut consumed);
|
|
|
|
assert_eq!(buffer.peek(1), &[0]);
|
|
assert_eq!(buffer.peek(1), &[0]);
|
|
|
|
core::mem::drop(buffer);
|
|
|
|
assert_eq!(consumed, 0);
|
|
}
|
|
|
|
#[test]
|
|
fn revert_undoes_pop() {
|
|
let mut consumed = 0;
|
|
|
|
let mut buffer = ReadBuffer::new(&[0, 1, 2, 3], &mut consumed);
|
|
|
|
assert_eq!(buffer.pop(4), &[0, 1, 2, 3]);
|
|
|
|
buffer.revert();
|
|
|
|
assert_eq!(consumed, 0);
|
|
}
|
|
}
|