AkurAI-Framework / crates/ws/src/base64.rs

Latest change f9fa0a48c9b9f40868041393d79336eeb32d8592 - Image generation API, app VERSION in health, base64 decode by Ólafur Búi Ólafsson

//! Standard base64 encoding (RFC 4648, `+/` alphabet, `=` padding) — just
//! enough to turn the 20-byte SHA-1 handshake digest into the textual
//! `Sec-WebSocket-Accept` header, and decode for upstream payloads that carry
//! binary as base64 (e.g. an image API's `b64_json`).

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

/// Encode bytes as standard, padded base64.
pub fn encode(input: &[u8]) -> String {
    let mut out = String::with_capacity(input.len().div_ceil(3) * 4);
    for chunk in input.chunks(3) {
        // Pack up to three bytes into a 24-bit big-endian group.
        let b0 = chunk[0] as u32;
        let b1 = *chunk.get(1).unwrap_or(&0) as u32;
        let b2 = *chunk.get(2).unwrap_or(&0) as u32;
        let group = (b0 << 16) | (b1 << 8) | b2;

        out.push(ALPHABET[(group >> 18 & 0x3f) as usize] as char);
        out.push(ALPHABET[(group >> 12 & 0x3f) as usize] as char);
        // The third and fourth symbols become '=' when their source bytes are
        // absent (1 or 2 trailing input bytes).
        out.push(if chunk.len() > 1 {
            ALPHABET[(group >> 6 & 0x3f) as usize] as char
        } else {
            '='
        });
        out.push(if chunk.len() > 2 {
            ALPHABET[(group & 0x3f) as usize] as char
        } else {
            '='
        });
    }
    out
}

/// Decode standard, padded base64. ASCII whitespace is skipped (wrapped
/// payloads); any other non-alphabet byte, or a length that is not a multiple
/// of four, is an error.
pub fn decode(input: &str) -> Result<Vec<u8>, String> {
    let sextet = |c: u8| ALPHABET.iter().position(|&a| a == c).map(|i| i as u32);
    let clean: Vec<u8> = input.bytes().filter(|b| !b.is_ascii_whitespace()).collect();
    if !clean.len().is_multiple_of(4) {
        return Err("base64 length is not a multiple of 4".into());
    }
    let mut out = Vec::with_capacity(clean.len() / 4 * 3);
    for quad in clean.chunks(4) {
        let pad = quad.iter().rev().take_while(|&&b| b == b'=').count();
        let mut group = 0u32;
        for &c in &quad[..4 - pad] {
            group =
                group << 6 | sextet(c).ok_or_else(|| format!("invalid base64 byte {c:#04x}"))?;
        }
        group <<= 6 * pad as u32;
        out.extend_from_slice(&group.to_be_bytes()[1..4 - pad]);
    }
    Ok(out)
}

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

    #[test]
    fn rfc4648_vectors() {
        // The canonical progression that exercises every padding case.
        assert_eq!(encode(b""), "");
        assert_eq!(encode(b"f"), "Zg==");
        assert_eq!(encode(b"fo"), "Zm8=");
        assert_eq!(encode(b"foo"), "Zm9v");
        assert_eq!(encode(b"foob"), "Zm9vYg==");
        assert_eq!(encode(b"fooba"), "Zm9vYmE=");
        assert_eq!(encode(b"foobar"), "Zm9vYmFy");
    }

    #[test]
    fn decode_round_trips_rfc4648_vectors() {
        for raw in [
            &b""[..],
            b"f",
            b"fo",
            b"foo",
            b"foob",
            b"fooba",
            b"foobar",
            &[0xff, 0, 0x7f],
        ] {
            assert_eq!(decode(&encode(raw)).unwrap(), raw);
        }
    }

    #[test]
    fn decode_skips_whitespace_and_rejects_garbage() {
        assert_eq!(decode("Zm9v\nYmFy").unwrap(), b"foobar");
        assert!(decode("Zm9v!").is_err());
        assert!(decode("Zm9").is_err());
    }

    #[test]
    fn encodes_high_bytes() {
        assert_eq!(encode(&[0xff, 0xff, 0xff]), "////");
        assert_eq!(encode(&[0x00, 0x00, 0x00]), "AAAA");
    }
}