base64_ng_tokio/
decoder_writer.rs

1use base64_ng::{Alphabet, Engine};
2use core::{
3    marker::PhantomData,
4    pin::Pin,
5    task::{Context, Poll, ready},
6};
7use tokio::io::{self, AsyncWrite};
8
9use crate::{decode_io_error, queue::OutputQueue, wipe_bytes};
10
11const DECODE_OUTPUT_CAP: usize = 1024;
12
13/// Async writer that accepts Base64 bytes and writes decoded bytes to `inner`.
14///
15/// The adapter validates strict Base64 quanta as input is accepted. Shutdown is
16/// the finalization boundary and validates any final unpadded tail.
17///
18/// # Security
19///
20/// Streaming decode is not atomic. Decoded bytes from valid leading quanta may
21/// already have been written before a later malformed quantum is observed. For
22/// atomic or secret-bearing frames, collect a bounded frame and use a
23/// non-streaming strict or `ct` decode path.
24///
25/// I/O errors from the wrapped writer during drain do not set [`Self::is_failed`];
26/// only internal protocol or capacity violations latch a permanent failure.
27pub struct DecoderWriter<W, A, const PAD: bool>
28where
29    A: Alphabet,
30{
31    inner: Option<W>,
32    engine: Engine<A, PAD>,
33    pending: [u8; 4],
34    pending_len: usize,
35    output: OutputQueue<DECODE_OUTPUT_CAP>,
36    terminal_padding: bool,
37    finalized: bool,
38    failed: bool,
39    _alphabet: PhantomData<A>,
40}
41
42impl<W, A, const PAD: bool> DecoderWriter<W, A, PAD>
43where
44    A: Alphabet,
45{
46    /// Creates a new async Base64 decoder writer.
47    #[must_use]
48    pub fn new(inner: W, engine: Engine<A, PAD>) -> Self {
49        Self {
50            inner: Some(inner),
51            engine,
52            pending: [0; 4],
53            pending_len: 0,
54            output: OutputQueue::new(),
55            terminal_padding: false,
56            finalized: false,
57            failed: false,
58            _alphabet: PhantomData,
59        }
60    }
61
62    /// Returns a shared reference to the wrapped writer.
63    #[must_use]
64    pub fn get_ref(&self) -> &W {
65        self.inner_ref()
66    }
67
68    /// Returns a mutable reference to the wrapped writer.
69    pub fn get_mut(&mut self) -> &mut W {
70        self.inner_mut()
71    }
72
73    /// Consumes the adapter and returns the wrapped writer.
74    ///
75    /// This does not validate final pending input. Prefer
76    /// [`AsyncWriteExt::shutdown`](tokio::io::AsyncWriteExt::shutdown) before
77    /// calling this when the decoded stream must be complete.
78    #[must_use]
79    pub fn into_inner(mut self) -> W {
80        self.take_inner()
81    }
82
83    /// Returns whether this adapter has encountered an unrecoverable error.
84    #[must_use]
85    pub const fn is_failed(&self) -> bool {
86        self.failed
87    }
88
89    /// Returns whether shutdown has finalized this adapter.
90    #[must_use]
91    pub const fn is_finalized(&self) -> bool {
92        self.finalized
93    }
94
95    /// Returns whether a terminal padded quantum has been accepted.
96    #[must_use]
97    pub const fn has_terminal_padding(&self) -> bool {
98        self.terminal_padding
99    }
100
101    /// Returns the number of encoded bytes buffered until a full decode
102    /// quantum is available.
103    #[must_use]
104    pub const fn pending_len(&self) -> usize {
105        self.pending_len
106    }
107
108    /// Returns the number of decoded bytes currently buffered for `inner`.
109    #[must_use]
110    pub const fn buffered_output_len(&self) -> usize {
111        self.output.len()
112    }
113
114    fn clear_pending(&mut self) {
115        wipe_bytes(&mut self.pending);
116        self.pending_len = 0;
117    }
118
119    fn clear_output(&mut self) {
120        self.output.clear_all();
121    }
122
123    fn inner_ref(&self) -> &W {
124        match &self.inner {
125            Some(inner) => inner,
126            None => unreachable!("tokio decoder writer inner writer was already taken"),
127        }
128    }
129
130    fn inner_mut(&mut self) -> &mut W {
131        match &mut self.inner {
132            Some(inner) => inner,
133            None => unreachable!("tokio decoder writer inner writer was already taken"),
134        }
135    }
136
137    fn take_inner(&mut self) -> W {
138        match self.inner.take() {
139            Some(inner) => inner,
140            None => unreachable!("tokio decoder writer inner writer was already taken"),
141        }
142    }
143
144    fn queue_decoded_temp(&mut self, input: &[u8], decoded: &mut [u8]) -> io::Result<usize> {
145        let written = match self.engine.decode_slice(input, decoded) {
146            Ok(written) => written,
147            Err(error) => {
148                wipe_bytes(decoded);
149                self.failed = true;
150                return Err(decode_io_error(error));
151            }
152        };
153
154        let result = self.output.push_slice(&decoded[..written]);
155        wipe_bytes(decoded);
156        if result.is_err() {
157            self.failed = true;
158        }
159        result?;
160        Ok(written)
161    }
162
163    fn queue_full_quad(&mut self, mut input: [u8; 4]) -> io::Result<()> {
164        let mut decoded = [0u8; 3];
165        let result = self.queue_decoded_temp(&input, &mut decoded);
166        wipe_bytes(&mut input);
167        let written = result?;
168        if written < 3 {
169            self.terminal_padding = true;
170        }
171        Ok(())
172    }
173
174    fn queue_pending_final(&mut self) -> io::Result<()> {
175        if self.pending_len == 0 {
176            return Ok(());
177        }
178
179        if PAD || self.pending_len == 1 {
180            self.failed = true;
181            self.clear_pending();
182            return Err(io::Error::new(
183                io::ErrorKind::InvalidData,
184                "base64-ng-tokio decoder writer received incomplete final quantum",
185            ));
186        }
187
188        let mut pending = [0u8; 4];
189        pending[..self.pending_len].copy_from_slice(&self.pending[..self.pending_len]);
190        let pending_len = self.pending_len;
191        let mut decoded = [0u8; 3];
192        let result = self.queue_decoded_temp(&pending[..pending_len], &mut decoded);
193        wipe_bytes(&mut pending);
194        result?;
195        self.clear_pending();
196        Ok(())
197    }
198
199    fn trailing_input_error(&mut self) -> io::Error {
200        self.failed = true;
201        io::Error::new(
202            io::ErrorKind::InvalidData,
203            "base64-ng-tokio decoder writer received trailing input after padding",
204        )
205    }
206
207    fn process_input(&mut self, input: &[u8]) -> io::Result<usize> {
208        if input.is_empty() {
209            return Ok(0);
210        }
211        if self.terminal_padding {
212            return Err(self.trailing_input_error());
213        }
214
215        let mut consumed = 0;
216        if self.pending_len > 0 {
217            let needed = 4 - self.pending_len;
218            if input.len() < needed {
219                self.pending[self.pending_len..self.pending_len + input.len()]
220                    .copy_from_slice(input);
221                self.pending_len += input.len();
222                return Ok(input.len());
223            }
224
225            let mut quad = [0u8; 4];
226            quad[..self.pending_len].copy_from_slice(&self.pending[..self.pending_len]);
227            quad[self.pending_len..].copy_from_slice(&input[..needed]);
228            let result = self.queue_full_quad(quad);
229            wipe_bytes(&mut quad);
230            if let Err(error) = result {
231                self.clear_pending();
232                return Err(error);
233            }
234            self.clear_pending();
235            consumed += needed;
236
237            if self.terminal_padding {
238                return Ok(consumed);
239            }
240        }
241
242        while input.len() - consumed >= 4 {
243            if self.output.available_capacity() < 3 {
244                return Ok(consumed);
245            }
246
247            let mut quad = [
248                input[consumed],
249                input[consumed + 1],
250                input[consumed + 2],
251                input[consumed + 3],
252            ];
253            let mut decoded = [0u8; 3];
254            let written = match self.engine.decode_slice(&quad, &mut decoded) {
255                Ok(written) => written,
256                Err(error) => {
257                    wipe_bytes(&mut quad);
258                    wipe_bytes(&mut decoded);
259                    if consumed > 0 {
260                        return Ok(consumed);
261                    }
262                    self.failed = true;
263                    return Err(decode_io_error(error));
264                }
265            };
266
267            let result = self.output.push_slice(&decoded[..written]);
268            wipe_bytes(&mut quad);
269            wipe_bytes(&mut decoded);
270            if result.is_err() {
271                self.failed = true;
272            }
273            result?;
274            consumed += 4;
275
276            if written < 3 {
277                self.terminal_padding = true;
278                return Ok(consumed);
279            }
280        }
281
282        let tail = &input[consumed..];
283        self.pending[..tail.len()].copy_from_slice(tail);
284        self.pending_len = tail.len();
285        consumed += tail.len();
286        Ok(consumed)
287    }
288}
289
290impl<W, A, const PAD: bool> Drop for DecoderWriter<W, A, PAD>
291where
292    A: Alphabet,
293{
294    fn drop(&mut self) {
295        self.clear_pending();
296        self.clear_output();
297    }
298}
299
300impl<W, A, const PAD: bool> DecoderWriter<W, A, PAD>
301where
302    W: AsyncWrite + Unpin,
303    A: Alphabet + Unpin,
304{
305    fn poll_drain_output(&mut self, context: &mut Context<'_>) -> Poll<io::Result<()>> {
306        let mut chunk = [0u8; DECODE_OUTPUT_CAP];
307        while !self.output.is_empty() {
308            let pending = self.output.copy_front(&mut chunk);
309            let result = Pin::new(self.inner_mut()).poll_write(context, &chunk[..pending]);
310            wipe_bytes(&mut chunk[..pending]);
311            match result {
312                Poll::Pending => return Poll::Pending,
313                Poll::Ready(Ok(0)) => {
314                    return Poll::Ready(Err(io::Error::new(
315                        io::ErrorKind::WriteZero,
316                        "base64-ng-tokio decoder writer could not drain buffered output",
317                    )));
318                }
319                Poll::Ready(Ok(written)) => {
320                    if written > pending {
321                        self.failed = true;
322                        return Poll::Ready(Err(io::Error::new(
323                            io::ErrorKind::InvalidData,
324                            "wrapped async writer reported more bytes than provided",
325                        )));
326                    }
327                    self.output.discard_front(written);
328                }
329                Poll::Ready(Err(error)) => return Poll::Ready(Err(error)),
330            }
331        }
332
333        Poll::Ready(Ok(()))
334    }
335}
336
337impl<W, A, const PAD: bool> AsyncWrite for DecoderWriter<W, A, PAD>
338where
339    W: AsyncWrite + Unpin,
340    A: Alphabet + Unpin,
341{
342    fn poll_write(
343        mut self: Pin<&mut Self>,
344        context: &mut Context<'_>,
345        input: &[u8],
346    ) -> Poll<io::Result<usize>> {
347        if self.failed {
348            return Poll::Ready(Err(io::Error::other(
349                "base64-ng-tokio decoder writer is failed",
350            )));
351        }
352
353        ready!(self.poll_drain_output(context))?;
354        if self.finalized {
355            return Poll::Ready(Err(io::Error::new(
356                io::ErrorKind::InvalidInput,
357                "base64-ng-tokio decoder writer received input after shutdown",
358            )));
359        }
360
361        Poll::Ready(self.process_input(input))
362    }
363
364    fn poll_flush(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<io::Result<()>> {
365        if self.failed {
366            return Poll::Ready(Err(io::Error::other(
367                "base64-ng-tokio decoder writer is failed",
368            )));
369        }
370
371        ready!(self.poll_drain_output(context))?;
372        Pin::new(self.inner_mut()).poll_flush(context)
373    }
374
375    fn poll_shutdown(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<io::Result<()>> {
376        if self.failed {
377            return Poll::Ready(Err(io::Error::other(
378                "base64-ng-tokio decoder writer is failed",
379            )));
380        }
381
382        ready!(self.poll_drain_output(context))?;
383        if !self.finalized {
384            if let Err(error) = self.queue_pending_final() {
385                self.clear_output();
386                return Poll::Ready(Err(error));
387            }
388            self.finalized = true;
389        }
390        ready!(self.poll_drain_output(context))?;
391        ready!(Pin::new(self.inner_mut()).poll_flush(context))?;
392        Pin::new(self.inner_mut()).poll_shutdown(context)
393    }
394}
base64_ng_tokio/decoder_writer.rs

base64_ng_tokio/
decoder_writer.rs
