diff --git a/src/chunked_encoder.rs b/src/chunked_encoder.rs index 0457259..bc3810a 100644 --- a/src/chunked_encoder.rs +++ b/src/chunked_encoder.rs @@ -49,7 +49,8 @@ impl<'e, E: Engine + ?Sized> ChunkedEncoder<'e, E> { if self.engine.config().encode_padding() && !more_input_left { // no more input, add padding if needed. Buffer will have room because // max_input_length leaves room for it. - b64_bytes_written += add_padding(bytes.len(), &mut encode_buf[b64_bytes_written..]); + b64_bytes_written += + add_padding(b64_bytes_written, &mut encode_buf[b64_bytes_written..]); } sink.write_encoded_bytes(&encode_buf[0..b64_bytes_written])?; diff --git a/src/encode.rs b/src/encode.rs index cb17650..15b903d 100644 --- a/src/encode.rs +++ b/src/encode.rs @@ -77,7 +77,7 @@ pub(crate) fn encode_with_padding( let b64_bytes_written = engine.internal_encode(input, output); let padding_bytes = if engine.config().encode_padding() { - add_padding(input.len(), &mut output[b64_bytes_written..]) + add_padding(b64_bytes_written, &mut output[b64_bytes_written..]) } else { 0 }; @@ -117,20 +117,20 @@ pub fn encoded_len(bytes_len: usize, padding: bool) -> Option { } /// Write padding characters. -/// `input_len` is the size of the original, not encoded, input. +/// `unpadded_output_len` is the size of the unpadded but base64 encoded data. /// `output` is the slice where padding should be written, of length at least 2. /// /// Returns the number of padding bytes written. -pub(crate) fn add_padding(input_len: usize, output: &mut [u8]) -> usize { - // TODO base on encoded len to use cheaper mod by 4 (aka & 7) - let rem = input_len % 3; - let mut bytes_written = 0; - for _ in 0..((3 - rem) % 3) { - output[bytes_written] = PAD_BYTE; - bytes_written += 1; +pub(crate) fn add_padding(unpadded_output_len: usize, output: &mut [u8]) -> usize { + let pad_bytes = (4 - (unpadded_output_len % 4)) % 4; + // for just a couple bytes, this has better performance than using + // .fill(), or iterating over mutable refs, which call memset() + #[allow(clippy::needless_range_loop)] + for i in 0..pad_bytes { + output[i] = PAD_BYTE; } - bytes_written + pad_bytes } /// Errors that can occur while encoding into a slice. @@ -434,18 +434,18 @@ mod tests { let mut rng = rand::rngs::SmallRng::from_entropy(); - // cover our bases for length % 3 - for input_len in 0..10 { + // cover our bases for length % 4 + for unpadded_output_len in 0..20 { output.clear(); // fill output with random - for _ in 0..10 { + for _ in 0..100 { output.push(rng.gen()); } let orig_output_buf = output.clone(); - let bytes_written = add_padding(input_len, &mut output); + let bytes_written = add_padding(unpadded_output_len, &mut output); // make sure the part beyond bytes_written is the same garbage it was before assert_eq!(orig_output_buf[bytes_written..], output[bytes_written..]); diff --git a/src/engine/tests.rs b/src/engine/tests.rs index 2b4d2e5..d2851ec 100644 --- a/src/engine/tests.rs +++ b/src/engine/tests.rs @@ -86,7 +86,7 @@ fn rfc_test_vectors_std_alphabet(engine_wrapper: E) { &encoded_without_padding, &std::str::from_utf8(&encode_buf[0..encode_len]).unwrap() ); - let pad_len = add_padding(orig.len(), &mut encode_buf[encode_len..]); + let pad_len = add_padding(encode_len, &mut encode_buf[encode_len..]); assert_eq!(encoded.as_bytes(), &encode_buf[..encode_len + pad_len]); let decode_len = engine @@ -195,7 +195,10 @@ fn encode_doesnt_write_extra_bytes(engine_wrapper: E) { // pad so we can decode it in case our random engine requires padding let pad_len = if padded { - add_padding(orig_len, &mut encode_buf[prefix_len + encoded_len_no_pad..]) + add_padding( + encoded_len_no_pad, + &mut encode_buf[prefix_len + encoded_len_no_pad..], + ) } else { 0 }; @@ -382,7 +385,7 @@ fn decode_detect_invalid_last_symbol_every_possible_two_symbols