vendor/base64/src/write/encoder_string_writer.rs

   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(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("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(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(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 = std::str::from_utf8(buf).expect("Input must be valid UTF-8");
 129
 130         self.str_consumer.consume(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 }