Skip to content

Commit 935027a

Browse files
jstar0jstar0
committed
fix(mcp): return valid UTF-8 snippets
Signed-off-by: King Star <mcxin.y@gmail.com>
1 parent 1519f86 commit 935027a

2 files changed

Lines changed: 143 additions & 1 deletion

File tree

‎src/mcp/mcp.c‎

Lines changed: 76 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2833,6 +2833,75 @@ static char *resolve_snippet_source(const char *root_path, const char *file_path
28332833
return NULL;
28342834
}
28352835

2836+
static bool utf8_is_cont(unsigned char c) {
2837+
return (c & 0xC0) == 0x80;
2838+
}
2839+
2840+
static char *sanitize_utf8_lossy(const char *s) {
2841+
enum {
2842+
UTF8_REPLACEMENT_LEN = 3,
2843+
UTF8_THREE_BYTE_LEN = 3,
2844+
UTF8_FOUR_BYTE_LEN = 4,
2845+
UTF8_FOURTH_BYTE = 3,
2846+
};
2847+
if (!s) {
2848+
return NULL;
2849+
}
2850+
size_t len = strlen(s);
2851+
if (len > (((size_t)-1) - SKIP_ONE) / UTF8_REPLACEMENT_LEN) {
2852+
return NULL;
2853+
}
2854+
char *out = malloc(len * UTF8_REPLACEMENT_LEN + SKIP_ONE);
2855+
if (!out) {
2856+
return NULL;
2857+
}
2858+
2859+
const unsigned char *p = (const unsigned char *)s;
2860+
const unsigned char *end = p + len;
2861+
unsigned char *dst = (unsigned char *)out;
2862+
while (p < end) {
2863+
unsigned char c = *p;
2864+
size_t n = 0;
2865+
if (c < 0x80) {
2866+
n = 1;
2867+
} else if (c >= 0xC2 && c <= 0xDF && p + 1 < end && utf8_is_cont(p[1])) {
2868+
n = 2;
2869+
} else if (c == 0xE0 && p + 2 < end && p[1] >= 0xA0 && p[1] <= 0xBF && utf8_is_cont(p[2])) {
2870+
n = UTF8_THREE_BYTE_LEN;
2871+
} else if (c >= 0xE1 && c <= 0xEC && p + 2 < end && utf8_is_cont(p[1]) &&
2872+
utf8_is_cont(p[2])) {
2873+
n = UTF8_THREE_BYTE_LEN;
2874+
} else if (c == 0xED && p + 2 < end && p[1] >= 0x80 && p[1] <= 0x9F && utf8_is_cont(p[2])) {
2875+
n = UTF8_THREE_BYTE_LEN;
2876+
} else if (c >= 0xEE && c <= 0xEF && p + 2 < end && utf8_is_cont(p[1]) &&
2877+
utf8_is_cont(p[2])) {
2878+
n = UTF8_THREE_BYTE_LEN;
2879+
} else if (c == 0xF0 && p + UTF8_FOURTH_BYTE < end && p[1] >= 0x90 && p[1] <= 0xBF &&
2880+
utf8_is_cont(p[2]) && utf8_is_cont(p[UTF8_FOURTH_BYTE])) {
2881+
n = UTF8_FOUR_BYTE_LEN;
2882+
} else if (c >= 0xF1 && c <= 0xF3 && p + UTF8_FOURTH_BYTE < end && utf8_is_cont(p[1]) &&
2883+
utf8_is_cont(p[2]) && utf8_is_cont(p[UTF8_FOURTH_BYTE])) {
2884+
n = UTF8_FOUR_BYTE_LEN;
2885+
} else if (c == 0xF4 && p + UTF8_FOURTH_BYTE < end && p[1] >= 0x80 && p[1] <= 0x8F &&
2886+
utf8_is_cont(p[2]) && utf8_is_cont(p[UTF8_FOURTH_BYTE])) {
2887+
n = UTF8_FOUR_BYTE_LEN;
2888+
}
2889+
2890+
if (n > 0) {
2891+
memcpy(dst, p, n);
2892+
dst += n;
2893+
p += n;
2894+
} else {
2895+
*dst++ = 0xEF;
2896+
*dst++ = 0xBF;
2897+
*dst++ = 0xBD;
2898+
p++;
2899+
}
2900+
}
2901+
*dst = '\0';
2902+
return out;
2903+
}
2904+
28362905
/* Build an enriched snippet response for a resolved node. */
28372906
/* Add a string array to a JSON object (no-op if count == 0). */
28382907
static void add_string_array(yyjson_mut_doc *doc, yyjson_mut_val *obj, const char *key,
@@ -2877,7 +2946,13 @@ static char *build_snippet_response(cbm_mcp_server_t *srv, cbm_node_t *node,
28772946
yyjson_mut_obj_add_int(doc, root_obj, "end_line", end);
28782947

28792948
if (source) {
2880-
yyjson_mut_obj_add_str(doc, root_obj, "source", source);
2949+
char *safe_source = sanitize_utf8_lossy(source);
2950+
if (safe_source) {
2951+
yyjson_mut_obj_add_strcpy(doc, root_obj, "source", safe_source);
2952+
free(safe_source);
2953+
} else {
2954+
yyjson_mut_obj_add_str(doc, root_obj, "source", "(source not available)");
2955+
}
28812956
} else {
28822957
yyjson_mut_obj_add_str(doc, root_obj, "source", "(source not available)");
28832958
}

‎tests/test_mcp.c‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
#include <yyjson/yyjson.h>
1212
#include <string.h>
1313
#include <stdlib.h>
14+
#include <stdbool.h>
1415

1516
/* ══════════════════════════════════════════════════════════════════
1617
* JSON-RPC PARSING
@@ -1291,6 +1292,31 @@ static char *call_snippet(cbm_mcp_server_t *srv, const char *args_json) {
12911292
return text;
12921293
}
12931294

1295+
static bool is_valid_json_response(const char *json) {
1296+
if (!json) {
1297+
return false;
1298+
}
1299+
yyjson_doc *doc = yyjson_read(json, strlen(json), 0);
1300+
if (!doc) {
1301+
return false;
1302+
}
1303+
yyjson_doc_free(doc);
1304+
return true;
1305+
}
1306+
1307+
static bool snippet_source_has_replacement(const char *json) {
1308+
yyjson_doc *doc = yyjson_read(json, strlen(json), 0);
1309+
if (!doc) {
1310+
return false;
1311+
}
1312+
yyjson_val *root = yyjson_doc_get_root(doc);
1313+
yyjson_val *source = yyjson_obj_get(root, "source");
1314+
const char *source_str = yyjson_get_str(source);
1315+
bool found = source_str && strstr(source_str, "\xEF\xBF\xBD");
1316+
yyjson_doc_free(doc);
1317+
return found;
1318+
}
1319+
12941320
/* ── TestSnippet_ExactQN ──────────────────────────────────────── */
12951321

12961322
TEST(snippet_exact_qn) {
@@ -1577,6 +1603,46 @@ TEST(snippet_include_neighbors_enabled) {
15771603
PASS();
15781604
}
15791605

1606+
/* ── TestSnippet_SourceInvalidUtf8 ────────────────────────────── */
1607+
1608+
TEST(snippet_source_invalid_utf8) {
1609+
char tmp[256];
1610+
cbm_mcp_server_t *srv = setup_snippet_server(tmp, sizeof(tmp));
1611+
ASSERT_NOT_NULL(srv);
1612+
1613+
char src_path[512];
1614+
snprintf(src_path, sizeof(src_path), "%s/project/main.go", tmp);
1615+
FILE *fp = fopen(src_path, "wb");
1616+
ASSERT_NOT_NULL(fp);
1617+
const unsigned char source[] = {
1618+
'p', 'a', 'c', 'k', 'a', 'g', 'e', ' ', 'm', 'a', 'i', 'n', '\n', '\n',
1619+
'f', 'u', 'n', 'c', ' ', 'H', 'a', 'n', 'd', 'l', 'e', 'R', 'e', 'q',
1620+
'u', 'e', 's', 't', '(', ')', ' ', 'e', 'r', 'r', 'o', 'r', ' ', '{',
1621+
'\n', '\t', '/', '/', ' ', 0xC0, 0xD4, 0xB7, 0xC2, '\n', '\t', 'r', 'e', 't',
1622+
'u', 'r', 'n', ' ', 'n', 'i', 'l', '\n', '}', '\n'};
1623+
ASSERT_EQ(fwrite(source, 1, sizeof(source), fp), sizeof(source));
1624+
ASSERT_EQ(fclose(fp), 0);
1625+
1626+
char *raw =
1627+
cbm_mcp_handle_tool(srv, "get_code_snippet",
1628+
"{\"qualified_name\":\"test-project.cmd.server.main.HandleRequest\","
1629+
"\"project\":\"test-project\"}");
1630+
ASSERT_TRUE(is_valid_json_response(raw));
1631+
char *resp = extract_text_content(raw);
1632+
ASSERT_NOT_NULL(resp);
1633+
ASSERT_TRUE(is_valid_json_response(resp));
1634+
ASSERT_NULL(strstr(resp, "\xC0\xD4"));
1635+
ASSERT_NOT_NULL(strstr(resp, "HandleRequest"));
1636+
ASSERT_NOT_NULL(strstr(resp, "return nil"));
1637+
ASSERT_TRUE(snippet_source_has_replacement(resp));
1638+
1639+
free(resp);
1640+
free(raw);
1641+
cbm_mcp_server_free(srv);
1642+
cleanup_snippet_dir(tmp);
1643+
PASS();
1644+
}
1645+
15801646
/* ══════════════════════════════════════════════════════════════════
15811647
* JSON-RPC PARSING — EDGE CASES
15821648
* ══════════════════════════════════════════════════════════════════ */
@@ -2129,5 +2195,6 @@ SUITE(mcp) {
21292195
RUN_TEST(snippet_auto_resolve_enabled);
21302196
RUN_TEST(snippet_include_neighbors_default);
21312197
RUN_TEST(snippet_include_neighbors_enabled);
2198+
RUN_TEST(snippet_source_invalid_utf8);
21322199
RUN_TEST(tool_bad_project_name_no_overflow_issue235);
21332200
}

0 commit comments

Comments
 (0)