krypteia-arcana 0.1.0

//! PEM armor: base64 encoding with `-----BEGIN/END ...-----` headers.

// ====================================================================
// Base64 (RFC 4648)
// ====================================================================

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

/// Encode `data` as base64.
pub fn base64_encode(data: &[u8]) -> String {
    let mut out = String::with_capacity((data.len() + 2) / 3 * 4);
    for chunk in data.chunks(3) {
        let b0 = chunk[0] as u32;
        let b1 = if chunk.len() > 1 { chunk[1] as u32 } else { 0 };
        let b2 = if chunk.len() > 2 { chunk[2] as u32 } else { 0 };
        let triple = (b0 << 16) | (b1 << 8) | b2;

        out.push(B64_CHARS[((triple >> 18) & 0x3F) as usize] as char);
        out.push(B64_CHARS[((triple >> 12) & 0x3F) as usize] as char);
        if chunk.len() > 1 {
            out.push(B64_CHARS[((triple >> 6) & 0x3F) as usize] as char);
        } else {
            out.push('=');
        }
        if chunk.len() > 2 {
            out.push(B64_CHARS[(triple & 0x3F) as usize] as char);
        } else {
            out.push('=');
        }
    }
    out
}

/// Decode a base64 string. Returns `None` on invalid input.
pub fn base64_decode(s: &str) -> Option<Vec<u8>> {
    let mut out = Vec::with_capacity(s.len() * 3 / 4);
    let mut buf = 0u32;
    let mut count = 0u32;

    for c in s.chars() {
        if c.is_whitespace() {
            continue;
        }
        if c == '=' {
            break;
        }
        let val = b64_char_value(c)? as u32;
        buf = (buf << 6) | val;
        count += 1;
        if count == 4 {
            out.push((buf >> 16) as u8);
            out.push((buf >> 8) as u8);
            out.push(buf as u8);
            buf = 0;
            count = 0;
        }
    }
    match count {
        2 => {
            buf <<= 12;
            out.push((buf >> 16) as u8);
        }
        3 => {
            buf <<= 6;
            out.push((buf >> 16) as u8);
            out.push((buf >> 8) as u8);
        }
        0 => {}
        _ => return None, // 1 leftover char is invalid
    }
    Some(out)
}

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

// ====================================================================
// PEM
// ====================================================================

/// Encode DER bytes as PEM with the given label (e.g. "RSA PUBLIC KEY").
pub fn pem_encode(label: &str, der: &[u8]) -> String {
    let b64 = base64_encode(der);
    let mut out = String::new();
    out.push_str("-----BEGIN ");
    out.push_str(label);
    out.push_str("-----\n");
    // Wrap at 64 characters per line (PEM standard).
    for line in b64.as_bytes().chunks(64) {
        out.push_str(std::str::from_utf8(line).unwrap());
        out.push('\n');
    }
    out.push_str("-----END ");
    out.push_str(label);
    out.push_str("-----\n");
    out
}

/// Decode a PEM block with the expected label. Returns the DER bytes.
/// Returns `None` if the label doesn't match or base64 is invalid.
pub fn pem_decode(label: &str, pem: &str) -> Option<Vec<u8>> {
    let begin = format!("-----BEGIN {}-----", label);
    let end = format!("-----END {}-----", label);

    let start = pem.find(&begin)? + begin.len();
    let stop = pem.find(&end)?;
    let b64_block = &pem[start..stop];
    base64_decode(b64_block)
}

// ====================================================================
// Tests
// ====================================================================

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

    #[test]
    fn base64_roundtrip() {
        let cases: &[(&[u8], &str)] = &[
            (b"", ""),
            (b"f", "Zg=="),
            (b"fo", "Zm8="),
            (b"foo", "Zm9v"),
            (b"foob", "Zm9vYg=="),
            (b"fooba", "Zm9vYmE="),
            (b"foobar", "Zm9vYmFy"),
        ];
        for (data, expected) in cases {
            let encoded = base64_encode(data);
            assert_eq!(encoded, *expected, "encode {:?}", data);
            let decoded = base64_decode(&encoded).unwrap();
            assert_eq!(decoded, *data, "decode {:?}", expected);
        }
    }

    #[test]
    fn base64_binary() {
        let data: Vec<u8> = (0..256).map(|i| i as u8).collect();
        let enc = base64_encode(&data);
        let dec = base64_decode(&enc).unwrap();
        assert_eq!(dec, data);
    }

    #[test]
    fn pem_roundtrip() {
        let der = vec![0x30, 0x03, 0x02, 0x01, 0x42]; // tiny DER
        let pem = pem_encode("TEST DATA", &der);
        assert!(pem.contains("-----BEGIN TEST DATA-----"));
        assert!(pem.contains("-----END TEST DATA-----"));
        let decoded = pem_decode("TEST DATA", &pem).unwrap();
        assert_eq!(decoded, der);
    }

    #[test]
    fn pem_wrong_label_returns_none() {
        let pem = pem_encode("FOO", &[0x01]);
        assert!(pem_decode("BAR", &pem).is_none());
    }
}