1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
use encode::encode_to_slice;
use std::io::{ErrorKind, Result, Write};
use std::{cmp, fmt};
use {encode_config_slice, Config};

pub(crate) const BUF_SIZE: usize = 1024;
/// The most bytes whose encoding will fit in `BUF_SIZE`
const MAX_INPUT_LEN: usize = BUF_SIZE / 4 * 3;
// 3 bytes of input = 4 bytes of base64, always (because we don't allow line wrapping)
const MIN_ENCODE_CHUNK_SIZE: usize = 3;

/// A `Write` implementation that base64 encodes data before delegating to the wrapped writer.
///
/// Because base64 has special handling for the end of the input data (padding, etc), there's a
/// `finish()` method on this type that encodes any leftover input bytes and adds padding if
/// appropriate. It's called automatically when deallocated (see the `Drop` implementation), but
/// any error that occurs when invoking the underlying writer will be suppressed. If you want to
/// handle such errors, call `finish()` yourself.
///
/// # Examples
///
/// ```
/// use std::io::Write;
///
/// // use a vec as the simplest possible `Write` -- in real code this is probably a file, etc.
/// let mut wrapped_writer = Vec::new();
/// {
///     let mut enc = base64::write::EncoderWriter::new(
///         &mut wrapped_writer, base64::STANDARD);
///
///     // handle errors as you normally would
///     enc.write_all(b"asdf").unwrap();
///     // could leave this out to be called by Drop, if you don't care
///     // about handling errors
///     enc.finish().unwrap();
///
/// }
///
/// // base64 was written to the writer
/// assert_eq!(b"YXNkZg==", &wrapped_writer[..]);
///
/// ```
///
/// # Panics
///
/// Calling `write()` after `finish()` is invalid and will panic.
///
/// # Errors
///
/// Base64 encoding itself does not generate errors, but errors from the wrapped writer will be
/// returned as per the contract of `Write`.
///
/// # Performance
///
/// It has some minor performance loss compared to encoding slices (a couple percent).
/// It does not do any heap allocation.
pub struct EncoderWriter<'a, W: 'a + Write> {
    config: Config,
    /// Where encoded data is written to
    w: &'a mut W,
    /// Holds a partial chunk, if any, after the last `write()`, so that we may then fill the chunk
    /// with the next `write()`, encode it, then proceed with the rest of the input normally.
    extra_input: [u8; MIN_ENCODE_CHUNK_SIZE],
    /// How much of `extra` is occupied, in `[0, MIN_ENCODE_CHUNK_SIZE]`.
    extra_input_occupied_len: usize,
    /// Buffer to encode into. May hold leftover encoded bytes from a previous write call that the underlying writer
    /// did not write last time.
    output: [u8; BUF_SIZE],
    /// How much of `output` is occupied with encoded data that couldn't be written last time
    output_occupied_len: usize,
    /// True iff padding / partial last chunk has been written.
    finished: bool,
    /// panic safety: don't write again in destructor if writer panicked while we were writing to it
    panicked: bool,
}

impl<'a, W: Write> fmt::Debug for EncoderWriter<'a, W> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "extra_input: {:?} extra_input_occupied_len:{:?} output[..5]: {:?} output_occupied_len: {:?}",
            self.extra_input,
            self.extra_input_occupied_len,
            &self.output[0..5],
            self.output_occupied_len
        )
    }
}

