AkurAI-Build / src/base64.rs

Latest change e8aa34a9be50d7d7ea72f66991ff82dfca7d23e1 - OLI-276: Git over SSH, forks and mirrors, code search through akurai-index, repository insights by Ólafur Búi Ólafsson

//! Standard base64 (RFC 4648) without padding tolerance tricks: encode with
//! padding, decode with or without it. Small enough to own rather than depend on.

const ALPHABET: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

pub fn encode(bytes: &[u8]) -> String {
    let mut out = String::with_capacity(bytes.len().div_ceil(3) * 4);
    for chunk in bytes.chunks(3) {
        let b = [
            chunk[0],
            chunk.get(1).copied().unwrap_or(0),
            chunk.get(2).copied().unwrap_or(0),
        ];
        let n = (u32::from(b[0]) << 16) | (u32::from(b[1]) << 8) | u32::from(b[2]);
        let chars = [
            ALPHABET[(n >> 18) as usize & 63],
            ALPHABET[(n >> 12) as usize & 63],
            ALPHABET[(n >> 6) as usize & 63],
            ALPHABET[n as usize & 63],
        ];
        out.push(chars[0] as char);
        out.push(chars[1] as char);
        out.push(if chunk.len() > 1 {
            chars[2] as char
        } else {
            '='
        });
        out.push(if chunk.len() > 2 {
            chars[3] as char
        } else {
            '='
        });
    }
    out
}

/// Like [`encode`] but without `=` padding (how OpenSSH prints fingerprints).
pub fn encode_unpadded(bytes: &[u8]) -> String {
    encode(bytes).trim_end_matches('=').to_owned()
}

fn value(c: u8) -> Option<u32> {
    Some(match c {
        b'A'..=b'Z' => u32::from(c - b'A'),
        b'a'..=b'z' => u32::from(c - b'a') + 26,
        b'0'..=b'9' => u32::from(c - b'0') + 52,
        b'+' => 62,
        b'/' => 63,
        _ => return None,
    })
}

/// Decode standard base64; padding is optional. `None` on any foreign byte.
pub fn decode(text: &str) -> Option<Vec<u8>> {
    let text = text.trim_end_matches('=');
    if text.len() % 4 == 1 {
        return None;
    }
    let mut out = Vec::with_capacity(text.len() * 3 / 4);
    let mut acc = 0u32;
    let mut bits = 0u32;
    for byte in text.bytes() {
        acc = (acc << 6) | value(byte)?;
        bits += 6;
        if bits >= 8 {
            bits -= 8;
            out.push((acc >> bits) as u8);
            acc &= (1 << bits) - 1;
        }
    }
    Some(out)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn round_trips_with_and_without_padding() {
        for input in ["", "f", "fo", "foo", "foob", "fooba", "foobar"] {
            let encoded = encode(input.as_bytes());
            assert_eq!(decode(&encoded).as_deref(), Some(input.as_bytes()));
            assert_eq!(
                decode(&encode_unpadded(input.as_bytes())).as_deref(),
                Some(input.as_bytes())
            );
        }
        assert_eq!(encode(b"foobar"), "Zm9vYmFy");
        assert_eq!(encode(b"fooba"), "Zm9vYmE=");
        assert_eq!(encode_unpadded(b"foob"), "Zm9vYg");
    }

    #[test]
    fn rejects_foreign_bytes_and_impossible_lengths() {
        assert_eq!(decode("Zm9v*mFy"), None);
        assert_eq!(decode("Z"), None);
    }
}