Base64 iterator adaptor for Rust

Question

As a Rust learning exercise, I'm working on implementing cryptopals challenges using iterators to perform lazy processing of data. For the first challenge, this is a base64 decoder implemented as an iterator adaptor:

use std::iter;

enum B64Result {
    Byte(u32),
    Pad,
    Invalid
}

fn b64_character_to_byte(byte: u8) -> B64Result {
    if byte >= b'A' && byte <= b'Z' {
        B64Result::Byte((byte - 65) as u32)
    } else if byte >= b'a' && byte <= b'z' {
        B64Result::Byte((byte - 97 + 26) as u32)
    } else if byte >= b'0' && byte <= b'9' {
        B64Result::Byte((byte - 48 + 52) as u32)
    } else if byte == b'+' {
        B64Result::Byte(62)
    } else if byte == b'/' {
        B64Result::Byte(63)
    } else if byte == b'=' {
        B64Result::Pad
    } else {
        B64Result::Invalid
    }
}

// Decodes a 4-character block of base64 characters, return the number of bytes decoded.
fn decode_block(input: &[u8; 4], output: &mut [u8; 3]) -> Result<u8, &'static str> {
    let mut accum = 0 as u32;
    let mut pad_count = 0;

    for (i, character) in input.iter().enumerate() {
        let shift = 18 - (i % 4) * 6;
        match b64_character_to_byte(*character) {
            B64Result::Byte(byte) => {
                if pad_count != 0 {
                    return Err("Invalid base64 string");
                }
                accum |= byte << shift;
            },
            B64Result::Pad => {
                if pad_count == 2 {
                    return Err("Invalid base64 string");
                }
                pad_count += 1;
            },
            B64Result::Invalid => return Err("Invalid base64 string"),
        }
    }

    output[0] = (accum >> 16) as u8;
    output[1] = ((accum >> 8) & 0xff) as u8;
    output[2] = (accum & 0xff) as u8;
    Ok(3 - pad_count)
}

pub struct Base64Decoder<'a, T>
    where T: Iterator<Item=u8> + 'a
{
    input: &'a mut T,
    buffer: [u8; 3],
    buffer_pos: u8,
    buffer_len: u8,
}

impl<'a, T> Base64Decoder<'a, T>
    where T: Iterator<Item=u8> + 'a
{
    pub fn new(input: &'a mut T) -> Base64Decoder<'a, T> {
        Base64Decoder::<'a, T> {
            input,
            buffer: [0; 3],
            buffer_pos: 0,
            buffer_len: 0,
        }
    }
}

impl<'a, T> iter::Iterator for Base64Decoder<'a, T>
    where T: Iterator<Item=u8> + 'a
{
    type Item = u8;

    fn next(&mut self) -> Option<Self::Item> {
        if self.buffer_pos == 0 {
            let mut buf = [0; 4];
            let mut len = 0;
            // fetch more from input
            while len < 4 {
                match self.input.next() {
                    Some(b) => {
                        if b != b'\n' {
                            buf[len] = b;
                            len += 1;
                        } else {
                            continue;
                        }
                    }
                    None => return None
                }
            }
            match decode_block(&buf, &mut self.buffer) {
                Ok(len) => self.buffer_len = len,
                Err(_) => return None
            }
        }

        let rv = Some(self.buffer[self.buffer_pos as usize]);
        self.buffer_pos = (self.buffer_pos + 1) % self.buffer_len;
        rv
    }
}

pub trait Base64Decodable: Iterator<Item=u8> {
    fn b64decode<'a>(&'a mut self) -> Base64Decoder<'a, Self>
        where Self: Sized + Iterator<Item=u8> + 'a
    {
        Base64Decoder::new(self)
    }
}

impl<'a, T> Base64Decodable for T
    where T: Iterator<Item=u8> + 'a
{ }

The idea is to be able to lazily decode base64 from any u8 iterator. This is an example usage:

// simple program to decode base64 from stdin
use std::io;
use std::io::Read;
use std::io::Write;

use matasano::util::base64::Base64Decodable;

fn main() {        
    for byte in io::stdin().bytes().filter(|b| b.is_ok()).map(|b| b.unwrap()).b64decode() {
            io::stdout().write([byte]).unwrap();
        }
    }    
}

Is this a good design? What can I do to improve, or what other tools/patterns can I use to implement lazy transformations like this?

Answer 1

Overall, I think this is a very clean implementation! My suggestions are minor:

1. In b64_character_to_byte, you use byte literals to check which range of ASCII codes the byte falls in, but then you use "magic numbers" in subtractions to refer to ASCII codes. You could use byte literals again; this would make the code clearer.

   fn b64_character_to_byte(byte: u8) -> B64Result {
       if byte >= b'A' && byte <= b'Z' {
           B64Result::Byte((byte - b'A') as u32)
       } else if byte >= b'a' && byte <= b'z' {
           B64Result::Byte((byte - b'a' + 26) as u32)
       } else if byte >= b'0' && byte <= b'9' {
           B64Result::Byte((byte - b'0' + 52) as u32)
   

2. In general, using &'static str for the error type is not a good idea, because taking a decision based on the contents of a string is not very robust. It's preferable to define a struct or a enum that implements the Error trait. The Error trait also makes it easy to compose errors in a complex program.

3. Normally, iterator adapters take ownership of the original iterator. This is true of the iterators in the standard library; for example, map takes self by value, and the Map struct doesn't have any lifetime parameters. The reason is that an iterator adapter often consumes the original iterator to perform its own iteration (assuming it is driven to completion).

   On the other hand, yours takes the original iterator by reference. This is less flexible, because it means that I can't write a function that returns a Base64Decoder unless I received an iterator of bytes as a parameter (I can't return a Base64Decoder that refers to a local byte iterator, because the byte iterator would be dropped when the function returns.)

   There are cases where it's useful to use an iterator adapter without giving away ownership of the original iterator. For these occasions, we can use by_ref to give a reference to an iterator to an iterator adapter. Mutable references to iterators are iterators themselves, thanks to a blanket implementation in the standard library (impl<'a, I> Iterator for &'a mut I where I: Iterator + ?Sized).

4. In your iterator's next method, the continue in the while loop in superfluous. I only write continue when necessary; when I see a continue, I expect to find some code after it, and when there isn't any code, it makes me wonder whether I missed something.

5. In your iterator's next method, you use buf to refer to the encoded string and self.buffer to refer to the decoded string. buf and buffer are similar terms that don't carry any meaning. It might be clearer to name them encoded and decoded, respectively.