diff --git a/Makefile b/Makefile index fdf5e2c2..c56a941e 100644 --- a/Makefile +++ b/Makefile @@ -238,7 +238,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) @@ -255,6 +255,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 @@ -275,6 +276,7 @@ else test/parse_protobuf_opts test/scan test/split + test/stack_depth test/summary test/summary_truncate # Output-based tests @@ -336,6 +338,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 6e103546..92cd2d6f 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 "18" #define PG_VERSION "18.4" 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 ec5134bd..693a05ad 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) { @@ -627,6 +630,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); @@ -812,7 +816,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 e5025ed6..0801ed47 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 @@ -289,6 +290,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 79402515..2088f2d4 100644 --- a/src/pg_query_parse_plpgsql.c +++ b/src/pg_query_parse_plpgsql.c @@ -140,7 +140,10 @@ static PLpgSQL_function *compile_do_stmt(DoStmt* stmt) } } - assert(proc_source != NULL); + if (proc_source == NULL) + ereport(ERROR, + (errcode(ERRCODE_SYNTAX_ERROR), + errmsg("no inline code specified"))); if(strcmp(language, "plpgsql") != 0) { return (PLpgSQL_function *) palloc0(sizeof(PLpgSQL_function)); @@ -431,11 +434,25 @@ compile_create_function_stmt_via_callback(CreateFunctionStmt *stmt) } } - assert(proc_source != NULL); - if (strcmp(language, "plpgsql") != 0) return (PLpgSQL_function *) palloc0(sizeof(PLpgSQL_function)); + /* Mirrors the checks in upstream interpret_AS_clause */ + if (stmt->sql_body == NULL && proc_source == NULL) + ereport(ERROR, + (errcode(ERRCODE_INVALID_FUNCTION_DEFINITION), + errmsg("no function body specified"))); + + if (stmt->sql_body != NULL && proc_source != NULL) + ereport(ERROR, + (errcode(ERRCODE_INVALID_FUNCTION_DEFINITION), + errmsg("duplicate function body specified"))); + + if (stmt->sql_body != NULL) + ereport(ERROR, + (errcode(ERRCODE_INVALID_FUNCTION_DEFINITION), + errmsg("inline SQL function body only valid for language SQL"))); + wrapper = pg_query_create_function(stmt, language, proc_source, &is_dml_trigger, &is_event_trigger); @@ -624,6 +641,7 @@ PgQueryPlpgsqlParseResult pg_query_parse_plpgsql(const char* input) if (statements.stmts_count == 0) { result.plpgsql_funcs = strdup("[]"); + free(parse_result.stderr_buffer); pg_query_exit_memory_context(ctx); return result; } @@ -639,6 +657,7 @@ PgQueryPlpgsqlParseResult pg_query_parse_plpgsql(const char* input) result.error = func_and_error.error; if (result.error != NULL) { + free(parse_result.stderr_buffer); pg_query_exit_memory_context(ctx); return result; } @@ -648,9 +667,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); @@ -671,6 +719,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/pg_query_scan.c b/src/pg_query_scan.c index 1d0e052e..50df4296 100644 --- a/src/pg_query_scan.c +++ b/src/pg_query_scan.c @@ -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; @@ -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); @@ -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"))); + } 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) { @@ -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 diff --git a/src/postgres/src_backend_utils_misc_stack_depth.c b/src/postgres/src_backend_utils_misc_stack_depth.c index b1b6690f..eecb6b28 100644 --- a/src/postgres/src_backend_utils_misc_stack_depth.c +++ b/src/postgres/src_backend_utils_misc_stack_depth.c @@ -146,3 +146,14 @@ stack_is_too_deep(void) #if defined(HAVE_GETRLIMIT) #else #endif + +/* 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 = (ssize_t) bytes; + max_stack_depth = (int) (bytes / 1024); +} diff --git a/src/postgres_deparse.c b/src/postgres_deparse.c index e2bdf87d..15fd2466 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" @@ -784,6 +785,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/parse_plpgsql.c b/test/parse_plpgsql.c index 50463691..feadd1fc 100644 --- a/test/parse_plpgsql.c +++ b/test/parse_plpgsql.c @@ -10,6 +10,40 @@ #include #include +static bool +expect_error(const char *input, const char *expected_message) +{ + PgQueryPlpgsqlParseResult result = pg_query_parse_plpgsql(input); + bool success = result.error != NULL && + strcmp(result.error->message, expected_message) == 0; + + if (!success) + { + printf("Expected error: %s\n", expected_message); + printf("Actual error: %s\n", + result.error == NULL ? "(none)" : result.error->message); + } + + pg_query_free_plpgsql_parse_result(result); + return success; +} + +static bool +expect_no_error(const char *input) +{ + PgQueryPlpgsqlParseResult result = pg_query_parse_plpgsql(input); + bool success = result.error == NULL; + + if (!success) + { + printf("Expected no error for: %s\n", input); + printf("Actual error: %s\n", result.error->message); + } + + pg_query_free_plpgsql_parse_result(result); + return success; +} + int main() { bool ret_code = EXIT_SUCCESS; char *sample_buffer; @@ -56,6 +90,31 @@ int main() { pg_query_free_plpgsql_parse_result(result); + if (!expect_error("CREATE FUNCTION f() RETURNS void LANGUAGE plpgsql;", + "no function body specified")) + ret_code = EXIT_FAILURE; + + if (!expect_error("DO LANGUAGE plpgsql;", + "no inline code specified")) + ret_code = EXIT_FAILURE; + + if (!expect_error("CREATE FUNCTION f() RETURNS int LANGUAGE plpgsql RETURN 1;", + "inline SQL function body only valid for language SQL")) + ret_code = EXIT_FAILURE; + + if (!expect_error("CREATE FUNCTION f() RETURNS int AS $$ BEGIN RETURN 1; END; $$ LANGUAGE plpgsql RETURN 1;", + "duplicate function body specified")) + ret_code = EXIT_FAILURE; + + if (!expect_no_error("CREATE FUNCTION f() RETURNS int LANGUAGE sql RETURN 1;")) + ret_code = EXIT_FAILURE; + + if (!expect_no_error("CREATE FUNCTION f() RETURNS int LANGUAGE sql BEGIN ATOMIC SELECT 1; END;")) + ret_code = EXIT_FAILURE; + + if (!expect_no_error("SELECT 1;")) + ret_code = EXIT_FAILURE; + pg_query_exit(); return ret_code; diff --git a/test/scan.c b/test/scan.c index 991ee3ab..25e0ab07 100644 --- a/test/scan.c +++ b/test/scan.c @@ -9,6 +9,58 @@ #include "protobuf/pg_query.pb-c.h" +#if defined(__linux__) && !defined(USE_VALGRIND) +#include +#include + +/* 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; @@ -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(); 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; +}