git.iko.soy/crush
on main help sign in
fix(mcp): handle raw/whitespace base64
df094bdChristian Rocha
internal/agent/tools/mcp/tools.go+46−8
compare with main
unifiedsplit
@@ -167,32 +167,70 @@ func filterDisabledTools(cfg *config.ConfigStore, mcpName string, tools []*Tool)
167 return filtered
167 return filtered
168}
168}
169
169
170−// ensureBase64 checks if data is valid base64 and returns it as-is if so,
170+// ensureBase64 normalizes valid base64 input and guarantees padded
171−// otherwise encodes the raw binary data to base64.
171+// base64.StdEncoding output; otherwise it encodes raw binary data.
172func ensureBase64(data []byte) []byte {
172func ensureBase64(data []byte) []byte {
173− // Check if the data is already valid base64 by attempting to decode it.
173+ if len(data) == 0 {
174− // Valid base64 should only contain ASCII characters (A-Z, a-z, 0-9, +, /, =).
175− if isValidBase64(data) {
176 return data
174 return data
177 }
175 }
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+
179 encoded := make([]byte, base64.StdEncoding.EncodedLen(len(data)))
184 encoded := make([]byte, base64.StdEncoding.EncodedLen(len(data)))
180 base64.StdEncoding.Encode(encoded, data)
185 base64.StdEncoding.Encode(encoded, data)
181 return encoded
186 return encoded
182}
187}
183
188
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+
184// isValidBase64 checks if the data appears to be valid base64-encoded content.
217// isValidBase64 checks if the data appears to be valid base64-encoded content.
185func isValidBase64(data []byte) bool {
218func isValidBase64(data []byte) bool {
186 if len(data) == 0 {
219 if len(data) == 0 {
187 return true
220 return true
188 }
221 }
222+
189 // Base64 strings should only contain ASCII characters.
223 // Base64 strings should only contain ASCII characters.
190 for _, b := range data {
224 for _, b := range data {
191 if b > 127 {
225 if b > 127 {
192 return false
226 return false
193 }
227 }
194 }
228 }
195− // Try to decode to verify it's valid base64.
229+
196− _, err := base64.StdEncoding.DecodeString(string(data))
230+ s := string(data)
231+ if _, err := base64.StdEncoding.DecodeString(s); err == nil {
232+ return true
233+ }
234+ _, err := base64.RawStdEncoding.DecodeString(s)
197 return err == nil
235 return err == nil
198}
236}
internal/agent/tools/mcp/tools_test.go+23
@@ -40,6 +40,16 @@ func TestEnsureBase64(t *testing.T) {
40 input: []byte("YQ=="), // "a" in base64
40 input: []byte("YQ=="), // "a" in base64
41 wantData: []byte("YQ=="),
41 wantData: []byte("YQ=="),
42 },
42 },
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+ },
43 }
53 }
44
54
45 for _, tt := range tests {
55 for _, tt := range tests {
@@ -51,6 +61,9 @@ func TestEnsureBase64(t *testing.T) {
51 // Verify the result is valid base64 that can be decoded.
61 // Verify the result is valid base64 that can be decoded.
52 if len(result) > 0 {
62 if len(result) > 0 {
53 _, err := base64.StdEncoding.DecodeString(string(result))
63 _, err := base64.StdEncoding.DecodeString(string(result))
64+ if err != nil {
65+ _, err = base64.RawStdEncoding.DecodeString(string(result))
66+ }
54 require.NoError(t, err, "result should be valid base64")
67 require.NoError(t, err, "result should be valid base64")
55 }
68 }
56 })
69 })
@@ -85,6 +98,16 @@ func TestIsValidBase64(t *testing.T) {
85 input: []byte{},
98 input: []byte{},
86 want: true,
99 want: true,
87 },
100 },
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+ },
88 {
111 {
89 name: "invalid base64 characters",
112 name: "invalid base64 characters",
90 input: []byte("SGVsbG8!@#$"),
113 input: []byte("SGVsbG8!@#$"),