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(string: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}
173