1 | use crate::{ |
2 | encode::add_padding, |
3 | engine::{Config, Engine}, |
4 | }; |
5 | #[cfg (any(feature = "alloc" , test))] |
6 | use alloc::string::String; |
7 | #[cfg (any(feature = "alloc" , test))] |
8 | use core::str; |
9 | |
10 | /// The output mechanism for ChunkedEncoder's encoded bytes. |
11 | pub 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. |
19 | pub struct ChunkedEncoder<'e, E: Engine + ?Sized> { |
20 | engine: &'e E, |
21 | } |
22 | |
23 | impl<'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))] |
51 | pub(crate) struct StringSink<'a> { |
52 | string: &'a mut String, |
53 | } |
54 | |
55 | #[cfg (any(feature = "alloc" , test))] |
56 | impl<'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))] |
63 | impl<'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)] |
74 | pub 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 | |