1 | use super::encoder::EncoderWriter; |
2 | use crate::engine::Engine; |
3 | use std::io; |
4 | |
5 | /// A `Write` implementation that base64-encodes data using the provided config and accumulates the |
6 | /// resulting base64 utf8 `&str` in a [StrConsumer] implementation (typically `String`), which is |
7 | /// then exposed via `into_inner()`. |
8 | /// |
9 | /// # Examples |
10 | /// |
11 | /// Buffer base64 in a new String: |
12 | /// |
13 | /// ``` |
14 | /// use std::io::Write; |
15 | /// use base64::engine::general_purpose; |
16 | /// |
17 | /// let mut enc = base64::write::EncoderStringWriter::new(&general_purpose::STANDARD); |
18 | /// |
19 | /// enc.write_all(b"asdf" ).unwrap(); |
20 | /// |
21 | /// // get the resulting String |
22 | /// let b64_string = enc.into_inner(); |
23 | /// |
24 | /// assert_eq!("YXNkZg==" , &b64_string); |
25 | /// ``` |
26 | /// |
27 | /// Or, append to an existing `String`, which implements `StrConsumer`: |
28 | /// |
29 | /// ``` |
30 | /// use std::io::Write; |
31 | /// use base64::engine::general_purpose; |
32 | /// |
33 | /// let mut buf = String::from("base64: " ); |
34 | /// |
35 | /// let mut enc = base64::write::EncoderStringWriter::from_consumer( |
36 | /// &mut buf, |
37 | /// &general_purpose::STANDARD); |
38 | /// |
39 | /// enc.write_all(b"asdf" ).unwrap(); |
40 | /// |
41 | /// // release the &mut reference on buf |
42 | /// let _ = enc.into_inner(); |
43 | /// |
44 | /// assert_eq!("base64: YXNkZg==" , &buf); |
45 | /// ``` |
46 | /// |
47 | /// # Panics |
48 | /// |
49 | /// Calling `write()` (or related methods) or `finish()` after `finish()` has completed without |
50 | /// error is invalid and will panic. |
51 | /// |
52 | /// # Performance |
53 | /// |
54 | /// Because it has to validate that the base64 is UTF-8, it is about 80% as fast as writing plain |
55 | /// bytes to a `io::Write`. |
56 | pub struct EncoderStringWriter<'e, E: Engine, S: StrConsumer> { |
57 | encoder: EncoderWriter<'e, E, Utf8SingleCodeUnitWriter<S>>, |
58 | } |
59 | |
60 | impl<'e, E: Engine, S: StrConsumer> EncoderStringWriter<'e, E, S> { |
61 | /// Create a EncoderStringWriter that will append to the provided `StrConsumer`. |
62 | pub fn from_consumer(str_consumer: S, engine: &'e E) -> Self { |
63 | EncoderStringWriter { |
64 | encoder: EncoderWriter::new(delegate:Utf8SingleCodeUnitWriter { str_consumer }, engine), |
65 | } |
66 | } |
67 | |
68 | /// Encode all remaining buffered data, including any trailing incomplete input triples and |
69 | /// associated padding. |
70 | /// |
71 | /// Returns the base64-encoded form of the accumulated written data. |
72 | pub fn into_inner(mut self) -> S { |
73 | self.encoder |
74 | .finish() |
75 | .expect(msg:"Writing to a consumer should never fail" ) |
76 | .str_consumer |
77 | } |
78 | } |
79 | |
80 | impl<'e, E: Engine> EncoderStringWriter<'e, E, String> { |
81 | /// Create a EncoderStringWriter that will encode into a new `String` with the provided config. |
82 | pub fn new(engine: &'e E) -> Self { |
83 | EncoderStringWriter::from_consumer(str_consumer:String::new(), engine) |
84 | } |
85 | } |
86 | |
87 | impl<'e, E: Engine, S: StrConsumer> io::Write for EncoderStringWriter<'e, E, S> { |
88 | fn write(&mut self, buf: &[u8]) -> io::Result<usize> { |
89 | self.encoder.write(buf) |
90 | } |
91 | |
92 | fn flush(&mut self) -> io::Result<()> { |
93 | self.encoder.flush() |
94 | } |
95 | } |
96 | |
97 | /// An abstraction around consuming `str`s produced by base64 encoding. |
98 | pub trait StrConsumer { |
99 | /// Consume the base64 encoded data in `buf` |
100 | fn consume(&mut self, buf: &str); |
101 | } |
102 | |
103 | /// As for io::Write, `StrConsumer` is implemented automatically for `&mut S`. |
104 | impl<S: StrConsumer + ?Sized> StrConsumer for &mut S { |
105 | fn consume(&mut self, buf: &str) { |
106 | (**self).consume(buf); |
107 | } |
108 | } |
109 | |
110 | /// Pushes the str onto the end of the String |
111 | impl StrConsumer for String { |
112 | fn consume(&mut self, buf: &str) { |
113 | self.push_str(string:buf); |
114 | } |
115 | } |
116 | |
117 | /// A `Write` that only can handle bytes that are valid single-byte UTF-8 code units. |
118 | /// |
119 | /// This is safe because we only use it when writing base64, which is always valid UTF-8. |
120 | struct Utf8SingleCodeUnitWriter<S: StrConsumer> { |
121 | str_consumer: S, |
122 | } |
123 | |
124 | impl<S: StrConsumer> io::Write for Utf8SingleCodeUnitWriter<S> { |
125 | fn write(&mut self, buf: &[u8]) -> io::Result<usize> { |
126 | // Because we expect all input to be valid utf-8 individual bytes, we can encode any buffer |
127 | // length |
128 | let s: &str = std::str::from_utf8(buf).expect(msg:"Input must be valid UTF-8" ); |
129 | |
130 | self.str_consumer.consume(buf:s); |
131 | |
132 | Ok(buf.len()) |
133 | } |
134 | |
135 | fn flush(&mut self) -> io::Result<()> { |
136 | // no op |
137 | Ok(()) |
138 | } |
139 | } |
140 | |
141 | #[cfg (test)] |
142 | mod tests { |
143 | use crate::{ |
144 | engine::Engine, tests::random_engine, write::encoder_string_writer::EncoderStringWriter, |
145 | }; |
146 | use rand::Rng; |
147 | use std::io::Write; |
148 | |
149 | #[test ] |
150 | fn every_possible_split_of_input() { |
151 | let mut rng = rand::thread_rng(); |
152 | let mut orig_data = Vec::<u8>::new(); |
153 | let mut normal_encoded = String::new(); |
154 | |
155 | let size = 5_000; |
156 | |
157 | for i in 0..size { |
158 | orig_data.clear(); |
159 | normal_encoded.clear(); |
160 | |
161 | for _ in 0..size { |
162 | orig_data.push(rng.gen()); |
163 | } |
164 | |
165 | let engine = random_engine(&mut rng); |
166 | engine.encode_string(&orig_data, &mut normal_encoded); |
167 | |
168 | let mut stream_encoder = EncoderStringWriter::new(&engine); |
169 | // Write the first i bytes, then the rest |
170 | stream_encoder.write_all(&orig_data[0..i]).unwrap(); |
171 | stream_encoder.write_all(&orig_data[i..]).unwrap(); |
172 | |
173 | let stream_encoded = stream_encoder.into_inner(); |
174 | |
175 | assert_eq!(normal_encoded, stream_encoded); |
176 | } |
177 | } |
178 | } |
179 | |