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, self.engine.padding(), &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 crate::{
76        alphabet::STANDARD,
77        engine::general_purpose::{GeneralPurpose, GeneralPurposeConfig, PAD},
78        tests::random_engine,
79    };
80    use rand::distr::{Distribution, Uniform};
81    use rand::{rngs, RngExt};
82
83    use super::*;
84
85    #[test]
86    fn chunked_encode_empty() {
87        assert_eq!("", chunked_encode_str(&[], PAD));
88    }
89
90    #[test]
91    fn chunked_encode_intermediate_fast_loop() {
92        // > 8 bytes input, will enter the pretty fast loop
93        assert_eq!("Zm9vYmFyYmF6cXV4", chunked_encode_str(b"foobarbazqux", PAD));
94    }
95
96    #[test]
97    fn chunked_encode_fast_loop() {
98        // > 32 bytes input, will enter the uber fast loop
99        assert_eq!(
100            "Zm9vYmFyYmF6cXV4cXV1eGNvcmdlZ3JhdWx0Z2FycGx5eg==",
101            chunked_encode_str(b"foobarbazquxquuxcorgegraultgarplyz", PAD)
102        );
103    }
104
105    #[test]
106    fn chunked_encode_slow_loop_only() {
107        // < 8 bytes input, slow loop only
108        assert_eq!("Zm9vYmFy", chunked_encode_str(b"foobar", PAD));
109    }
110
111    #[test]
112    fn chunked_encode_matches_normal_encode_random_string_sink() {
113        let helper = StringSinkTestHelper;
114        chunked_encode_matches_normal_encode_random(&helper);
115    }
116
117    pub fn chunked_encode_matches_normal_encode_random<S: SinkTestHelper>(sink_test_helper: &S) {
118        let mut input_buf: Vec<u8> = Vec::new();
119        let mut output_buf = String::new();
120        let mut rng = rand::make_rng::<rngs::SmallRng>();
121        let input_len_range = Uniform::new(1, 10_000).unwrap();
122
123        for _ in 0..20_000 {
124            input_buf.clear();
125            output_buf.clear();
126
127            let buf_len = input_len_range.sample(&mut rng);
128            for _ in 0..buf_len {
129                input_buf.push(rng.random());
130            }
131
132            let engine = random_engine(&mut rng);
133
134            let chunk_encoded_string = sink_test_helper.encode_to_string(&engine, &input_buf);
135            engine.encode_string(&input_buf, &mut output_buf);
136
137            assert_eq!(output_buf, chunk_encoded_string, "input len={}", buf_len);
138        }
139    }
140
141    fn chunked_encode_str(bytes: &[u8], config: GeneralPurposeConfig) -> String {
142        let mut s = String::new();
143
144        let mut sink = StringSink::new(&mut s);
145        let engine = GeneralPurpose::new(&STANDARD, config);
146        let encoder = ChunkedEncoder::new(&engine);
147        encoder.encode(bytes, &mut sink).unwrap();
148
149        s
150    }
151
152    // An abstraction around sinks so that we can have tests that easily to any sink implementation
153    pub trait SinkTestHelper {
154        fn encode_to_string<E: Engine>(&self, engine: &E, bytes: &[u8]) -> String;
155    }
156
157    struct StringSinkTestHelper;
158
159    impl SinkTestHelper for StringSinkTestHelper {
160        fn encode_to_string<E: Engine>(&self, engine: &E, bytes: &[u8]) -> String {
161            let encoder = ChunkedEncoder::new(engine);
162            let mut s = String::new();
163            let mut sink = StringSink::new(&mut s);
164            encoder.encode(bytes, &mut sink).unwrap();
165
166            s
167        }
168    }
169}
base64/chunked_encoder.rs

base64/
chunked_encoder.rs