impl<'a, W: Write> EncoderWriter<'a, W> {
    /// Create a new encoder that will write to the provided delegate writer `w`.
    pub fn new(w: &'a mut W, config: Config) -> EncoderWriter<'a, W> {
        EncoderWriter {
            config,
            w,
            extra_input: [0u8; MIN_ENCODE_CHUNK_SIZE],
            extra_input_occupied_len: 0,
            output: [0u8; BUF_SIZE],
            output_occupied_len: 0,
            finished: false,
            panicked: false,
        }
    }

    /// Encode all remaining buffered data and write it, including any trailing incomplete input
    /// triples and associated padding.
    ///
    /// Once this succeeds, no further writes can be performed, as that would produce invalid
    /// base64.
    ///
    /// This may write to the delegate writer multiple times if the delegate writer does not accept all input provided
    /// to its `write` each invocation.
    ///
    /// # Errors
    ///
    /// The first error that is not of [`ErrorKind::Interrupted`] will be returned.
    pub fn finish(&mut self) -> Result<()> {
        if self.finished {
            return Ok(());
        };

        self.write_all_encoded_output()?;

        if self.extra_input_occupied_len > 0 {
            let encoded_len = encode_config_slice(
                &self.extra_input[..self.extra_input_occupied_len],
                self.config,
                &mut self.output[..],
            );

            self.output_occupied_len = encoded_len;

            self.write_all_encoded_output()?;

            // write succeeded, do not write the encoding of extra again if finish() is retried
            self.extra_input_occupied_len = 0;
        }

        self.finished = true;
        Ok(())
    }

    /// Write as much of the encoded output to the delegate writer as it will accept, and store the
    /// leftovers to be attempted at the next write() call. Updates `self.output_occupied_len`.
    ///
    /// # Errors
    ///
    /// Errors from the delegate writer are returned. In the case of an error,
    /// `self.output_occupied_len` will not be updated, as errors from `write` are specified to mean
    /// that no write took place.
    fn write_to_delegate(&mut self, current_output_len: usize) -> Result<()> {
        self.panicked = true;
        let res = self.w.write(&self.output[..current_output_len]);
        self.panicked = false;

        return res.map(|consumed| {
            debug_assert!(consumed <= current_output_len);

            if consumed < current_output_len {
                self.output_occupied_len = current_output_len.checked_sub(consumed).unwrap();
                // If we're blocking on I/O, the minor inefficiency of copying bytes to the
                // start of the buffer is the least of our concerns...
                // Rotate moves more than we need to, but copy_within isn't stabilized yet.
                self.output.rotate_left(consumed);
            } else {
                self.output_occupied_len = 0;
            }

            ()
        });
    }

    /// Write all buffered encoded output. If this returns `Ok`, `self.output_occupied_len` is `0`.
    ///
    /// This is basically write_all for the remaining buffered data but without the undesirable
    /// abort-on-`Ok(0)` behavior.
    ///
    /// # Errors
    ///
    /// Any error emitted by the delegate writer abort the write loop and is returned, unless it's
    /// `Interrupted`, in which case the error is ignored and writes will continue.
    fn write_all_encoded_output(&mut self) -> Result<()> {
        while self.output_occupied_len > 0 {
            let remaining_len = self.output_occupied_len;
            match self.write_to_delegate(remaining_len) {
                // try again on interrupts ala write_all
                Err(ref e) if e.kind() == ErrorKind::Interrupted => {}
                // other errors return
                Err(e) => return Err(e),
                // success no-ops because remaining length is already updated
                Ok(_) => {}
            };
        }

        debug_assert_eq!(0, self.output_occupied_len);
        Ok(())
    }
}

impl<'a, W: Write> Write for EncoderWriter<'a, W> {
    /// Encode input and then write to the delegate writer.
    ///
    /// Under non-error circumstances, this returns `Ok` with the value being the number of bytes
    /// of `input` consumed. The value may be `0`, which interacts poorly with `write_all`, which
    /// interprets `Ok(0)` as an error, despite it being allowed by the contract of `write`. See
    /// https://github.com/rust-lang/rust/issues/56889 for more on that.
    ///
    /// If the previous call to `write` provided more (encoded) data than the delegate writer could
    /// accept in a single call to its `write`, the remaining data is buffered. As long as buffered
    /// data is present, subsequent calls to `write` will try to write the remaining buffered data
    /// to the delegate and return either `Ok(0)` -- and therefore not consume any of `input` -- or
    /// an error.
    ///
    /// # Errors
    ///
    /// Any errors emitted by the delegate writer are returned.
    fn write(&mut self, input: &[u8]) -> Result<usize> {
        if self.finished {
            panic!("Cannot write more after calling finish()");
        }

        if input.is_empty() {
            return Ok(0);
        }

        // The contract of `Write::write` places some constraints on this implementation:
        // - a call to `write()` represents at most one call to a wrapped `Write`, so we can't
        // iterate over the input and encode multiple chunks.
        // - Errors mean that "no bytes were written to this writer", so we need to reset the
        // internal state to what it was before the error occurred

        // before reading any input, write any leftover encoded output from last time
        if self.output_occupied_len > 0 {
            let current_len = self.output_occupied_len;
            return self.write_to_delegate(current_len)
                // did not read any input
                .map(|_| 0)

        }

        debug_assert_eq!(0, self.output_occupied_len);

        // how many bytes, if any, were read into `extra` to create a triple to encode
        let mut extra_input_read_len = 0;
        let mut input = input;

        let orig_extra_len = self.extra_input_occupied_len;

        let mut encoded_size = 0;
        // always a multiple of MIN_ENCODE_CHUNK_SIZE
        let mut max_input_len = MAX_INPUT_LEN;

        // process leftover un-encoded input from last write
        if self.extra_input_occupied_len > 0 {
            debug_assert!(self.extra_input_occupied_len < 3);
            if input.len() + self.extra_input_occupied_len >= MIN_ENCODE_CHUNK_SIZE {
                // Fill up `extra`, encode that into `output`, and consume as much of the rest of
                // `input` as possible.
                // We could write just the encoding of `extra` by itself but then we'd have to
                // return after writing only 4 bytes, which is inefficient if the underlying writer
                // would make a syscall.
                extra_input_read_len = MIN_ENCODE_CHUNK_SIZE - self.extra_input_occupied_len;
                debug_assert!(extra_input_read_len > 0);
                // overwrite only bytes that weren't already used. If we need to rollback extra_len
                // (when the subsequent write errors), the old leading bytes will still be there.
                self.extra_input[self.extra_input_occupied_len..MIN_ENCODE_CHUNK_SIZE]
                    .copy_from_slice(&input[0..extra_input_read_len]);

                let len = encode_to_slice(
                    &self.extra_input[0..MIN_ENCODE_CHUNK_SIZE],
                    &mut self.output[..],
                    self.config.char_set.encode_table(),
                );
                debug_assert_eq!(4, len);

                input = &input[extra_input_read_len..];

                // consider extra to be used up, since we encoded it
                self.extra_input_occupied_len = 0;
                // don't clobber where we just encoded to
                encoded_size = 4;
                // and don't read more than can be encoded
                max_input_len = MAX_INPUT_LEN - MIN_ENCODE_CHUNK_SIZE;

                // fall through to normal encoding
            } else {
                // `extra` and `input` are non empty, but `|extra| + |input| < 3`, so there must be
                // 1 byte in each.
                debug_assert_eq!(1, input.len());
                debug_assert_eq!(1, self.extra_input_occupied_len);

                self.extra_input[self.extra_input_occupied_len] = input[0];
                self.extra_input_occupied_len += 1;
                return Ok(1);
            };
        } else if input.len() < MIN_ENCODE_CHUNK_SIZE {
            // `extra` is empty, and `input` fits inside it
            self.extra_input[0..input.len()].copy_from_slice(input);
            self.extra_input_occupied_len = input.len();
            return Ok(input.len());
        };

        // either 0 or 1 complete chunks encoded from extra
        debug_assert!(encoded_size == 0 || encoded_size == 4);
        debug_assert!(
            // didn't encode extra input
            MAX_INPUT_LEN == max_input_len
                // encoded one triple
                || MAX_INPUT_LEN == max_input_len + MIN_ENCODE_CHUNK_SIZE
        );

        // encode complete triples only
        let input_complete_chunks_len = input.len() - (input.len() % MIN_ENCODE_CHUNK_SIZE);
        let input_chunks_to_encode_len = cmp::min(input_complete_chunks_len, max_input_len);
        debug_assert_eq!(0, max_input_len % MIN_ENCODE_CHUNK_SIZE);
        debug_assert_eq!(0, input_chunks_to_encode_len % MIN_ENCODE_CHUNK_SIZE);

        encoded_size += encode_to_slice(
            &input[..(input_chunks_to_encode_len)],
            &mut self.output[encoded_size..],
            self.config.char_set.encode_table(),
        );

        // not updating `self.output_occupied_len` here because if the below write fails, it should
        // "never take place" -- the buffer contents we encoded are ignored and perhaps retried
        // later, if the consumer chooses.

        self.write_to_delegate(encoded_size)
            // no matter whether we wrote the full encoded buffer or not, we consumed the same
            // input
            .map(|_| extra_input_read_len + input_chunks_to_encode_len)
            .map_err( |e| {
                // in case we filled and encoded `extra`, reset extra_len
                self.extra_input_occupied_len = orig_extra_len;

                e
            })
    }

    /// Because this is usually treated as OK to call multiple times, it will *not* flush any
    /// incomplete chunks of input or write padding.
    fn flush(&mut self) -> Result<()> {
        self.write_all_encoded_output()?;
        self.w.flush()
    }
}

impl<'a, W: Write> Drop for EncoderWriter<'a, W> {
    fn drop(&mut self) {
        if !self.panicked {
            // like `BufWriter`, ignore errors during drop
            let _ = self.finish();
        }
    }
}