Skip to content

Commit f82c6ae

Browse files
DeusDataDeusData
authored
Merge pull request #526 from jstar0/fix/snippet-valid-utf8
fix(mcp): return valid UTF-8 snippets
2 parents 96c1fc3 + 935027a commit f82c6ae

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
@@ -3008,6 +3008,75 @@ static char *resolve_snippet_source(const char *root_path, const char *file_path
30083008
return NULL;
30093009
}
30103010

3011+
static bool utf8_is_cont(unsigned char c) {
3012+
return (c & 0xC0) == 0x80;
3013+
}
3014+
3015+
static char *sanitize_utf8_lossy(const char *s) {
3016+
enum {
3017+
UTF8_REPLACEMENT_LEN = 3,
3018+
UTF8_THREE_BYTE_LEN = 3,
3019+
UTF8_FOUR_BYTE_LEN = 4,
3020+
UTF8_FOURTH_BYTE = 3,
3021+
};
3022+
if (!s) {
3023+
return NULL;
3024+
}
3025+
size_t len = strlen(s);
3026+
if (len > (((size_t)-1) - SKIP_ONE) / UTF8_REPLACEMENT_LEN) {
3027+
return NULL;
3028+
}
3029+
char *out = malloc(len * UTF8_REPLACEMENT_LEN + SKIP_ONE);
3030+
if (!out) {
3031+
return NULL;
3032+
}
3033+
3034+
const unsigned char *p = (const unsigned char *)s;
3035+
const unsigned char *end = p + len;
3036+
unsigned char *dst = (unsigned char *)out;
3037+
while (p < end) {
3038+
unsigned char c = *p;
3039+
size_t n = 0;
3040+
if (c < 0x80) {
3041+
n = 1;
3042+
} else if (c >= 0xC2 && c <= 0xDF && p + 1 < end && utf8_is_cont(p[1])) {
3043+
n = 2;
3044+
} else if (c == 0xE0 && p + 2 < end && p[1] >= 0xA0 && p[1] <= 0xBF && utf8_is_cont(p[2])) {
3045+
n = UTF8_THREE_BYTE_LEN;
3046+
} else if (c >= 0xE1 && c <= 0xEC && p + 2 < end && utf8_is_cont(p[1]) &&
3047+
utf8_is_cont(p[2])) {
3048+
n = UTF8_THREE_BYTE_LEN;
3049+
} else if (c == 0xED && p + 2 < end && p[1] >= 0x80 && p[1] <= 0x9F && utf8_is_cont(p[2])) {
3050+
n = UTF8_THREE_BYTE_LEN;
3051+
} else if (c >= 0xEE && c <= 0xEF && p + 2 < end && utf8_is_cont(p[1]) &&
3052+
utf8_is_cont(p[2])) {
3053+
n = UTF8_THREE_BYTE_LEN;
3054+
} else if (c == 0xF0 && p + UTF8_FOURTH_BYTE < end && p[1] >= 0x90 && p[1] <= 0xBF &&
3055+
utf8_is_cont(p[2]) && utf8_is_cont(p[UTF8_FOURTH_BYTE])) {
3056+
n = UTF8_FOUR_BYTE_LEN;
3057+
} else if (c >= 0xF1 && c <= 0xF3 && p + UTF8_FOURTH_BYTE < end && utf8_is_cont(p[1]) &&
3058+
utf8_is_cont(p[2]) && utf8_is_cont(p[UTF8_FOURTH_BYTE])) {
3059+
n = UTF8_FOUR_BYTE_LEN;
3060+
} else if (c == 0xF4 && p + UTF8_FOURTH_BYTE < end && p[1] >= 0x80 && p[1] <= 0x8F &&
3061+
utf8_is_cont(p[2]) && utf8_is_cont(p[UTF8_FOURTH_BYTE])) {
3062+
n = UTF8_FOUR_BYTE_LEN;
3063+
}
3064+
3065+
if (n > 0) {
3066+
memcpy(dst, p, n);
3067+
dst += n;
3068+
p += n;
3069+
} else {
3070+
*dst++ = 0xEF;
3071+
*dst++ = 0xBF;
3072+
*dst++ = 0xBD;
3073+
p++;
3074+
}
3075+
}
3076+
*dst = '\0';
3077+
return out;
3078+
}
3079+
30113080
/* Build an enriched snippet response for a resolved node. */
30123081
/* Add a string array to a JSON object (no-op if count == 0). */
30133082
static void add_string_array(yyjson_mut_doc *doc, yyjson_mut_val *obj, const char *key,
@@ -3052,7 +3121,13 @@ static char *build_snippet_response(cbm_mcp_server_t *srv, cbm_node_t *node,
30523121
yyjson_mut_obj_add_int(doc, root_obj, "end_line", end);
30533122

30543123
if (source) {
3055-
yyjson_mut_obj_add_str(doc, root_obj, "source", source);
3124+
char *safe_source = sanitize_utf8_lossy(source);
3125+
if (safe_source) {
3126+
yyjson_mut_obj_add_strcpy(doc, root_obj, "source", safe_source);
3127+
free(safe_source);
3128+
} else {
3129+
yyjson_mut_obj_add_str(doc, root_obj, "source", "(source not available)");
3130+
}
30563131
} else {
30573132
yyjson_mut_obj_add_str(doc, root_obj, "source", "(source not available)");
30583133
}

‎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
@@ -1444,6 +1445,31 @@ static char *call_snippet(cbm_mcp_server_t *srv, const char *args_json) {
14441445
return text;
14451446
}
14461447

1448+
static bool is_valid_json_response(const char *json) {
1449+
if (!json) {
1450+
return false;
1451+
}
1452+
yyjson_doc *doc = yyjson_read(json, strlen(json), 0);
1453+
if (!doc) {
1454+
return false;
1455+
}
1456+
yyjson_doc_free(doc);
1457+
return true;
1458+
}
1459+
1460+
static bool snippet_source_has_replacement(const char *json) {
1461+
yyjson_doc *doc = yyjson_read(json, strlen(json), 0);
1462+
if (!doc) {
1463+
return false;
1464+
}
1465+
yyjson_val *root = yyjson_doc_get_root(doc);
1466+
yyjson_val *source = yyjson_obj_get(root, "source");
1467+
const char *source_str = yyjson_get_str(source);
1468+
bool found = source_str && strstr(source_str, "\xEF\xBF\xBD");
1469+
yyjson_doc_free(doc);
1470+
return found;
1471+
}
1472+
14471473
/* ── TestSnippet_ExactQN ──────────────────────────────────────── */
14481474

14491475
TEST(snippet_exact_qn) {
@@ -1730,6 +1756,46 @@ TEST(snippet_include_neighbors_enabled) {
17301756
PASS();
17311757
}
17321758

1759+
/* ── TestSnippet_SourceInvalidUtf8 ────────────────────────────── */
1760+
1761+
TEST(snippet_source_invalid_utf8) {
1762+
char tmp[256];
1763+
cbm_mcp_server_t *srv = setup_snippet_server(tmp, sizeof(tmp));
1764+
ASSERT_NOT_NULL(srv);
1765+
1766+
char src_path[512];
1767+
snprintf(src_path, sizeof(src_path), "%s/project/main.go", tmp);
1768+
FILE *fp = fopen(src_path, "wb");
1769+
ASSERT_NOT_NULL(fp);
1770+
const unsigned char source[] = {
1771+
'p', 'a', 'c', 'k', 'a', 'g', 'e', ' ', 'm', 'a', 'i', 'n', '\n', '\n',
1772+
'f', 'u', 'n', 'c', ' ', 'H', 'a', 'n', 'd', 'l', 'e', 'R', 'e', 'q',
1773+
'u', 'e', 's', 't', '(', ')', ' ', 'e', 'r', 'r', 'o', 'r', ' ', '{',
1774+
'\n', '\t', '/', '/', ' ', 0xC0, 0xD4, 0xB7, 0xC2, '\n', '\t', 'r', 'e', 't',
1775+
'u', 'r', 'n', ' ', 'n', 'i', 'l', '\n', '}', '\n'};
1776+
ASSERT_EQ(fwrite(source, 1, sizeof(source), fp), sizeof(source));
1777+
ASSERT_EQ(fclose(fp), 0);
1778+
1779+
char *raw =
1780+
cbm_mcp_handle_tool(srv, "get_code_snippet",
1781+
"{\"qualified_name\":\"test-project.cmd.server.main.HandleRequest\","
1782+
"\"project\":\"test-project\"}");
1783+
ASSERT_TRUE(is_valid_json_response(raw));
1784+
char *resp = extract_text_content(raw);
1785+
ASSERT_NOT_NULL(resp);
1786+
ASSERT_TRUE(is_valid_json_response(resp));
1787+
ASSERT_NULL(strstr(resp, "\xC0\xD4"));
1788+
ASSERT_NOT_NULL(strstr(resp, "HandleRequest"));
1789+
ASSERT_NOT_NULL(strstr(resp, "return nil"));
1790+
ASSERT_TRUE(snippet_source_has_replacement(resp));
1791+
1792+
free(resp);
1793+
free(raw);
1794+
cbm_mcp_server_free(srv);
1795+
cleanup_snippet_dir(tmp);
1796+
PASS();
1797+
}
1798+
17331799
/* ══════════════════════════════════════════════════════════════════
17341800
* JSON-RPC PARSING — EDGE CASES
17351801
* ══════════════════════════════════════════════════════════════════ */
@@ -2285,5 +2351,6 @@ SUITE(mcp) {
22852351
RUN_TEST(snippet_auto_resolve_enabled);
22862352
RUN_TEST(snippet_include_neighbors_default);
22872353
RUN_TEST(snippet_include_neighbors_enabled);
2354+
RUN_TEST(snippet_source_invalid_utf8);
22882355
RUN_TEST(tool_bad_project_name_no_overflow_issue235);
22892356
}

0 commit comments

Comments
 (0)