Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 23 additions & 4 deletions src/pg_query_scan.c
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,15 @@ struct yyguts_t
size_t yyleng_r;
};

static void free_scan_tokens(PgQuery__ScanToken **tokens, size_t n_tokens)
{
size_t i;

for (i = 0; i < n_tokens; i++)
free(tokens[i]);
free(tokens);
}

PgQueryScanResult pg_query_scan(const char* input)
{
MemoryContext ctx = NULL;
Expand Down Expand Up @@ -77,6 +86,8 @@ PgQueryScanResult pg_query_scan(const char* input)
scanner_finish(yyscanner);

output_tokens = malloc(sizeof(PgQuery__ScanToken *) * token_count);
if (output_tokens == NULL && token_count > 0)
ereport(ERROR, (errcode(ERRCODE_OUT_OF_MEMORY), errmsg("out of memory")));

/* initialize the flex scanner --- should match raw_parser() */
yyscanner = scanner_init(input, &yyextra, &ScanKeywords, ScanKeywordTokens);
Expand All @@ -91,6 +102,11 @@ PgQueryScanResult pg_query_scan(const char* input)
if (tok == 0) break;

output_tokens[i] = malloc(sizeof(PgQuery__ScanToken));
if (output_tokens[i] == NULL)
{
free_scan_tokens(output_tokens, i);
ereport(ERROR, (errcode(ERRCODE_OUT_OF_MEMORY), errmsg("out of memory")));
}
Comment on lines 104 to +109

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 bug · medium

Error paths leak the token array and yyscanner

When core_yylex raises an error mid-loop — invalid byte sequence, unterminated string, or overlong identifier — the ereport(ERROR) longjmps to PG_CATCH(), which never frees the partially built output_tokens array or finishes the second yyscanner (src/pg_query_scan.c:104-109). In the per-token OOM path the tokens are freed but the flex scanner opened at line 93 is likewise left un-finalized. Because flex scanner buffers and the token array are plain malloc allocations outside the PostgreSQL memory context, pg_query_exit_memory_context cannot reclaim them, so every failed scan leaks the token array plus scanner buffers.

📋 Prompt for AI Agents

In src/pg_query_scan.c, move the PgQuery__ScanToken **output_tokens declaration above PG_TRY() and initialize it to NULL. In the PG_CATCH() block (around lines 151-170), after MemoryContextSwitchTo(parse_context), add a guarded cleanup: if output_tokens != NULL, call free_scan_tokens(output_tokens, token_count) and set output_tokens = NULL; if yyscanner != NULL, call scanner_finish(yyscanner) and set yyscanner = NULL. On the success path set output_tokens = NULL and yyscanner = NULL right after the existing free_scan_tokens and scanner_finish calls (lines 127 and 142) so the catch-block cleanup cannot double-free. This closes the leak where an ereport(ERROR) from core_yylex or the per-token malloc failure longjmps past cleanup, leaking the malloc-allocated token array and the flex scanner buffers.

pg_query__scan_token__init(output_tokens[i]);
output_tokens[i]->start = yylloc;
if (tok == SCONST || tok == USCONST || tok == BCONST || tok == XCONST || tok == IDENT || tok == UIDENT || tok == C_COMMENT) {
Expand All @@ -115,12 +131,15 @@ PgQueryScanResult pg_query_scan(const char* input)
scan_result.tokens = output_tokens;
result.pbuf.len = pg_query__scan_result__get_packed_size(&scan_result);
result.pbuf.data = malloc(result.pbuf.len);
if (result.pbuf.data == NULL)
{
result.pbuf.len = 0;
free_scan_tokens(output_tokens, token_count);
ereport(ERROR, (errcode(ERRCODE_OUT_OF_MEMORY), errmsg("out of memory")));
}
pg_query__scan_result__pack(&scan_result, (void*) result.pbuf.data);

for (i = 0; i < token_count; i++) {
free(output_tokens[i]);
}
free(output_tokens);
free_scan_tokens(output_tokens, token_count);

#ifndef DEBUG
// Save stderr for result
Expand Down
57 changes: 57 additions & 0 deletions test/scan.c
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,58 @@

#include "protobuf/pg_query.pb-c.h"

#if defined(__linux__) && !defined(USE_VALGRIND)
#include <sys/resource.h>
#include <unistd.h>

/* Scanning under a tight address-space limit must report "out of memory"
* instead of writing tokens through a NULL pointer. */
static bool test_scan_out_of_memory(void)
{
const char *stmt = "SELECT 1;\n";
const size_t repeat = 2000000;
size_t len = strlen(stmt);
char *input = malloc(len * repeat + 1);
long pages = 0;
FILE *statm;
struct rlimit old_limit;
struct rlimit limit;
PgQueryScanResult result;
bool ok;
size_t i;

for (i = 0; i < repeat; i++)
memcpy(input + i * len, stmt, len);
input[len * repeat] = '\0';

statm = fopen("/proc/self/statm", "r");
if (statm == NULL || fscanf(statm, "%ld", &pages) != 1) {
printf("INVALID result for out of memory scan: cannot read /proc/self/statm\n");
free(input);
return false;
}
fclose(statm);

getrlimit(RLIMIT_AS, &old_limit);
limit = old_limit;
limit.rlim_cur = (rlim_t) pages * sysconf(_SC_PAGESIZE) + 128 * 1024 * 1024;
setrlimit(RLIMIT_AS, &limit);
result = pg_query_scan(input);
setrlimit(RLIMIT_AS, &old_limit);

ok = result.error != NULL && strcmp(result.error->message, "out of memory") == 0;
if (ok) {
printf(".");
} else {
printf("INVALID result for out of memory scan: %s\n", result.error ? result.error->message : "no error");
}

pg_query_free_scan_result(result);
free(input);
return ok;
}
#endif

int main() {
size_t i;
size_t j;
Expand Down Expand Up @@ -61,6 +113,11 @@ int main() {
pg_query_free_scan_result(result);
}

#if defined(__linux__) && !defined(USE_VALGRIND)
if (!test_scan_out_of_memory())
ret_code = -1;
#endif

printf("\n");

pg_query_exit();
Expand Down
Loading