diff --git a/crates/develop/src/segmask.rs b/crates/develop/src/segmask.rs index 172d53fce..de0aede09 100644 --- a/crates/develop/src/segmask.rs +++ b/crates/develop/src/segmask.rs @@ -104,9 +104,14 @@ fn base64_encode(bytes: &[u8]) -> String { fn base64_decode(s: &str) -> Option> { let mut out = Vec::with_capacity(s.len() / 4 * 3); let (mut acc, mut bits) = (0u32, 0u32); + let mut padded = false; for c in s.bytes().filter(|c| !c.is_ascii_whitespace()) { if c == b'=' { - break; + padded = true; + continue; + } + if padded { + return None; } let v = B64.iter().position(|b| *b == c)? as u32; acc = (acc << 6) | v; @@ -117,6 +122,9 @@ fn base64_decode(s: &str) -> Option> { acc &= (1 << bits) - 1; } } + if bits == 6 { + return None; + } Some(out) } @@ -133,6 +141,10 @@ mod tests { assert_eq!(base64_encode(b"Man"), "TWFu"); assert_eq!(base64_encode(b"Ma"), "TWE="); assert!(base64_decode("TW!u").is_none()); + assert!(base64_decode("Z").is_none()); + assert!(base64_decode("Zm9vY").is_none()); + assert!(base64_decode("Zg==!!!not-base64").is_none()); + assert_eq!(base64_decode("Zg==\n").unwrap(), b"f"); } #[test] diff --git a/crates/mcp/src/base64.rs b/crates/mcp/src/base64.rs index 9c02ac327..12ccf45bb 100644 --- a/crates/mcp/src/base64.rs +++ b/crates/mcp/src/base64.rs @@ -20,9 +20,14 @@ pub fn base64_encode(bytes: &[u8]) -> String { pub fn base64_decode(text: &str) -> Option> { let mut out = Vec::with_capacity(text.len() / 4 * 3); let (mut acc, mut bits) = (0u32, 0u32); + let mut padded = false; for c in text.bytes().filter(|c| !c.is_ascii_whitespace()) { if c == b'=' { - break; + padded = true; + continue; + } + if padded { + return None; } let v = ALPHABET.iter().position(|&a| a == c)? as u32; acc = (acc << 6) | v; @@ -33,6 +38,9 @@ pub fn base64_decode(text: &str) -> Option> { acc &= (1 << bits) - 1; } } + if bits == 6 { + return None; + } Some(out) } @@ -57,4 +65,12 @@ mod tests { } assert!(base64_decode("not base64!").is_none()); } + + #[test] + fn rejects_malformed() { + assert!(base64_decode("Z").is_none()); + assert!(base64_decode("Zm9vY").is_none()); + assert!(base64_decode("Zg==!!!not-base64").is_none()); + assert_eq!(base64_decode("Zg==\n").unwrap(), b"f"); + } }