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
17pub 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 #[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 #[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
234pub 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 #[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 #[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}