1 files changed, 150 insertions, 9 deletions

diff --git a/src/read/decoder_tests.rs b/src/read/decoder_tests.rs
index 65d58d8..099dd63 100644
--- a/src/read/decoder_tests.rs
+++ b/src/read/decoder_tests.rs
@@ -8,9 +8,10 @@ use rand::{Rng as _, RngCore as _};
 
 use super::decoder::{DecoderReader, BUF_SIZE};
 use crate::{
+    alphabet,
     engine::{general_purpose::STANDARD, Engine, GeneralPurpose},
     tests::{random_alphabet, random_config, random_engine},
-    DecodeError,
+    DecodeError, PAD_BYTE,
 };
 
 #[test]
@@ -75,7 +76,7 @@ fn trailing_junk() {
                         saw_error = true;
                         break;
                     }
-                    Ok(read) if read == 0 => break,
+                    Ok(0) => break,
                     Ok(_) => (),
                 }
             }
@@ -247,19 +248,21 @@ fn reports_invalid_byte_correctly() {
     let mut rng = rand::thread_rng();
     let mut bytes = Vec::new();
     let mut b64 = String::new();
-    let mut decoded = Vec::new();
+    let mut stream_decoded = Vec::new();
+    let mut bulk_decoded = Vec::new();
 
     for _ in 0..10_000 {
         bytes.clear();
         b64.clear();
-        decoded.clear();
+        stream_decoded.clear();
+        bulk_decoded.clear();
 
         let size = rng.gen_range(1..(10 * BUF_SIZE));
         bytes.extend(iter::repeat(0).take(size));
         rng.fill_bytes(&mut bytes[..size]);
         assert_eq!(size, bytes.len());
 
-        let engine = random_engine(&mut rng);
+        let engine = GeneralPurpose::new(&alphabet::STANDARD, random_config(&mut rng));
 
         engine.encode_string(&bytes[..], &mut b64);
         // replace one byte, somewhere, with '*', which is invalid
@@ -270,9 +273,8 @@ fn reports_invalid_byte_correctly() {
         let mut wrapped_reader = io::Cursor::new(b64_bytes.clone());
         let mut decoder = DecoderReader::new(&mut wrapped_reader, &engine);
 
-        // some gymnastics to avoid double-moving the io::Error, which is not Copy
         let read_decode_err = decoder
-            .read_to_end(&mut decoded)
+            .read_to_end(&mut stream_decoded)
             .map_err(|e| {
                 let kind = e.kind();
                 let inner = e
@@ -283,8 +285,7 @@ fn reports_invalid_byte_correctly() {
             .err()
             .and_then(|o| o);
 
-        let mut bulk_buf = Vec::new();
-        let bulk_decode_err = engine.decode_vec(&b64_bytes[..], &mut bulk_buf).err();
+        let bulk_decode_err = engine.decode_vec(&b64_bytes[..], &mut bulk_decoded).err();
 
         // it's tricky to predict where the invalid data's offset will be since if it's in the last
         // chunk it will be reported at the first padding location because it's treated as invalid
@@ -296,6 +297,134 @@ fn reports_invalid_byte_correctly() {
     }
 }
 
+#[test]
+fn internal_padding_error_with_short_read_concatenated_texts_invalid_byte_error() {
+    let mut rng = rand::thread_rng();
+    let mut bytes = Vec::new();
+    let mut b64 = String::new();
+    let mut reader_decoded = Vec::new();
+    let mut bulk_decoded = Vec::new();
+
+    // encodes with padding, requires that padding be present so we don't get InvalidPadding
+    // just because padding is there at all
+    let engine = STANDARD;
+
+    for _ in 0..10_000 {
+        bytes.clear();
+        b64.clear();
+        reader_decoded.clear();
+        bulk_decoded.clear();
+
+        // at least 2 bytes so there can be a split point between bytes
+        let size = rng.gen_range(2..(10 * BUF_SIZE));
+        bytes.resize(size, 0);
+        rng.fill_bytes(&mut bytes[..size]);
+
+        // Concatenate two valid b64s, yielding padding in the middle.
+        // This avoids scenarios that are challenging to assert on, like random padding location
+        // that might be InvalidLastSymbol when decoded at certain buffer sizes but InvalidByte
+        // when done all at once.
+        let split = loop {
+            // find a split point that will produce padding on the first part
+            let s = rng.gen_range(1..size);
+            if s % 3 != 0 {
+                // short enough to need padding
+                break s;
+            };
+        };
+
+        engine.encode_string(&bytes[..split], &mut b64);
+        assert!(b64.contains('='), "split: {}, b64: {}", split, b64);
+        let bad_byte_pos = b64.find('=').unwrap();
+        engine.encode_string(&bytes[split..], &mut b64);
+        let b64_bytes = b64.as_bytes();
+
+        // short read to make it plausible for padding to happen on a read boundary
+        let read_len = rng.gen_range(1..10);
+        let mut wrapped_reader = ShortRead {
+            max_read_len: read_len,
+            delegate: io::Cursor::new(&b64_bytes),
+        };
+
+        let mut decoder = DecoderReader::new(&mut wrapped_reader, &engine);
+
+        let read_decode_err = decoder
+            .read_to_end(&mut reader_decoded)
+            .map_err(|e| {
+                *e.into_inner()
+                    .and_then(|e| e.downcast::<DecodeError>().ok())
+                    .unwrap()
+            })
+            .unwrap_err();
+
+        let bulk_decode_err = engine.decode_vec(b64_bytes, &mut bulk_decoded).unwrap_err();
+
+        assert_eq!(
+            bulk_decode_err,
+            read_decode_err,
+            "read len: {}, bad byte pos: {}, b64: {}",
+            read_len,
+            bad_byte_pos,
+            std::str::from_utf8(b64_bytes).unwrap()
+        );
+        assert_eq!(
+            DecodeError::InvalidByte(
+                split / 3 * 4
+                    + match split % 3 {
+                        1 => 2,
+                        2 => 3,
+                        _ => unreachable!(),
+                    },
+                PAD_BYTE
+            ),
+            read_decode_err
+        );
+    }
+}
+
+#[test]
+fn internal_padding_anywhere_error() {
+    let mut rng = rand::thread_rng();
+    let mut bytes = Vec::new();
+    let mut b64 = String::new();
+    let mut reader_decoded = Vec::new();
+
+    // encodes with padding, requires that padding be present so we don't get InvalidPadding
+    // just because padding is there at all
+    let engine = STANDARD;
+
+    for _ in 0..10_000 {
+        bytes.clear();
+        b64.clear();
+        reader_decoded.clear();
+
+        bytes.resize(10 * BUF_SIZE, 0);
+        rng.fill_bytes(&mut bytes[..]);
+
+        // Just shove a padding byte in there somewhere.
+        // The specific error to expect is challenging to predict precisely because it
+        // will vary based on the position of the padding in the quad and the read buffer
+        // length, but SOMETHING should go wrong.
+
+        engine.encode_string(&bytes[..], &mut b64);
+        let mut b64_bytes = b64.as_bytes().to_vec();
+        // put padding somewhere other than the last quad
+        b64_bytes[rng.gen_range(0..bytes.len() - 4)] = PAD_BYTE;
+
+        // short read to make it plausible for padding to happen on a read boundary
+        let read_len = rng.gen_range(1..10);
+        let mut wrapped_reader = ShortRead {
+            max_read_len: read_len,
+            delegate: io::Cursor::new(&b64_bytes),
+        };
+
+        let mut decoder = DecoderReader::new(&mut wrapped_reader, &engine);
+
+        let result = decoder.read_to_end(&mut reader_decoded);
+        assert!(result.is_err());
+    }
+}
+
 fn consume_with_short_reads_and_validate<R: io::Read>(
     rng: &mut rand::rngs::ThreadRng,
     expected_bytes: &[u8],
@@ -344,3 +473,15 @@ impl<'a, 'b, R: io::Read, N: rand::Rng> io::Read for RandomShortRead<'a, 'b, R,
         self.delegate.read(&mut buf[..effective_len])
     }
 }
+
+struct ShortRead<R: io::Read> {
+    delegate: R,
+    max_read_len: usize,
+}
+
+impl<R: io::Read> io::Read for ShortRead<R> {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        let len = self.max_read_len.max(buf.len());
+        self.delegate.read(&mut buf[..len])
+    }
+}