From a0ac81ae4ec160726b1004e4f68ba83114265753 Mon Sep 17 00:00:00 2001 From: Dan Lynch Date: Thu, 17 Sep 2026 07:20:54 +0000 Subject: [PATCH] fix(scan): size build_scan_json buffer from token text length, never write past it Fixes #167 --- templates/full/wasm_wrapper.c | 37 ++++++++++++++++++++-------------- versions/18/src/wasm_wrapper.c | 37 ++++++++++++++++++++-------------- versions/18/test/scan.test.js | 28 +++++++++++++++++++++++++ 3 files changed, 72 insertions(+), 30 deletions(-) diff --git a/templates/full/wasm_wrapper.c b/templates/full/wasm_wrapper.c index 3844d960..268d2893 100644 --- a/templates/full/wasm_wrapper.c +++ b/templates/full/wasm_wrapper.c @@ -271,15 +271,23 @@ static char* build_scan_json(PgQuery__ScanResult *scan_result, const char* origi return safe_strdup("{\"version\":0,\"tokens\":[]}"); } - // Calculate rough JSON size estimate - size_t estimated_size = 1024 + (scan_result->n_tokens * 200); + // Size the buffer from token text length, not token count: each token + // contributes at most ~170 bytes of fixed JSON plus up to 2 bytes per input + // byte once escaped. + size_t estimated_size = 1024; + for (size_t i = 0; i < scan_result->n_tokens; i++) { + PgQuery__ScanToken *token = scan_result->tokens[i]; + int token_length = token->end - token->start; + if (token_length < 0) token_length = 0; + estimated_size += (size_t) token_length * 2 + 200; + } char* json = safe_malloc(estimated_size); if (!json) { return safe_strdup("{\"version\":0,\"tokens\":[]}"); } // Start building JSON - int pos = snprintf(json, estimated_size, "{\"version\":%d,\"tokens\":[", scan_result->version); + size_t pos = snprintf(json, estimated_size, "{\"version\":%d,\"tokens\":[", scan_result->version); for (size_t i = 0; i < scan_result->n_tokens; i++) { PgQuery__ScanToken *token = scan_result->tokens[i]; @@ -329,25 +337,24 @@ static char* build_scan_json(PgQuery__ScanResult *scan_result, const char* origi const char* keyword_name = get_keyword_name(token->keyword_kind); // Add comma if not first token - if (i > 0) { - pos += snprintf(json + pos, estimated_size - pos, ","); + if (i > 0 && pos < estimated_size) { + int n = snprintf(json + pos, estimated_size - pos, ","); + pos += (n > 0) ? (size_t) n : 0; } // Add token object to JSON - pos += snprintf(json + pos, estimated_size - pos, - "{\"start\":%d,\"end\":%d,\"text\":\"%s\",\"tokenType\":%d,\"tokenName\":\"%s\",\"keywordKind\":%d,\"keywordName\":\"%s\"}", - token->start, token->end, escaped_text, token->token, token_name, token->keyword_kind, keyword_name); + if (pos < estimated_size) { + int n = snprintf(json + pos, estimated_size - pos, + "{\"start\":%d,\"end\":%d,\"text\":\"%s\",\"tokenType\":%d,\"tokenName\":\"%s\",\"keywordKind\":%d,\"keywordName\":\"%s\"}", + token->start, token->end, escaped_text, token->token, token_name, token->keyword_kind, keyword_name); + pos += (n > 0) ? (size_t) n : 0; + } free(token_text); free(escaped_text); - // Check if we're running out of space - if (pos >= estimated_size - 200) { - char* new_json = realloc(json, estimated_size * 2); - if (!new_json) break; - json = new_json; - estimated_size *= 2; - } + // snprintf returns the would-be length; never let pos run past the buffer + if (pos >= estimated_size) pos = estimated_size - 1; } // Close JSON diff --git a/versions/18/src/wasm_wrapper.c b/versions/18/src/wasm_wrapper.c index ca6ea740..63ae5dab 100644 --- a/versions/18/src/wasm_wrapper.c +++ b/versions/18/src/wasm_wrapper.c @@ -278,15 +278,23 @@ static char* build_scan_json(PgQuery__ScanResult *scan_result, const char* origi return safe_strdup("{\"version\":0,\"tokens\":[]}"); } - // Calculate rough JSON size estimate - size_t estimated_size = 1024 + (scan_result->n_tokens * 200); + // Size the buffer from token text length, not token count: each token + // contributes at most ~170 bytes of fixed JSON plus up to 2 bytes per input + // byte once escaped. + size_t estimated_size = 1024; + for (size_t i = 0; i < scan_result->n_tokens; i++) { + PgQuery__ScanToken *token = scan_result->tokens[i]; + int token_length = token->end - token->start; + if (token_length < 0) token_length = 0; + estimated_size += (size_t) token_length * 2 + 200; + } char* json = safe_malloc(estimated_size); if (!json) { return safe_strdup("{\"version\":0,\"tokens\":[]}"); } // Start building JSON - int pos = snprintf(json, estimated_size, "{\"version\":%d,\"tokens\":[", scan_result->version); + size_t pos = snprintf(json, estimated_size, "{\"version\":%d,\"tokens\":[", scan_result->version); for (size_t i = 0; i < scan_result->n_tokens; i++) { PgQuery__ScanToken *token = scan_result->tokens[i]; @@ -336,25 +344,24 @@ static char* build_scan_json(PgQuery__ScanResult *scan_result, const char* origi const char* keyword_name = get_keyword_name(token->keyword_kind); // Add comma if not first token - if (i > 0) { - pos += snprintf(json + pos, estimated_size - pos, ","); + if (i > 0 && pos < estimated_size) { + int n = snprintf(json + pos, estimated_size - pos, ","); + pos += (n > 0) ? (size_t) n : 0; } // Add token object to JSON - pos += snprintf(json + pos, estimated_size - pos, - "{\"start\":%d,\"end\":%d,\"text\":\"%s\",\"tokenType\":%d,\"tokenName\":\"%s\",\"keywordKind\":%d,\"keywordName\":\"%s\"}", - token->start, token->end, escaped_text, token->token, token_name, token->keyword_kind, keyword_name); + if (pos < estimated_size) { + int n = snprintf(json + pos, estimated_size - pos, + "{\"start\":%d,\"end\":%d,\"text\":\"%s\",\"tokenType\":%d,\"tokenName\":\"%s\",\"keywordKind\":%d,\"keywordName\":\"%s\"}", + token->start, token->end, escaped_text, token->token, token_name, token->keyword_kind, keyword_name); + pos += (n > 0) ? (size_t) n : 0; + } free(token_text); free(escaped_text); - // Check if we're running out of space - if (pos >= estimated_size - 200) { - char* new_json = realloc(json, estimated_size * 2); - if (!new_json) break; - json = new_json; - estimated_size *= 2; - } + // snprintf returns the would-be length; never let pos run past the buffer + if (pos >= estimated_size) pos = estimated_size - 1; } // Close JSON diff --git a/versions/18/test/scan.test.js b/versions/18/test/scan.test.js index 24d7b4cf..bd927874 100644 --- a/versions/18/test/scan.test.js +++ b/versions/18/test/scan.test.js @@ -275,5 +275,33 @@ comment */ SELECT 2`; assert.ok(commentToken, "should have a C_COMMENT token"); assert.ok(commentToken.text.includes('\n'), "comment text should preserve newlines"); }); + + it("should handle tokens far larger than the per-token JSON budget", () => { + // A single long literal used to overflow the count-based buffer estimate, + // truncating the JSON output (and writing past the buffer). + for (const length of [1400, 5000, 300000]) { + const literal = "a".repeat(length); + const sql = `SELECT '${literal}'; SELECT 1; SELECT 2; SELECT 3;`; + const result = query.scanSync(sql); + + const sconst = result.tokens.find(t => t.tokenName === "SCONST"); + assert.ok(sconst, `should have an SCONST token (length ${length})`); + assert.equal(sconst.text, `'${literal}'`); + assert.equal(result.tokens[result.tokens.length - 1].text, ";"); + assert.equal(result.tokens.filter(t => t.text === "SELECT").length, 4); + } + }); + + it("should handle many long tokens that need JSON escaping", () => { + const literal = ('"\\\n\t').repeat(2000); + const sql = `SELECT $$${literal}$$, $$${literal}$$`; + const result = query.scanSync(sql); + + const dollarTokens = result.tokens.filter(t => t.text.startsWith("$$")); + assert.equal(dollarTokens.length, 2); + for (const t of dollarTokens) { + assert.equal(t.text, `$$${literal}$$`); + } + }); }); });