static_buffer/
lib.rs

1#![no_std]
2
3extern crate generic_array;
4extern crate typenum;
5
6use core::fmt;
7use core::cmp;
8
9use generic_array::{ArrayLength, GenericArray};
10use typenum::uint::Unsigned;
11use typenum::consts::{U64, U128};
12
13/// A FixedBuffer, likes its name implies, is a fixed size buffer. When the buffer becomes full, it
14/// must be processed. The input() method takes care of processing and then clearing the buffer
15/// automatically. However, other methods do not and require the caller to process the buffer. Any
16/// method that modifies the buffer directory or provides the caller with bytes that can be modifies
17/// results in those bytes being marked as used by the buffer.
18pub trait FixedBuf {
19    /// Input a vector of bytes. If the buffer becomes full, process it with the provided
20    /// function and then clear the buffer.
21    fn input<F: FnMut(&[u8])>(&mut self, input: &[u8], func: F);
22
23    /// Reset the buffer.
24    fn reset(&mut self);
25
26    /// Zero the buffer up until the specified index. The buffer position currently must not be
27    /// greater than that index.
28    fn zero_until(&mut self, idx: usize);
29
30    /// Get a slice of the buffer of the specified size. There must be at least that many bytes
31    /// remaining in the buffer.
32    fn next(&mut self, len: usize) -> &mut [u8];
33
34    /// Get the current buffer. The buffer must already be full. This clears the buffer as well.
35    fn full_buffer(&mut self) -> &mut [u8];
36
37    /// Get the current buffer.
38    fn current_buffer(&self) -> &[u8];
39
40    /// Get the current position of the buffer.
41    fn position(&self) -> usize;
42
43    /// Get the number of bytes remaining in the buffer until it is full.
44    fn remaining(&self) -> usize;
45
46    /// Get the size of the buffer
47    fn size() -> usize;
48}
49
50pub struct FixedBuffer<N: ArrayLength<u8>> {
51    buffer: GenericArray<u8, N>,
52    position: usize,
53}
54
55pub type FixedBuffer64 = FixedBuffer<U64>;
56pub type FixedBuffer128 = FixedBuffer<U128>;
57
58impl<N: ArrayLength<u8>> FixedBuffer<N> {
59    /// Create a new buffer
60    #[inline(always)]
61    pub fn new() -> Self {
62        FixedBuffer {
63            buffer: GenericArray::new(),
64            position: 0,
65        }
66    }
67
68    fn fill_up<'a>(&mut self, input: &'a [u8]) -> &'a [u8] {
69        let count = cmp::min(input.len(), self.remaining());
70
71        self.buffer[self.position..(self.position + count)].copy_from_slice(&input[..count]);
72        self.position += count;
73
74        &input[count..]
75    }
76}
77
78impl<N: ArrayLength<u8>> FixedBuf for FixedBuffer<N> {
79    #[inline(always)]
80    fn input<F: FnMut(&[u8])>(&mut self, input: &[u8], mut func: F) {
81        let rest = self.fill_up(input);
82        if rest.is_empty() { return }
83
84        for chunk in rest.chunks(Self::size()) {
85            func(&self.buffer);
86            self.reset();
87
88            self.fill_up(chunk);
89        }
90    }
91
92    #[inline(always)]
93    fn reset(&mut self) {
94        self.position = 0;
95    }
96
97    #[inline(always)]
98    fn zero_until(&mut self, idx: usize) {
99        assert!(idx >= self.position);
100        for b in &mut self.buffer[self.position..idx] {
101            *b = 0
102        }
103        self.position = idx;
104    }
105
106    #[inline(always)]
107    fn next(&mut self, len: usize) -> &mut [u8] {
108        assert!(self.position + len <= Self::size());
109        self.position += len;
110        &mut self.buffer[self.position - len..self.position]
111    }
112
113    fn full_buffer(&mut self) -> &mut [u8] {
114        assert!(self.position == Self::size());
115        self.position = 0;
116        &mut self.buffer[..]
117    }
118
119    #[inline(always)]
120    fn current_buffer(&self) -> &[u8] {
121        &self.buffer[..self.position]
122    }
123
124    #[inline(always)]
125    fn position(&self) -> usize {
126        self.position
127    }
128
129    #[inline(always)]
130    fn remaining(&self) -> usize {
131        Self::size() - self.position
132    }
133
134    #[inline(always)]
135    fn size() -> usize {
136        <N as Unsigned>::to_usize()
137    }
138}
139
140impl<N: ArrayLength<u8>> Clone for FixedBuffer<N> {
141    fn clone(&self) -> Self {
142        FixedBuffer {
143            buffer: GenericArray::from_slice(&self.buffer),
144            position: self.position,
145        }
146    }
147}
148
149impl<N: ArrayLength<u8>> fmt::Debug for FixedBuffer<N> {
150    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
151        write!(f, "{:?}", &self.buffer[..self.position])
152    }
153}
154
155/// The StandardPadding trait adds a method useful for various hash algorithms to a FixedBuffer
156/// struct.
157pub trait StandardPadding {
158    /// Add standard padding to the buffer. The buffer must not be full when this method is called
159    /// and is guaranteed to have exactly rem remaining bytes when it returns. If there are not at
160    /// least rem bytes available, the buffer will be zero padded, processed, cleared, and then
161    /// filled with zeros again until only rem bytes are remaining.
162    fn pad<F: FnMut(&[u8])>(&mut self, padding: u8, rem: usize, func: F);
163
164    #[inline(always)]
165    fn standard_padding<F: FnMut(&[u8])>(&mut self, rem: usize, func: F) {
166        self.pad(0b10000000, rem, func)
167    }
168}
169
170impl<T: FixedBuf> StandardPadding for T {
171    #[inline(always)]
172    fn pad<F: FnMut(&[u8])>(&mut self, padding: u8, rem: usize, mut func: F) {
173        let size = T::size();
174
175        self.input(&[padding], |args| func(args));
176
177        if self.remaining() < rem {
178            self.zero_until(size);
179            func(self.full_buffer());
180        }
181
182        self.zero_until(size - rem);
183    }
184}
185
186#[cfg(test)]
187mod tests {
188    use super::*;
189    use typenum::consts::U2;
190
191    #[test]
192    fn doesnt_fire_function_when_buffer_is_not_filled() {
193        let mut buf = FixedBuffer::<U2>::new();
194
195        buf.input(&[1], |_| assert!(false, "This should not happen."));
196    }
197
198    #[test]
199    fn do_fire_function_when_buffer_is_filled() {
200        let mut buf = FixedBuffer::<U2>::new();
201
202        buf.input(&[1, 2], |_| assert!(false, "This should not happen."));
203    }
204
205    #[test]
206    fn fire_function_when_buffer_is_filled_and_there_is_still_data() {
207        let mut buf = FixedBuffer::<U2>::new();
208        let mut i = false;
209
210        buf.input(&[1, 2, 3], |_| i = true);
211
212        assert!(i, "Passed method should be called");
213    }
214}