18: stack-depth guard, PL/pgSQL alias leak, scan out of memory (combined)

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)
 

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"

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) {

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

src/pg_query_outfuncs_json.c

@@ -1,6 +1,7 @@
 #include "pg_query_outfuncs.h"
 
 #include "postgres.h"
+#include "miscadmin.h"
 
 #include <ctype.h>
 
@@ -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");

src/pg_query_outfuncs_protobuf.c

@@ -1,6 +1,7 @@
 #include "pg_query_outfuncs.h"
 
 #include "postgres.h"
+#include "miscadmin.h"
 #include <ctype.h>
 #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
 

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);
 

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;