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    /// Returns the inner writer, consuming this wrapper.
72    ///
73    /// Call [`flush`](std::io::Write::flush) before this method to ensure all
74    /// buffered data is written.
75    #[inline]
76    pub fn into_inner(self) -> W {
77        self.inner
78    }
79}
80
81impl<W: Write, const N: usize> Write for FromBase64Writer<W, N> {
82    fn write(&mut self, mut buf: &[u8]) -> Result<usize, io::Error> {
83        let original_buf_length = buf.len();
84
85        if self.buf_length == 0 {
86            while buf.len() >= 4 {
87                let max_available_buf_length = (buf.len() & !0b11).min((N / 3) << 2); // (N / 3) * 4
88
89                let decode_length = self
90                    .engine
91                    .decode_slice(&buf[..max_available_buf_length], &mut self.temp)
92                    .map_err(decode_error_to_io_error)?;
93
94                buf = &buf[max_available_buf_length..];
95
96                self.inner.write_all(&self.temp[..decode_length])?;
97            }
98
99            let buf_length = buf.len();
100
101            if buf_length > 0 {
102                self.buf[..buf_length].copy_from_slice(&buf[..buf_length]);
103
104                self.buf_length = buf_length;
105            }
106        } else {
107            debug_assert!(self.buf_length < 4);
108
109            let r = 4 - self.buf_length;
110
111            let buf_length = buf.len();
112
113            let drain_length = r.min(buf_length);
114
115            self.buf[self.buf_length..self.buf_length + drain_length]
116                .copy_from_slice(&buf[..drain_length]);
117
118            buf = &buf[drain_length..];
119
120            self.buf_length += drain_length;
121
122            if self.buf_length == 4 {
123                self.drain_block()?;
124
125                if buf_length > r {
126                    self.write_all(buf)?;
127                }
128            }
129        }
130
131        Ok(original_buf_length)
132    }
133
134    #[inline]
135    fn flush(&mut self) -> Result<(), io::Error> {
136        if self.buf_length > 0 {
137            self.drain_block()?;
138        }
139
140        self.inner.flush()
141    }
142}
143
144impl<W: Write> From<W> for FromBase64Writer<W> {
145    #[inline]
146    fn from(reader: W) -> Self {
147        FromBase64Writer::new(reader)
148    }
149}
base64_stream/from_base64_writer.rs

base64_stream/
from_base64_writer.rs
