base64/write/
encoder_string_writer.rs

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