on main help sign in
fix(mcp): handle raw/whitespace base64
commitdf094bdeca6717cb36c127fddda69a041b4e38f1 authorChristian Rocha <christian@rocha.is> dateThu 19 Mar 2026, 11:09 (7 months ago) parent00cefc9
2 files changed+69−8
internal/agent/tools/mcp/tools.go+46−8 internal/agent/tools/mcp/tools_test.go+23
internal/agent/tools/mcp/tools.go+46−8
@@ -167,32 +167,70 @@ func filterDisabledTools(cfg *config.ConfigStore, mcpName string, tools []*Tool)
167167 return filtered
168168}
169169
170−// ensureBase64 checks if data is valid base64 and returns it as-is if so,
171−// otherwise encodes the raw binary data to base64.
170+// ensureBase64 normalizes valid base64 input and guarantees padded
171+// base64.StdEncoding output; otherwise it encodes raw binary data.
172172func ensureBase64(data []byte) []byte {
173− // Check if the data is already valid base64 by attempting to decode it.
174− // Valid base64 should only contain ASCII characters (A-Z, a-z, 0-9, +, /, =).
175− if isValidBase64(data) {
173+ if len(data) == 0 {
176174 return data
177175 }
178− // Data is raw binary, encode it to base64.
176+
177+ normalized := normalizeBase64Input(data)
178+ if decoded, ok := decodeBase64(normalized); ok {
179+ encoded := make([]byte, base64.StdEncoding.EncodedLen(len(decoded)))
180+ base64.StdEncoding.Encode(encoded, decoded)
181+ return encoded
182+ }
183+
179184 encoded := make([]byte, base64.StdEncoding.EncodedLen(len(data)))
180185 base64.StdEncoding.Encode(encoded, data)
181186 return encoded
182187}
183188
189+func normalizeBase64Input(data []byte) []byte {
190+ normalized := strings.Join(strings.Fields(string(data)), "")
191+ return []byte(normalized)
192+}
193+
194+func decodeBase64(data []byte) ([]byte, bool) {
195+ if len(data) == 0 {
196+ return data, true
197+ }
198+
199+ for _, b := range data {
200+ if b > 127 {
201+ return nil, false
202+ }
203+ }
204+
205+ s := string(data)
206+ decoded, err := base64.StdEncoding.DecodeString(s)
207+ if err == nil {
208+ return decoded, true
209+ }
210+ decoded, err = base64.RawStdEncoding.DecodeString(s)
211+ if err == nil {
212+ return decoded, true
213+ }
214+ return nil, false
215+}
216+
184217// isValidBase64 checks if the data appears to be valid base64-encoded content.
185218func isValidBase64(data []byte) bool {
186219 if len(data) == 0 {
187220 return true
188221 }
222+
189223 // Base64 strings should only contain ASCII characters.
190224 for _, b := range data {
191225 if b > 127 {
192226 return false
193227 }
194228 }
195− // Try to decode to verify it's valid base64.
196− _, err := base64.StdEncoding.DecodeString(string(data))
229+
230+ s := string(data)
231+ if _, err := base64.StdEncoding.DecodeString(s); err == nil {
232+ return true
233+ }
234+ _, err := base64.RawStdEncoding.DecodeString(s)
197235 return err == nil
198236}
internal/agent/tools/mcp/tools_test.go+23
@@ -40,6 +40,16 @@ func TestEnsureBase64(t *testing.T) {
4040 input: []byte("YQ=="), // "a" in base64
4141 wantData: []byte("YQ=="),
4242 },
43+ {
44+ name: "base64 without padding",
45+ input: []byte("YQ"),
46+ wantData: []byte("YQ=="),
47+ },
48+ {
49+ name: "base64 with whitespace",
50+ input: []byte("U0dWc2JHOGdWMjl5YkdRaA==\n"),
51+ wantData: []byte("U0dWc2JHOGdWMjl5YkdRaA=="),
52+ },
4353 }
4454
4555 for _, tt := range tests {
@@ -51,6 +61,9 @@ func TestEnsureBase64(t *testing.T) {
5161 // Verify the result is valid base64 that can be decoded.
5262 if len(result) > 0 {
5363 _, err := base64.StdEncoding.DecodeString(string(result))
64+ if err != nil {
65+ _, err = base64.RawStdEncoding.DecodeString(string(result))
66+ }
5467 require.NoError(t, err, "result should be valid base64")
5568 }
5669 })
@@ -85,6 +98,16 @@ func TestIsValidBase64(t *testing.T) {
8598 input: []byte{},
8699 want: true,
87100 },
101+ {
102+ name: "valid raw base64 without padding",
103+ input: []byte("YQ"),
104+ want: true,
105+ },
106+ {
107+ name: "valid base64 with whitespace",
108+ input: normalizeBase64Input([]byte("U0dWc2JHOGdWMjl5YkdRaA==\n")),
109+ want: true,
110+ },
88111 {
89112 name: "invalid base64 characters",
90113 input: []byte("SGVsbG8!@#$"),