From bdf535e954f60e89a66a169cb075b291f7e39b34 Mon Sep 17 00:00:00 2001
From: NetterSilver <133535268+Silvertree2010@users.noreply.github.com>
Date: Thu, 29 Oct 2026 12:42:30 +0000
Subject: [PATCH] Reject malformed base64 in MCP and segmask decoders

Found while testing crates/mcp/src/base64.rs:20 (base64_decode)

Co-Authored-By: Claude Haiku 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
---
 crates/develop/src/segmask.rs | 14 +++++++++++++-
 crates/mcp/src/base64.rs      | 18 +++++++++++++++++-
 2 files changed, 30 insertions(+), 2 deletions(-)

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<Vec<u8>> {
     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<Vec<u8>> {
             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<Vec<u8>> {
     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<Vec<u8>> {
             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");
+    }
 }