base64_stream/
from_base64_writer.rs

1use std::{
2    fmt,
3    io::{self, ErrorKind, Write},
4};
5
6use base64::{
7    DecodeSliceError, Engine,
8    engine::{GeneralPurpose, general_purpose::STANDARD},
9};
10
11#[inline]
12fn decode_error_to_io_error(error: DecodeSliceError) -> io::Error {
13    io::Error::new(ErrorKind::InvalidData, error)
14}
15
16/// Write base64 data and decode them to plain data.
17pub struct FromBase64Writer<W: Write, const N: usize = 4096> {
18    inner:      W,
19    buf:        [u8; 4],
20    buf_length: usize,
21    temp:       [u8; N],
22    engine:     &'static GeneralPurpose,
23}
24
25impl<W: Write, const N: usize> fmt::Debug for FromBase64Writer<W, N> {
26    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
27        f.debug_struct("FromBase64Writer")
28            .field("buf", &&self.buf[..self.buf_length])
29            .field("buf_length", &self.buf_length)
30            .finish_non_exhaustive()
31    }
32}
33
34impl<W: Write> FromBase64Writer<W> {
35    #[inline]
36    pub fn new(writer: W) -> FromBase64Writer<W> {
37        Self::new2(writer)
38    }
39}
40
41impl<W: Write, const N: usize> FromBase64Writer<W, N> {
42    #[inline]
43    pub fn new2(writer: W) -> FromBase64Writer<W, N> {
44        const { assert!(N >= 4, "buffer size N must be at least 4") };
45        FromBase64Writer {
46            inner:      writer,
47            buf:        [0; 4],
48            buf_length: 0,
49            temp:       [0u8; N],
50            engine:     &STANDARD,
51        }
52    }
53}
54
55impl<W: Write, const N: usize> FromBase64Writer<W, N> {
56    fn drain_block(&mut self) -> Result<(), io::Error> {
57        debug_assert!(self.buf_length > 0);
58
59        let decode_length = self
60            .engine
61            .decode_slice(&self.buf[..self.buf_length], &mut self.temp)
62            .map_err(decode_error_to_io_error)?;
63
64        self.inner.write_all(&self.temp[..decode_length])?;
65
66        self.buf_length = 0;
67
68        Ok(())
69    }
70
71    /// Finishes decoding buffered data and returns the inner writer.
72    #[inline]
73    pub fn finish(mut self) -> Result<W, io::Error> {
74        if self.buf_length > 0 {
75            self.drain_block()?;
76        }
77
78        self.inner.flush()?;
79
80        Ok(self.inner)
81    }
82
83    /// Returns the inner writer, consuming this wrapper without finishing it.
84    ///
85    /// Call [`finish`](Self::finish) instead to decode the final buffered data and flush the inner writer.
86    #[inline]
87    pub fn into_inner(self) -> W {
88        self.inner
89    }
90}
91
92impl<W: Write, const N: usize> Write for FromBase64Writer<W, N> {
93    fn write(&mut self, mut buf: &[u8]) -> Result<usize, io::Error> {
94        let original_buf_length = buf.len();
95
96        if self.buf_length == 0 {
97            while buf.len() >= 4 {
98                let max_available_buf_length = (buf.len() & !0b11).min((N / 3) << 2); // (N / 3) * 4
99
100                let decode_length = self
101                    .engine
102                    .decode_slice(&buf[..max_available_buf_length], &mut self.temp)
103                    .map_err(decode_error_to_io_error)?;
104
105                buf = &buf[max_available_buf_length..];
106
107                self.inner.write_all(&self.temp[..decode_length])?;
108            }
109
110            let buf_length = buf.len();
111
112            if buf_length > 0 {
113                self.buf[..buf_length].copy_from_slice(&buf[..buf_length]);
114
115                self.buf_length = buf_length;
116            }
117        } else {
118            debug_assert!(self.buf_length < 4);
119
120            let r = 4 - self.buf_length;
121
122            let buf_length = buf.len();
123
124            let drain_length = r.min(buf_length);
125
126            self.buf[self.buf_length..self.buf_length + drain_length]
127                .copy_from_slice(&buf[..drain_length]);
128
129            buf = &buf[drain_length..];
130
131            self.buf_length += drain_length;
132
133            if self.buf_length == 4 {
134                self.drain_block()?;
135
136                if buf_length > r {
137                    self.write_all(buf)?;
138                }
139            }
140        }
141
142        Ok(original_buf_length)
143    }
144
145    #[inline]
146    fn flush(&mut self) -> Result<(), io::Error> {
147        self.inner.flush()
148    }
149}
150
151impl<W: Write> From<W> for FromBase64Writer<W> {
152    #[inline]
153    fn from(reader: W) -> Self {
154        FromBase64Writer::new(reader)
155    }
156}
base64_stream/from_base64_writer.rs

base64_stream/
from_base64_writer.rs
