base64/
chunked_encoder.rs

1use crate::{
2    encode::add_padding,
3    engine::{Config, Engine},
4};
5#[cfg(any(feature = "alloc", test))]
6use alloc::string::String;
7#[cfg(any(feature = "alloc", test))]
8use core::str;
9
10/// The output mechanism for ChunkedEncoder's encoded bytes.
11pub trait Sink {
12    type Error;
13
14    /// Handle a chunk of encoded base64 data (as UTF-8 bytes)
15    fn write_encoded_bytes(&mut self, encoded: &[u8]) -> Result<(), Self::Error>;
16}
17
18/// A base64 encoder that emits encoded bytes in chunks without heap allocation.
19pub struct ChunkedEncoder<'e, E: Engine + ?Sized> {
20    engine: &'e E,
21}
22
23impl<'e, E: Engine + ?Sized> ChunkedEncoder<'e, E> {
24    pub fn new(engine: &'e E) -> ChunkedEncoder<'e, E> {
25        ChunkedEncoder { engine }
26    }
27
28    pub fn encode<S: Sink>(&self, bytes: &[u8], sink: &mut S) -> Result<(), S::Error> {
29        const BUF_SIZE: usize = 1024;
30        const CHUNK_SIZE: usize = BUF_SIZE / 4 * 3;
31
32        let mut buf = [0; BUF_SIZE];
33        for chunk in bytes.chunks(CHUNK_SIZE) {
34            let mut len = self.engine.internal_encode(chunk, &mut buf);
35            if chunk.len() != CHUNK_SIZE && self.engine.config().encode_padding() {
36                // Final, potentially partial, chunk.
37                // Only need to consider if padding is needed on a partial chunk since full chunk
38                // is a multiple of 3, which therefore won't be padded.
39                // Pad output to multiple of four bytes if required by config.
40                len += add_padding(len, &mut buf[len..]);
41            }
42            sink.write_encoded_bytes(&buf[..len])?;
43        }
44
45        Ok(())
46    }
47}
48
49// A really simple sink that just appends to a string
50#[cfg(any(feature = "alloc", test))]
51pub(crate) struct StringSink<'a> {
52    string: &'a mut String,
53}
54
55#[cfg(any(feature = "alloc", test))]
56impl<'a> StringSink<'a> {
57    pub(crate) fn new(s: &mut String) -> StringSink {
58        StringSink { string: s }
59    }
60}
61
62#[cfg(any(feature = "alloc", test))]
63impl<'a> Sink for StringSink<'a> {
64    type Error = ();
65
66    fn write_encoded_bytes(&mut self, s: &[u8]) -> Result<(), Self::Error> {
67        self.string.push_str(str::from_utf8(s).unwrap());
68
69        Ok(())
70    }
71}
72
73#[cfg(test)]
74pub mod tests {
75    use rand::{
76        distributions::{Distribution, Uniform},
77        Rng, SeedableRng,
78    };
79
80    use crate::{
81        alphabet::STANDARD,
82        engine::general_purpose::{GeneralPurpose, GeneralPurposeConfig, PAD},
83        tests::random_engine,
84    };
85
86    use super::*;
87
88    #[test]
89    fn chunked_encode_empty() {
90        assert_eq!("", chunked_encode_str(&[], PAD));
91    }
92
93    #[test]
94    fn chunked_encode_intermediate_fast_loop() {
95        // > 8 bytes input, will enter the pretty fast loop
96        assert_eq!("Zm9vYmFyYmF6cXV4", chunked_encode_str(b"foobarbazqux", PAD));
97    }
98
99    #[test]
100    fn chunked_encode_fast_loop() {
101        // > 32 bytes input, will enter the uber fast loop
102        assert_eq!(
103            "Zm9vYmFyYmF6cXV4cXV1eGNvcmdlZ3JhdWx0Z2FycGx5eg==",
104            chunked_encode_str(b"foobarbazquxquuxcorgegraultgarplyz", PAD)
105        );
106    }
107
108    #[test]
109    fn chunked_encode_slow_loop_only() {
110        // < 8 bytes input, slow loop only
111        assert_eq!("Zm9vYmFy", chunked_encode_str(b"foobar", PAD));
112    }
113
114    #[test]
115    fn chunked_encode_matches_normal_encode_random_string_sink() {
116        let helper = StringSinkTestHelper;
117        chunked_encode_matches_normal_encode_random(&helper);
118    }
119
120    pub fn chunked_encode_matches_normal_encode_random<S: SinkTestHelper>(sink_test_helper: &S) {
121        let mut input_buf: Vec<u8> = Vec::new();
122        let mut output_buf = String::new();
123        let mut rng = rand::rngs::SmallRng::from_entropy();
124        let input_len_range = Uniform::new(1, 10_000);
125
126        for _ in 0..20_000 {
127            input_buf.clear();
128            output_buf.clear();
129
130            let buf_len = input_len_range.sample(&mut rng);
131            for _ in 0..buf_len {
132                input_buf.push(rng.gen());
133            }
134
135            let engine = random_engine(&mut rng);
136
137            let chunk_encoded_string = sink_test_helper.encode_to_string(&engine, &input_buf);
138            engine.encode_string(&input_buf, &mut output_buf);
139
140            assert_eq!(output_buf, chunk_encoded_string, "input len={}", buf_len);
141        }
142    }
143
144    fn chunked_encode_str(bytes: &[u8], config: GeneralPurposeConfig) -> String {
145        let mut s = String::new();
146
147        let mut sink = StringSink::new(&mut s);
148        let engine = GeneralPurpose::new(&STANDARD, config);
149        let encoder = ChunkedEncoder::new(&engine);
150        encoder.encode(bytes, &mut sink).unwrap();
151
152        s
153    }
154
155    // An abstraction around sinks so that we can have tests that easily to any sink implementation
156    pub trait SinkTestHelper {
157        fn encode_to_string<E: Engine>(&self, engine: &E, bytes: &[u8]) -> String;
158    }
159
160    struct StringSinkTestHelper;
161
162    impl SinkTestHelper for StringSinkTestHelper {
163        fn encode_to_string<E: Engine>(&self, engine: &E, bytes: &[u8]) -> String {
164            let encoder = ChunkedEncoder::new(engine);
165            let mut s = String::new();
166            let mut sink = StringSink::new(&mut s);
167            encoder.encode(bytes, &mut sink).unwrap();
168
169            s
170        }
171    }
172}
base64/chunked_encoder.rs

base64/
chunked_encoder.rs
