diff --git a/Makefile b/Makefile index a36f6119..4bec02ba 100644 --- a/Makefile +++ b/Makefile @@ -237,7 +237,7 @@ examples/normalize_error: examples/normalize_error.c $(ARLIB) examples/simple_plpgsql: examples/simple_plpgsql.c $(ARLIB) $(CC) $(TEST_CFLAGS) -o $@ -g examples/simple_plpgsql.c $(ARLIB) $(TEST_LDFLAGS) -TESTS = test/complex test/concurrency test/deparse test/fingerprint test/fingerprint_opts test/is_utility_stmt test/normalize test/normalize_utility test/parse test/parse_opts test/parse_protobuf test/parse_protobuf_opts test/parse_plpgsql test/scan test/split test/summary test/summary_truncate +TESTS = test/complex test/concurrency test/deparse test/fingerprint test/fingerprint_opts test/is_utility_stmt test/normalize test/normalize_utility test/parse test/parse_opts test/parse_protobuf test/parse_protobuf_opts test/parse_plpgsql test/scan test/split test/stack_depth test/summary test/summary_truncate test: $(TESTS) ifeq ($(VALGRIND),1) $(VALGRIND_MEMCHECK) test/complex || (cat test/valgrind.log && false) @@ -254,6 +254,7 @@ ifeq ($(VALGRIND),1) $(VALGRIND_MEMCHECK) test/parse_protobuf_opts || (cat test/valgrind.log && false) $(VALGRIND_MEMCHECK) test/scan || (cat test/valgrind.log && false) $(VALGRIND_MEMCHECK) test/split || (cat test/valgrind.log && false) + $(VALGRIND_MEMCHECK) test/stack_depth || (cat test/valgrind.log && false) $(VALGRIND_MEMCHECK) test/summary || (cat test/valgrind.log && false) $(VALGRIND_MEMCHECK) test/summary_truncate || (cat test/valgrind.log && false) # Output-based tests @@ -274,6 +275,7 @@ else test/parse_protobuf_opts test/scan test/split + test/stack_depth test/summary test/summary_truncate # Output-based tests @@ -335,6 +337,9 @@ test/parse_protobuf_opts: test/parse_protobuf_opts.c test/parse_opts_tests.c $(A test/scan: test/scan.c test/scan_tests.c $(ARLIB) $(CC) $(TEST_CFLAGS) -o $@ test/scan.c $(ARLIB) $(TEST_LDFLAGS) +test/stack_depth: test/stack_depth.c $(ARLIB) + $(CC) $(TEST_CFLAGS) -o $@ test/stack_depth.c $(ARLIB) $(TEST_LDFLAGS) + test/split: test/split.c test/split_tests.c $(ARLIB) $(CC) $(TEST_CFLAGS) -o $@ test/split.c $(ARLIB) $(TEST_LDFLAGS) diff --git a/pg_query.h b/pg_query.h index def0254e..210ce18f 100644 --- a/pg_query.h +++ b/pg_query.h @@ -161,6 +161,11 @@ void pg_query_free_summary_parse_result(PgQuerySummaryParseResult result); // Optional, cleans up the top-level memory context (automatically done for threads that exit) void pg_query_exit(void); +// Optional, makes deeply nested input fail with "stack depth limit exceeded" +// instead of overflowing the native stack. Allows `bytes` of stack below the +// caller's frame, so call it from an outer frame. The setting is per thread. +void pg_query_set_stack_limit(size_t bytes); + // Postgres version information #define PG_MAJORVERSION "17" #define PG_VERSION "17.7" diff --git a/src/pg_query_fingerprint.c b/src/pg_query_fingerprint.c index 8e15b821..5679c8e9 100644 --- a/src/pg_query_fingerprint.c +++ b/src/pg_query_fingerprint.c @@ -5,6 +5,8 @@ #include "pg_query_fingerprint.h" #include "postgres.h" +#include "miscadmin.h" +#define XXH_STATIC_LINKING_ONLY #include "xxhash/xxhash.h" #include "lib/ilist.h" @@ -225,6 +227,33 @@ _fingerprintList(FingerprintContext *ctx, const List *node, const void *parent, } } +/* + * Hash states come from the current memory context rather than malloc, so an + * error thrown mid-fingerprint (e.g. stack depth exceeded) cannot leak them. + * Keeps XXH3's 64-byte alignment the same way XXH_alignedMalloc does. + */ +static XXH3_state_t * +_fingerprintCreateState(void) +{ + char *base = palloc0(sizeof(XXH3_state_t) + 64); + size_t offset = 64 - ((uintptr_t) base & 63); + char *ptr = base + offset; + + ptr[-1] = (char) offset; + return (XXH3_state_t *) ptr; +} + +static void +_fingerprintFreeState(XXH3_state_t *state) +{ + char *ptr = (char *) state; + + pfree(ptr - (unsigned char) ptr[-1]); +} + +#define XXH3_createState _fingerprintCreateState +#define XXH3_freeState _fingerprintFreeState + static void _fingerprintInitContext(FingerprintContext *ctx, FingerprintContext *parent, bool write_tokens) { @@ -263,6 +292,7 @@ _fingerprintFreeContext(FingerprintContext *ctx) { void _fingerprintNode(FingerprintContext *ctx, const void *obj, const void *parent, char *field_name, unsigned int depth) { + check_stack_depth(); // Some queries are overly complex in their parsetree - lets consistently cut them off at 100 nodes deep if (depth >= 100) { return; @@ -327,10 +357,12 @@ uint64_t pg_query_fingerprint_node(const void *node) PgQueryFingerprintResult pg_query_fingerprint_with_opts(const char* input, int parser_options, bool printTokens) { MemoryContext ctx = NULL; + MemoryContext mem_ctx; PgQueryInternalParsetreeAndError parsetree_and_error; PgQueryFingerprintResult result = {0}; ctx = pg_query_enter_memory_context(); + mem_ctx = ctx; parsetree_and_error = pg_query_raw_parse(input, parser_options); @@ -345,7 +377,34 @@ PgQueryFingerprintResult pg_query_fingerprint_with_opts(const char* input, int p _fingerprintInitContext(&ctx, NULL, printTokens); if (parsetree_and_error.tree != NULL) { - _fingerprintNode(&ctx, parsetree_and_error.tree, NULL, NULL, 0); + PG_TRY(); + { + _fingerprintNode(&ctx, parsetree_and_error.tree, NULL, NULL, 0); + } + PG_CATCH(); + { + ErrorData* error_data; + PgQueryError* error; + + MemoryContextSwitchTo(mem_ctx); + error_data = CopyErrorData(); + error = malloc(sizeof(PgQueryError)); + error->message = strdup(error_data->message); + error->filename = strdup(error_data->filename); + error->funcname = strdup(error_data->funcname); + error->context = NULL; + error->lineno = error_data->lineno; + error->cursorpos = error_data->cursorpos; + result.error = error; + FlushErrorState(); + } + PG_END_TRY(); + + if (result.error != NULL) { + _fingerprintFreeContext(&ctx); + pg_query_exit_memory_context(mem_ctx); + return result; + } } if (printTokens) { diff --git a/src/pg_query_json_plpgsql.c b/src/pg_query_json_plpgsql.c index 349068d6..d4015318 100644 --- a/src/pg_query_json_plpgsql.c +++ b/src/pg_query_json_plpgsql.c @@ -1,4 +1,6 @@ #include "pg_query.h" +#include "postgres.h" +#include "miscadmin.h" #include "pg_query_json_plpgsql.h" #include "pg_query_json_helper.c" @@ -138,6 +140,7 @@ static void dump_type(StringInfo out, PLpgSQL_type *node); static void dump_stmt(StringInfo out, PLpgSQL_stmt *node) { + check_stack_depth(); appendStringInfoChar(out, '{'); switch (node->cmd_type) { @@ -625,6 +628,7 @@ dump_getdiag_item(StringInfo out, PLpgSQL_diag_item *node) static void dump_expr(StringInfo out, PLpgSQL_expr *node) { + check_stack_depth(); WRITE_NODE_TYPE("PLpgSQL_expr"); WRITE_STRING_FIELD(query, query, query); @@ -806,7 +810,10 @@ pg_query_plpgsql_reset_aliases(void) int i; for (i = 0; i < plpgsql_naliases; i++) free(plpgsql_aliases[i].name); + free(plpgsql_aliases); + plpgsql_aliases = NULL; plpgsql_naliases = 0; + plpgsql_aliases_size = 0; } void diff --git a/src/pg_query_outfuncs_json.c b/src/pg_query_outfuncs_json.c index d4d74bff..318af000 100644 --- a/src/pg_query_outfuncs_json.c +++ b/src/pg_query_outfuncs_json.c @@ -1,6 +1,7 @@ #include "pg_query_outfuncs.h" #include "postgres.h" +#include "miscadmin.h" #include @@ -282,6 +283,7 @@ _outAConst(StringInfo out, const A_Const *node) static void _outNode(StringInfo out, const void *obj) { + check_stack_depth(); if (obj == NULL) { appendStringInfoString(out, "null"); diff --git a/src/pg_query_outfuncs_protobuf.c b/src/pg_query_outfuncs_protobuf.c index 64ccdf1b..c7bcd34a 100644 --- a/src/pg_query_outfuncs_protobuf.c +++ b/src/pg_query_outfuncs_protobuf.c @@ -1,6 +1,7 @@ #include "pg_query_outfuncs.h" #include "postgres.h" +#include "miscadmin.h" #include #include "access/relation.h" #include "nodes/parsenodes.h" @@ -246,6 +247,7 @@ _outAConst(PgQuery__AConst* out, const A_Const *node) static void _outNode(PgQuery__Node* out, const void *obj) { + check_stack_depth(); if (obj == NULL) return; // Keep out as NULL diff --git a/src/pg_query_parse.c b/src/pg_query_parse.c index 3f2c111a..31f851b5 100644 --- a/src/pg_query_parse.c +++ b/src/pg_query_parse.c @@ -135,9 +135,39 @@ PgQueryParseResult pg_query_parse_opts(const char* input, int parser_options) result.stderr_buffer = parsetree_and_error.stderr_buffer; result.error = parsetree_and_error.error; - tree_json = pg_query_nodes_to_json(parsetree_and_error.tree); - result.parse_tree = strdup(tree_json); - pfree(tree_json); + if (result.error == NULL) + { + PG_TRY(); + { + tree_json = pg_query_nodes_to_json(parsetree_and_error.tree); + result.parse_tree = strdup(tree_json); + pfree(tree_json); + } + PG_CATCH(); + { + ErrorData* error_data; + PgQueryError* error; + + MemoryContextSwitchTo(ctx); + error_data = CopyErrorData(); + error = malloc(sizeof(PgQueryError)); + error->message = strdup(error_data->message); + error->filename = strdup(error_data->filename); + error->funcname = strdup(error_data->funcname); + error->context = NULL; + error->lineno = error_data->lineno; + error->cursorpos = error_data->cursorpos; + result.error = error; + FlushErrorState(); + } + PG_END_TRY(); + } + else + { + tree_json = pg_query_nodes_to_json(parsetree_and_error.tree); + result.parse_tree = strdup(tree_json); + pfree(tree_json); + } pg_query_exit_memory_context(ctx); @@ -162,7 +192,28 @@ PgQueryProtobufParseResult pg_query_parse_protobuf_opts(const char* input, int p // These are all malloc-ed and will survive exiting the memory context, the caller is responsible to free them now result.stderr_buffer = parsetree_and_error.stderr_buffer; result.error = parsetree_and_error.error; - result.parse_tree = pg_query_nodes_to_protobuf(parsetree_and_error.tree); + PG_TRY(); + { + result.parse_tree = pg_query_nodes_to_protobuf(parsetree_and_error.tree); + } + PG_CATCH(); + { + ErrorData* error_data; + PgQueryError* error; + + MemoryContextSwitchTo(ctx); + error_data = CopyErrorData(); + error = malloc(sizeof(PgQueryError)); + error->message = strdup(error_data->message); + error->filename = strdup(error_data->filename); + error->funcname = strdup(error_data->funcname); + error->context = NULL; + error->lineno = error_data->lineno; + error->cursorpos = error_data->cursorpos; + result.error = error; + FlushErrorState(); + } + PG_END_TRY(); pg_query_exit_memory_context(ctx); diff --git a/src/pg_query_parse_plpgsql.c b/src/pg_query_parse_plpgsql.c index 3c169343..69516e1b 100644 --- a/src/pg_query_parse_plpgsql.c +++ b/src/pg_query_parse_plpgsql.c @@ -489,9 +489,38 @@ PgQueryPlpgsqlParseResult pg_query_parse_plpgsql(const char* input) char *new_out; size_t new_out_len; - func_json = plpgsqlToJSON(func_and_error.func); + PG_TRY(); + { + func_json = plpgsqlToJSON(func_and_error.func); + } + PG_CATCH(); + { + ErrorData* error_data; + PgQueryError* error; + + MemoryContextSwitchTo(ctx); + error_data = CopyErrorData(); + error = malloc(sizeof(PgQueryError)); + error->message = strdup(error_data->message); + error->filename = strdup(error_data->filename); + error->funcname = strdup(error_data->funcname); + error->context = NULL; + error->lineno = error_data->lineno; + error->cursorpos = error_data->cursorpos; + result.error = error; + FlushErrorState(); + } + PG_END_TRY(); + plpgsql_free_function_memory(func_and_error.func); + if (result.error != NULL) { + free(parse_result.stderr_buffer); + pg_query_plpgsql_reset_aliases(); + pg_query_exit_memory_context(ctx); + return result; + } + new_out_len = strlen(result.plpgsql_funcs) + strlen(func_json) + 3; new_out = malloc(new_out_len); int n = snprintf(new_out, new_out_len, "%s%s,\n", result.plpgsql_funcs, func_json); @@ -512,6 +541,8 @@ PgQueryPlpgsqlParseResult pg_query_parse_plpgsql(const char* input) result.plpgsql_funcs[strlen(result.plpgsql_funcs) - 1] = ']'; free(parse_result.stderr_buffer); + /* Drop the aliases recorded during this compile so nothing outlives the call */ + pg_query_plpgsql_reset_aliases(); pg_query_exit_memory_context(ctx); return result; diff --git a/src/postgres/src_backend_tcop_postgres.c b/src/postgres/src_backend_tcop_postgres.c index 1af69bd7..62cae168 100644 --- a/src/postgres/src_backend_tcop_postgres.c +++ b/src/postgres/src_backend_tcop_postgres.c @@ -826,3 +826,14 @@ stack_is_too_deep(void) * Disable statement timeout, if active. */ + +/* libpg-query-node: set the stack-depth reference point and limit from the embedder. */ +void +pg_query_set_stack_limit(size_t bytes) +{ + char here; + + stack_base_ptr = &here; + max_stack_depth_bytes = (long) bytes; + max_stack_depth = (int) (bytes / 1024); +} diff --git a/src/postgres_deparse.c b/src/postgres_deparse.c index d4d8237a..139b5144 100644 --- a/src/postgres_deparse.c +++ b/src/postgres_deparse.c @@ -1,6 +1,7 @@ #include "postgres_deparse.h" #include "postgres.h" +#include "miscadmin.h" #include "catalog/index.h" #include "catalog/pg_am.h" #include "catalog/pg_attribute.h" @@ -783,6 +784,7 @@ static void deparseCExpr(DeparseState *state, Node *node); // "a_expr" in gram.y static void deparseExpr(DeparseState *state, Node *node, DeparseNodeContext context) { + check_stack_depth(); if (node == NULL) return; switch (nodeTag(node)) diff --git a/test/stack_depth.c b/test/stack_depth.c new file mode 100644 index 00000000..98df18ef --- /dev/null +++ b/test/stack_depth.c @@ -0,0 +1,127 @@ +#include + +#include +#include +#include +#include + +/* With pg_query_set_stack_limit() configured, input nested deeper than the + * limit must come back as a "stack depth limit exceeded" error. Before these + * paths were wrapped in PG_TRY, the error had no handler, so it was promoted + * to FATAL and the whole process exited. */ + +#define STACK_LIMIT (16 * 1024) + +static char *repeat_into(char *out, const char *s, size_t n) +{ + size_t len = strlen(s); + size_t i; + + for (i = 0; i < n; i++, out += len) + memcpy(out, s, len); + *out = '\0'; + return out; +} + +static char *copy_into(char *out, const char *s) +{ + strcpy(out, s); + return out + strlen(s); +} + +static char *deep_expression(size_t depth) +{ + char *sql = malloc(strlen("SELECT 1") + depth * 2 + 1); + repeat_into(copy_into(sql, "SELECT 1"), "+1", depth); + return sql; +} + +static char *deep_plpgsql(size_t depth) +{ + const char *head = "CREATE FUNCTION f() RETURNS void LANGUAGE plpgsql AS $$ BEGIN "; + const char *tail = "END $$;"; + char *sql = malloc(strlen(head) + depth * (strlen("IF true THEN ") + strlen("END IF; ")) + strlen("NULL; ") + strlen(tail) + 1); + char *p = copy_into(sql, head); + + p = repeat_into(p, "IF true THEN ", depth); + p = copy_into(p, "NULL; "); + p = repeat_into(p, "END IF; ", depth); + strcpy(p, tail); + return sql; +} + +static bool expect_stack_error(const char *name, PgQueryError *error) +{ + if (error != NULL && strstr(error->message, "stack depth limit exceeded") != NULL) { + printf("."); + return true; + } + printf("INVALID result for %s: %s\n", name, error ? error->message : "no error"); + return false; +} + +static bool test_fingerprint(void) +{ + char *sql = deep_expression(3000); + PgQueryFingerprintResult result; + bool ok; + + /* Fingerprinting stops at 100 nodes deep, so it needs a tighter limit */ + pg_query_set_stack_limit(4 * 1024); + result = pg_query_fingerprint(sql); + ok = expect_stack_error("deep fingerprint", result.error); + pg_query_free_fingerprint_result(result); + free(sql); + return ok; +} + +static bool test_parse_protobuf(void) +{ + char *sql = deep_expression(3000); + PgQueryProtobufParseResult result; + bool ok; + + pg_query_set_stack_limit(STACK_LIMIT); + result = pg_query_parse_protobuf(sql); + /* The experimental protobuf-cpp serializer has no stack guard and returns the tree */ + if (result.error == NULL && result.parse_tree.len > 0) { + printf("."); + ok = true; + } else { + ok = expect_stack_error("deep protobuf parse", result.error); + } + pg_query_free_protobuf_parse_result(result); + free(sql); + return ok; +} + +static bool test_parse_plpgsql(void) +{ + char *sql = deep_plpgsql(2000); + PgQueryPlpgsqlParseResult result; + bool ok; + + pg_query_set_stack_limit(STACK_LIMIT); + result = pg_query_parse_plpgsql(sql); + ok = expect_stack_error("deep plpgsql parse", result.error); + pg_query_free_plpgsql_parse_result(result); + free(sql); + return ok; +} + +int main(int argc, char **argv) +{ + const char *only = argc > 1 ? argv[1] : NULL; + int ret_code = 0; + + if ((!only || strcmp(only, "fingerprint") == 0) && !test_fingerprint()) + ret_code = 1; + if ((!only || strcmp(only, "protobuf") == 0) && !test_parse_protobuf()) + ret_code = 1; + if ((!only || strcmp(only, "plpgsql") == 0) && !test_parse_plpgsql()) + ret_code = 1; + + printf("\n"); + pg_query_exit(); + return ret_code; +}