Skip to main content

base64_stream/
to_base64_writer.rs

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