Skip to main content

base64_ng_tokio/
readers.rs

1use base64_ng::{Alphabet, Engine};
2use core::{
3    cmp,
4    marker::PhantomData,
5    pin::Pin,
6    task::{Context, Poll},
7};
8use tokio::io::{self, AsyncRead, ReadBuf};
9
10use crate::{decode_io_error, encode_io_error, wipe_bytes};
11
12const ENCODE_INPUT_CAP: usize = 768;
13const ENCODE_OUTPUT_CAP: usize = 1024;
14const DECODE_INPUT_CAP: usize = 1024;
15const DECODE_OUTPUT_CAP: usize = 768;
16
17/// Async reader that streams raw bytes as Base64.
18///
19/// This adapter reads from `inner` in bounded chunks, preserves at most two
20/// pending raw bytes between polls, and clears its pending/output buffers on
21/// drop. It is cancellation-resumable: if a read future is dropped after
22/// returning [`Poll::Pending`], polling the same adapter again continues from
23/// the same internal state without duplicating or dropping bytes.
24///
25/// # Security
26///
27/// Internal cleanup is best-effort and limited to this adapter's fixed buffers.
28/// It cannot clear copies held by the wrapped reader, the caller's output
29/// buffer, registers, caches, swap, or crash dumps.
30pub struct EncoderReader<R, A, const PAD: bool>
31where
32    A: Alphabet,
33{
34    inner: R,
35    engine: Engine<A, PAD>,
36    pending: [u8; 2],
37    pending_len: usize,
38    output: [u8; ENCODE_OUTPUT_CAP],
39    output_pos: usize,
40    output_len: usize,
41    finished: bool,
42    failed: bool,
43    _alphabet: PhantomData<A>,
44}
45
46impl<R, A, const PAD: bool> EncoderReader<R, A, PAD>
47where
48    A: Alphabet,
49{
50    /// Creates a new async Base64 encoder reader.
51    #[must_use]
52    pub fn new(inner: R, engine: Engine<A, PAD>) -> Self {
53        Self {
54            inner,
55            engine,
56            pending: [0; 2],
57            pending_len: 0,
58            output: [0; ENCODE_OUTPUT_CAP],
59            output_pos: 0,
60            output_len: 0,
61            finished: false,
62            failed: false,
63            _alphabet: PhantomData,
64        }
65    }
66
67    /// Returns whether the adapter has encountered an unrecoverable error.
68    #[must_use]
69    pub const fn is_failed(&self) -> bool {
70        self.failed
71    }
72
73    fn clear_buffers(&mut self) {
74        wipe_bytes(&mut self.pending);
75        self.pending_len = 0;
76        wipe_bytes(&mut self.output);
77        self.output_pos = 0;
78        self.output_len = 0;
79    }
80
81    fn drain_output(&mut self, destination: &mut ReadBuf<'_>) -> bool {
82        let available = self.output_len.saturating_sub(self.output_pos);
83        if available == 0 || destination.remaining() == 0 {
84            return false;
85        }
86
87        let count = cmp::min(available, destination.remaining());
88        destination.put_slice(&self.output[self.output_pos..self.output_pos + count]);
89        wipe_bytes(&mut self.output[self.output_pos..self.output_pos + count]);
90        self.output_pos += count;
91        if self.output_pos == self.output_len {
92            self.output_pos = 0;
93            self.output_len = 0;
94        }
95        true
96    }
97
98    fn append_encoded(&mut self, input: &[u8]) -> io::Result<()> {
99        let written = self
100            .engine
101            .encode_slice(input, &mut self.output[self.output_len..])
102            .map_err(encode_io_error)?;
103        self.output_len += written;
104        Ok(())
105    }
106
107    fn process_input(&mut self, input: &[u8]) -> io::Result<()> {
108        let mut read = 0;
109
110        if self.pending_len != 0 {
111            let needed = 3 - self.pending_len;
112            if input.len() < needed {
113                self.pending[self.pending_len..self.pending_len + input.len()]
114                    .copy_from_slice(input);
115                self.pending_len += input.len();
116                return Ok(());
117            }
118
119            let mut quantum = [0u8; 3];
120            quantum[..self.pending_len].copy_from_slice(&self.pending[..self.pending_len]);
121            quantum[self.pending_len..].copy_from_slice(&input[..needed]);
122            let result = self.append_encoded(&quantum);
123            wipe_bytes(&mut quantum);
124            wipe_bytes(&mut self.pending);
125            result?;
126            self.pending_len = 0;
127            read += needed;
128        }
129
130        let remaining = &input[read..];
131        let full_len = remaining.len() / 3 * 3;
132        if full_len != 0 {
133            self.append_encoded(&remaining[..full_len])?;
134        }
135
136        let tail = &remaining[full_len..];
137        if !tail.is_empty() {
138            self.pending[..tail.len()].copy_from_slice(tail);
139            self.pending_len = tail.len();
140        }
141
142        Ok(())
143    }
144
145    fn finish(&mut self) -> io::Result<()> {
146        if self.pending_len != 0 {
147            let mut tail = [0u8; 2];
148            tail[..self.pending_len].copy_from_slice(&self.pending[..self.pending_len]);
149            let pending_len = self.pending_len;
150            let result = self.append_encoded(&tail[..pending_len]);
151            wipe_bytes(&mut tail);
152            wipe_bytes(&mut self.pending);
153            result?;
154            self.pending_len = 0;
155        }
156        self.finished = true;
157        Ok(())
158    }
159}
160
161impl<R, A, const PAD: bool> AsyncRead for EncoderReader<R, A, PAD>
162where
163    R: AsyncRead + Unpin,
164    A: Alphabet + Unpin,
165{
166    fn poll_read(
167        mut self: Pin<&mut Self>,
168        context: &mut Context<'_>,
169        destination: &mut ReadBuf<'_>,
170    ) -> Poll<io::Result<()>> {
171        if self.failed {
172            return Poll::Ready(Err(io::Error::other(
173                "base64-ng-tokio encoder reader is failed",
174            )));
175        }
176
177        if self.drain_output(destination) || destination.remaining() == 0 {
178            return Poll::Ready(Ok(()));
179        }
180
181        if self.finished {
182            return Poll::Ready(Ok(()));
183        }
184
185        loop {
186            let mut input = [0u8; ENCODE_INPUT_CAP];
187            let mut input_buf = ReadBuf::new(&mut input);
188            match Pin::new(&mut self.inner).poll_read(context, &mut input_buf) {
189                Poll::Pending => return Poll::Pending,
190                Poll::Ready(Err(error)) => {
191                    wipe_bytes(&mut input);
192                    self.failed = true;
193                    self.clear_buffers();
194                    return Poll::Ready(Err(error));
195                }
196                Poll::Ready(Ok(())) => {
197                    let read = input_buf.filled().len();
198                    if read == 0 {
199                        let result = self.finish();
200                        wipe_bytes(&mut input);
201                        if let Err(error) = result {
202                            self.failed = true;
203                            self.clear_buffers();
204                            return Poll::Ready(Err(error));
205                        }
206                    } else {
207                        let result = self.process_input(&input[..read]);
208                        wipe_bytes(&mut input);
209                        if let Err(error) = result {
210                            self.failed = true;
211                            self.clear_buffers();
212                            return Poll::Ready(Err(error));
213                        }
214                    }
215
216                    if self.drain_output(destination) || self.finished {
217                        return Poll::Ready(Ok(()));
218                    }
219                }
220            }
221        }
222    }
223}
224
225impl<R, A, const PAD: bool> Drop for EncoderReader<R, A, PAD>
226where
227    A: Alphabet,
228{
229    fn drop(&mut self) {
230        self.clear_buffers();
231    }
232}
233
234/// Async reader that streams Base64 input as decoded bytes.
235///
236/// This adapter decodes strict Base64 quanta as they arrive and preserves at
237/// most three pending encoded bytes between polls. If malformed input is
238/// observed, the adapter fails closed and clears its internal buffers.
239///
240/// # Security
241///
242/// Streaming decode is not atomic. Decoded bytes from valid leading quanta may
243/// already have been returned before a later malformed quantum is observed. For
244/// atomic secret-bearing frames, use `decode_reader_to_writer_limited` or a
245/// `ct` staged decode after collecting a bounded frame.
246pub struct DecoderReader<R, A, const PAD: bool>
247where
248    A: Alphabet,
249{
250    inner: R,
251    engine: Engine<A, PAD>,
252    pending: [u8; 4],
253    pending_len: usize,
254    output: [u8; DECODE_OUTPUT_CAP],
255    output_pos: usize,
256    output_len: usize,
257    finished: bool,
258    failed: bool,
259    terminal_padding: bool,
260    _alphabet: PhantomData<A>,
261}
262
263impl<R, A, const PAD: bool> DecoderReader<R, A, PAD>
264where
265    A: Alphabet,
266{
267    /// Creates a new async Base64 decoder reader.
268    #[must_use]
269    pub fn new(inner: R, engine: Engine<A, PAD>) -> Self {
270        Self {
271            inner,
272            engine,
273            pending: [0; 4],
274            pending_len: 0,
275            output: [0; DECODE_OUTPUT_CAP],
276            output_pos: 0,
277            output_len: 0,
278            finished: false,
279            failed: false,
280            terminal_padding: false,
281            _alphabet: PhantomData,
282        }
283    }
284
285    /// Returns whether the adapter has encountered an unrecoverable error.
286    #[must_use]
287    pub const fn is_failed(&self) -> bool {
288        self.failed
289    }
290
291    fn clear_buffers(&mut self) {
292        wipe_bytes(&mut self.pending);
293        self.pending_len = 0;
294        wipe_bytes(&mut self.output);
295        self.output_pos = 0;
296        self.output_len = 0;
297    }
298
299    fn drain_output(&mut self, destination: &mut ReadBuf<'_>) -> bool {
300        let available = self.output_len.saturating_sub(self.output_pos);
301        if available == 0 || destination.remaining() == 0 {
302            return false;
303        }
304
305        let count = cmp::min(available, destination.remaining());
306        destination.put_slice(&self.output[self.output_pos..self.output_pos + count]);
307        wipe_bytes(&mut self.output[self.output_pos..self.output_pos + count]);
308        self.output_pos += count;
309        if self.output_pos == self.output_len {
310            self.output_pos = 0;
311            self.output_len = 0;
312        }
313        true
314    }
315
316    fn append_decoded(&mut self, input: &[u8]) -> io::Result<usize> {
317        let written = self
318            .engine
319            .decode_slice(input, &mut self.output[self.output_len..])
320            .map_err(decode_io_error)?;
321        self.output_len += written;
322        Ok(written)
323    }
324
325    fn process_quad(&mut self, mut quad: [u8; 4]) -> io::Result<()> {
326        if self.terminal_padding {
327            wipe_bytes(&mut quad);
328            return Err(io::Error::new(
329                io::ErrorKind::InvalidData,
330                "base64-ng-tokio decoder reader received trailing input after padding",
331            ));
332        }
333
334        let result = self.append_decoded(&quad);
335        let saw_terminal = quad.contains(&b'=');
336        wipe_bytes(&mut quad);
337        result?;
338
339        if saw_terminal {
340            self.terminal_padding = true;
341        }
342        Ok(())
343    }
344
345    fn process_input(&mut self, input: &[u8]) -> io::Result<()> {
346        if self.terminal_padding && !input.is_empty() {
347            return Err(io::Error::new(
348                io::ErrorKind::InvalidData,
349                "base64-ng-tokio decoder reader received trailing input after padding",
350            ));
351        }
352
353        let mut read = 0;
354
355        if self.pending_len != 0 {
356            let needed = 4 - self.pending_len;
357            if input.len() < needed {
358                self.pending[self.pending_len..self.pending_len + input.len()]
359                    .copy_from_slice(input);
360                self.pending_len += input.len();
361                return Ok(());
362            }
363
364            let mut quad = [0u8; 4];
365            quad[..self.pending_len].copy_from_slice(&self.pending[..self.pending_len]);
366            quad[self.pending_len..].copy_from_slice(&input[..needed]);
367            self.process_quad(quad)?;
368            wipe_bytes(&mut self.pending);
369            self.pending_len = 0;
370            read += needed;
371        }
372
373        while read + 4 <= input.len() {
374            let quad = [
375                input[read],
376                input[read + 1],
377                input[read + 2],
378                input[read + 3],
379            ];
380            self.process_quad(quad)?;
381            let saw_terminal = self.terminal_padding;
382            read += 4;
383
384            if saw_terminal && read != input.len() {
385                return Err(io::Error::new(
386                    io::ErrorKind::InvalidData,
387                    "base64-ng-tokio decoder reader received trailing input after padding",
388                ));
389            }
390        }
391
392        let tail = &input[read..];
393        if !tail.is_empty() {
394            self.pending[..tail.len()].copy_from_slice(tail);
395            self.pending_len = tail.len();
396        }
397
398        Ok(())
399    }
400
401    fn finish(&mut self) -> io::Result<()> {
402        if self.pending_len != 0 {
403            if PAD || self.pending_len == 1 {
404                return Err(io::Error::new(
405                    io::ErrorKind::InvalidData,
406                    "base64-ng-tokio decoder reader received incomplete final quantum",
407                ));
408            }
409
410            let mut tail = [0u8; 4];
411            tail[..self.pending_len].copy_from_slice(&self.pending[..self.pending_len]);
412            let pending_len = self.pending_len;
413            self.append_decoded(&tail[..pending_len])?;
414            wipe_bytes(&mut tail);
415            wipe_bytes(&mut self.pending);
416            self.pending_len = 0;
417        }
418        self.finished = true;
419        Ok(())
420    }
421}
422
423impl<R, A, const PAD: bool> AsyncRead for DecoderReader<R, A, PAD>
424where
425    R: AsyncRead + Unpin,
426    A: Alphabet + Unpin,
427{
428    fn poll_read(
429        mut self: Pin<&mut Self>,
430        context: &mut Context<'_>,
431        destination: &mut ReadBuf<'_>,
432    ) -> Poll<io::Result<()>> {
433        if self.failed {
434            return Poll::Ready(Err(io::Error::other(
435                "base64-ng-tokio decoder reader is failed",
436            )));
437        }
438
439        if self.drain_output(destination) || destination.remaining() == 0 {
440            return Poll::Ready(Ok(()));
441        }
442
443        if self.finished {
444            return Poll::Ready(Ok(()));
445        }
446
447        loop {
448            let mut input = [0u8; DECODE_INPUT_CAP];
449            let mut input_buf = ReadBuf::new(&mut input);
450            match Pin::new(&mut self.inner).poll_read(context, &mut input_buf) {
451                Poll::Pending => return Poll::Pending,
452                Poll::Ready(Err(error)) => {
453                    wipe_bytes(&mut input);
454                    self.failed = true;
455                    self.clear_buffers();
456                    return Poll::Ready(Err(error));
457                }
458                Poll::Ready(Ok(())) => {
459                    let read = input_buf.filled().len();
460                    if read == 0 {
461                        let result = self.finish();
462                        wipe_bytes(&mut input);
463                        if let Err(error) = result {
464                            self.failed = true;
465                            self.clear_buffers();
466                            return Poll::Ready(Err(error));
467                        }
468                    } else {
469                        let result = self.process_input(&input[..read]);
470                        wipe_bytes(&mut input);
471                        if let Err(error) = result {
472                            self.failed = true;
473                            self.clear_buffers();
474                            return Poll::Ready(Err(error));
475                        }
476                    }
477
478                    if self.drain_output(destination) || self.finished {
479                        return Poll::Ready(Ok(()));
480                    }
481                }
482            }
483        }
484    }
485}
486
487impl<R, A, const PAD: bool> Drop for DecoderReader<R, A, PAD>
488where
489    A: Alphabet,
490{
491    fn drop(&mut self) {
492        self.clear_buffers();
493    }
494}