1#[cfg(any(feature = "alloc", feature = "std", test))]
2use alloc::string::String;
3use core::cmp;
4#[cfg(any(feature = "alloc", feature = "std", test))]
5use core::str;
6
7use crate::encode::add_padding;
8use crate::engine::{Config, Engine};
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
18const BUF_SIZE: usize = 1024;
19
20/// A base64 encoder that emits encoded bytes in chunks without heap allocation.
21pub struct ChunkedEncoder<'e, E: Engine + ?Sized> {
22 engine: &'e E,
23 max_input_chunk_len: usize,
24}
25
26impl<'e, E: Engine + ?Sized> ChunkedEncoder<'e, E> {
27 pub fn new(engine: &'e E) -> ChunkedEncoder<'e, E> {
28 ChunkedEncoder {
29 engine,
30 max_input_chunk_len: max_input_length(BUF_SIZE, engine.config().encode_padding()),
31 }
32 }
33
34 pub fn encode<S: Sink>(&self, bytes: &[u8], sink: &mut S) -> Result<(), S::Error> {
35 let mut encode_buf: [u8; BUF_SIZE] = [0; BUF_SIZE];
36 let mut input_index = 0;
37
38 while input_index < bytes.len() {
39 // either the full input chunk size, or it's the last iteration
40 let input_chunk_len = cmp::min(self.max_input_chunk_len, bytes.len() - input_index);
41
42 let chunk = &bytes[input_index..(input_index + input_chunk_len)];
43
44 let mut b64_bytes_written = self.engine.internal_encode(chunk, &mut encode_buf);
45
46 input_index += input_chunk_len;
47 let more_input_left = input_index < bytes.len();
48
49 if self.engine.config().encode_padding() && !more_input_left {
50 // no more input, add padding if needed. Buffer will have room because
51 // max_input_length leaves room for it.
52 b64_bytes_written +=
53 add_padding(b64_bytes_written, &mut encode_buf[b64_bytes_written..]);
54 }
55
56 sink.write_encoded_bytes(&encode_buf[0..b64_bytes_written])?;
57 }
58
59 Ok(())
60 }
61}
62
63/// Calculate the longest input that can be encoded for the given output buffer size.
64///
65/// If the config requires padding, two bytes of buffer space will be set aside so that the last
66/// chunk of input can be encoded safely.
67///
68/// The input length will always be a multiple of 3 so that no encoding state has to be carried over
69/// between chunks.
70fn max_input_length(encoded_buf_len: usize, padded: bool) -> usize {
71 let effective_buf_len: usize = if padded {
72 // make room for padding
73 encoded_buf_len
74 .checked_sub(2)
75 .expect(msg:"Don't use a tiny buffer")
76 } else {
77 encoded_buf_len
78 };
79
80 // No padding, so just normal base64 expansion.
81 (effective_buf_len / 4) * 3
82}
83
84// A really simple sink that just appends to a string
85#[cfg(any(feature = "alloc", feature = "std", test))]
86pub(crate) struct StringSink<'a> {
87 string: &'a mut String,
88}
89
90#[cfg(any(feature = "alloc", feature = "std", test))]
91impl<'a> StringSink<'a> {
92 pub(crate) fn new(s: &mut String) -> StringSink {
93 StringSink { string: s }
94 }
95}
96
97#[cfg(any(feature = "alloc", feature = "std", test))]
98impl<'a> Sink for StringSink<'a> {
99 type Error = ();
100
101 fn write_encoded_bytes(&mut self, s: &[u8]) -> Result<(), Self::Error> {
102 self.string.push_str(string:str::from_utf8(s).unwrap());
103
104 Ok(())
105 }
106}
107
108#[cfg(test)]
109pub mod tests {
110 use rand::{
111 distributions::{Distribution, Uniform},
112 Rng, SeedableRng,
113 };
114
115 use crate::{
116 alphabet::STANDARD,
117 engine::general_purpose::{GeneralPurpose, GeneralPurposeConfig, PAD},
118 tests::random_engine,
119 };
120
121 use super::*;
122
123 #[test]
124 fn chunked_encode_empty() {
125 assert_eq!("", chunked_encode_str(&[], PAD));
126 }
127
128 #[test]
129 fn chunked_encode_intermediate_fast_loop() {
130 // > 8 bytes input, will enter the pretty fast loop
131 assert_eq!("Zm9vYmFyYmF6cXV4", chunked_encode_str(b"foobarbazqux", PAD));
132 }
133
134 #[test]
135 fn chunked_encode_fast_loop() {
136 // > 32 bytes input, will enter the uber fast loop
137 assert_eq!(
138 "Zm9vYmFyYmF6cXV4cXV1eGNvcmdlZ3JhdWx0Z2FycGx5eg==",
139 chunked_encode_str(b"foobarbazquxquuxcorgegraultgarplyz", PAD)
140 );
141 }
142
143 #[test]
144 fn chunked_encode_slow_loop_only() {
145 // < 8 bytes input, slow loop only
146 assert_eq!("Zm9vYmFy", chunked_encode_str(b"foobar", PAD));
147 }
148
149 #[test]
150 fn chunked_encode_matches_normal_encode_random_string_sink() {
151 let helper = StringSinkTestHelper;
152 chunked_encode_matches_normal_encode_random(&helper);
153 }
154
155 #[test]
156 fn max_input_length_no_pad() {
157 assert_eq!(768, max_input_length(1024, false));
158 }
159
160 #[test]
161 fn max_input_length_with_pad_decrements_one_triple() {
162 assert_eq!(765, max_input_length(1024, true));
163 }
164
165 #[test]
166 fn max_input_length_with_pad_one_byte_short() {
167 assert_eq!(765, max_input_length(1025, true));
168 }
169
170 #[test]
171 fn max_input_length_with_pad_fits_exactly() {
172 assert_eq!(768, max_input_length(1026, true));
173 }
174
175 #[test]
176 fn max_input_length_cant_use_extra_single_encoded_byte() {
177 assert_eq!(300, max_input_length(401, false));
178 }
179
180 pub fn chunked_encode_matches_normal_encode_random<S: SinkTestHelper>(sink_test_helper: &S) {
181 let mut input_buf: Vec<u8> = Vec::new();
182 let mut output_buf = String::new();
183 let mut rng = rand::rngs::SmallRng::from_entropy();
184 let input_len_range = Uniform::new(1, 10_000);
185
186 for _ in 0..5_000 {
187 input_buf.clear();
188 output_buf.clear();
189
190 let buf_len = input_len_range.sample(&mut rng);
191 for _ in 0..buf_len {
192 input_buf.push(rng.gen());
193 }
194
195 let engine = random_engine(&mut rng);
196
197 let chunk_encoded_string = sink_test_helper.encode_to_string(&engine, &input_buf);
198 engine.encode_string(&input_buf, &mut output_buf);
199
200 assert_eq!(output_buf, chunk_encoded_string, "input len={}", buf_len);
201 }
202 }
203
204 fn chunked_encode_str(bytes: &[u8], config: GeneralPurposeConfig) -> String {
205 let mut s = String::new();
206
207 let mut sink = StringSink::new(&mut s);
208 let engine = GeneralPurpose::new(&STANDARD, config);
209 let encoder = ChunkedEncoder::new(&engine);
210 encoder.encode(bytes, &mut sink).unwrap();
211
212 s
213 }
214
215 // An abstraction around sinks so that we can have tests that easily to any sink implementation
216 pub trait SinkTestHelper {
217 fn encode_to_string<E: Engine>(&self, engine: &E, bytes: &[u8]) -> String;
218 }
219
220 struct StringSinkTestHelper;
221
222 impl SinkTestHelper for StringSinkTestHelper {
223 fn encode_to_string<E: Engine>(&self, engine: &E, bytes: &[u8]) -> String {
224 let encoder = ChunkedEncoder::new(engine);
225 let mut s = String::new();
226 let mut sink = StringSink::new(&mut s);
227 encoder.encode(bytes, &mut sink).unwrap();
228
229 s
230 }
231 }
232}
233