Add stack overflow crash protection, error out instead

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
@@ -339,6 +341,9 @@ test/scan: test/scan.c test/scan_tests.c $(ARLIB)
 test/split: test/split.c test/split_tests.c $(ARLIB)
 	$(CC) $(TEST_CFLAGS) -o $@ test/split.c $(ARLIB) $(TEST_LDFLAGS)
 
+test/stack_depth: test/stack_depth.c $(ARLIB)
+	$(CC) $(TEST_CFLAGS) -o $@ test/stack_depth.c $(ARLIB) $(TEST_LDFLAGS)
+
 prefix = /usr/local
 libdir = $(prefix)/lib
 includedir = $(prefix)/include

scripts/extract_source.rb

@@ -674,6 +674,8 @@ def write_out
 # Other required functions
 runner.deep_resolve('pg_printf')
 runner.deep_resolve('pg_strncasecmp')
+runner.deep_resolve('set_stack_base')
+runner.deep_resolve('assign_max_stack_depth')
 
 # Retain these functions for optional 32-bit support
 # (see BITS_PER_BITMAPWORD checks in bitmapset.c)

src/pg_query.c

@@ -1,14 +1,31 @@
+#if defined(__linux__) && !defined(_GNU_SOURCE)
+/* Allow using pthread_getattr_np on glibc/musl. Must precede any includes. */
+#define _GNU_SOURCE
+#endif
+
 #include "pg_query.h"
 #include "pg_query_internal.h"
 
+#include <miscadmin.h>
 #include <mb/pg_wchar.h>
 #include <utils/memutils.h>
 #include <utils/memdebug.h>
+#include <utils/guc_hooks.h>
 
 #ifdef HAVE_PTHREAD
 #include <pthread.h>
 #endif
 
+#if defined(HAVE_PTHREAD) && (defined(__FreeBSD__) || defined(__NetBSD__))
+/* For pthread_attr_get_np() */
+#include <pthread_np.h>
+#endif
+
+#if defined(_WIN32)
+/* For GetCurrentThreadStackLimits(); port/win32.h already sets _WIN32_WINNT. */
+#include <windows.h>
+#endif
+
 #include <signal.h>
 
 const char* progname = "pg_query";
@@ -20,8 +37,34 @@ static pthread_key_t pg_query_thread_exit_key;
 static void pg_query_thread_exit(void *key);
 #endif
 
+static __thread char *pg_query_lowest_stackaddr = NULL;
+static char *get_lowest_stackaddr(void);
+static void set_max_stack_depth(void);
+
 void pg_query_init(void)
 {
+	/*
+	 * Set the reference point for stack depth checking, similiar to how
+	 * Postgres does in its main() function. We have to re-do this on every
+	 * execution since we may be called from different parts in our host
+	 * program.
+	 *
+	 * The check_stack_depth() function called by recursive tree walkers
+	 * measures the current stack depth against this base and raises a
+	 * "stack depth limit exceeded" error before we overflow and crash.
+	 */
+	set_stack_base();
+
+	/*
+	 * Size the budget to this thread's actual stack; must follow
+	 * set_stack_base() so both measure from the same point.
+	 */
+	set_max_stack_depth();
+
+	/*
+	 * Later parts of initialization don't have to re-run if we execute again
+	 * in the same thread, since they are not affected by the calling location.
+	 */
 	if (pg_query_initialized != 0) return;
 	pg_query_initialized = 1;
 
@@ -111,3 +154,115 @@ void pg_query_free_error(PgQueryError *error)
 
 	free(error);
 }
+
+/*
+ * Get this thread's lowest addressable stack address and cache it for
+ * subsequent calls. Returns NULL if address can't be determined.
+ */
+static char *get_lowest_stackaddr(void)
+{
+	if (pg_query_lowest_stackaddr)
+		return pg_query_lowest_stackaddr;
+
+#if defined(_WIN32)
+	{
+		ULONG_PTR	low, high;
+
+		GetCurrentThreadStackLimits(&low, &high);
+
+		if (low == 0)
+			return NULL;
+
+		pg_query_lowest_stackaddr = (char *) low;
+	}
+#elif defined(HAVE_PTHREAD) && defined(__linux__)
+	{
+		pthread_attr_t attr;
+		void	   *stack_addr = NULL;
+		size_t		stack_size = 0;
+
+		if (pthread_getattr_np(pthread_self(), &attr) != 0)
+			return NULL;
+
+		if (pthread_attr_getstack(&attr, &stack_addr, &stack_size) == 0 && stack_addr != NULL)
+			pg_query_lowest_stackaddr = (char *) stack_addr;
+
+		pthread_attr_destroy(&attr);
+	}
+#elif defined(HAVE_PTHREAD) && defined(__APPLE__)
+	{
+		char	   *stack_high = (char *) pthread_get_stackaddr_np(pthread_self());
+		size_t		stack_size = pthread_get_stacksize_np(pthread_self());
+
+		if (stack_size == 0)
+			return NULL;
+
+		pg_query_lowest_stackaddr = stack_high - stack_size;
+	}
+#elif defined(HAVE_PTHREAD) && (defined(__FreeBSD__) || defined(__NetBSD__))
+	{
+		pthread_attr_t attr;
+		void	   *stack_addr = NULL;
+		size_t		stack_size = 0;
+
+		if (pthread_attr_init(&attr) != 0)
+			return NULL;
+
+		if (pthread_attr_get_np(pthread_self(), &attr) == 0 &&
+			pthread_attr_getstack(&attr, &stack_addr, &stack_size) == 0 &&
+			stack_addr != NULL)
+		{
+			pg_query_lowest_stackaddr = (char *) stack_addr;
+		}
+
+		pthread_attr_destroy(&attr);
+	}
+#endif
+
+	return pg_query_lowest_stackaddr;
+}
+
+/*
+ * Raise max_stack_depth to fit the headroom on the calling thread's stack,
+ * capped at 2MB, similar to how InitializeGUCOptionsFromEnvironment sets
+ * max_stack_depth.
+ *
+ * Re-evaluated for each call, since libpg_query gets embedded into other
+ * programs, so we have to be conscious of our place in the stack, and measure
+ * the remaining stack space we can use below the current address.
+ */
+static void set_max_stack_depth(void)
+{
+	char		stack_here;
+	char*		lowest_stackaddr = get_lowest_stackaddr();
+	ssize_t		usable_stack_depth;
+	ssize_t		new_limit;
+
+	if (lowest_stackaddr == NULL)
+		return;
+
+	usable_stack_depth = (ssize_t) (&stack_here - lowest_stackaddr);
+	if (usable_stack_depth < 0)
+		return;
+
+	/*
+	 * Determine new limit keeping a similar padding to Postgres via
+	 * STACK_DEPTH_SLOP (512 kB) to account for error handling and places
+	 * where we don't check stack depth frequently.
+	 */
+	new_limit = (usable_stack_depth - STACK_DEPTH_SLOP) / 1024;
+
+	/*
+	 * In case the new limit would put us below 100 kB, ignore it, since that
+	 * is our required minimum and the default of the setting.
+	 */
+	if (new_limit <= 100)
+		return;
+
+	/* Clamp at 2MB, like Postgres auto-detection logic. */
+	if (new_limit > 2048)
+		new_limit = 2048;
+
+	max_stack_depth = (int) new_limit;
+	assign_max_stack_depth((int) new_limit, NULL);
+}

src/pg_query_outfuncs_json.c

@@ -9,6 +9,7 @@
 #include "nodes/plannodes.h"
 #include "nodes/value.h"
 #include "utils/datum.h"
+#include "miscadmin.h"
 
 #include "pg_query_json_helper.c"
 
@@ -289,6 +290,8 @@ _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

@@ -7,6 +7,7 @@
 #include "nodes/plannodes.h"
 #include "nodes/value.h"
 #include "utils/datum.h"
+#include "miscadmin.h"
 
 #include "protobuf/pg_query.pb-c.h"
 
@@ -246,6 +247,8 @@ _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_outfuncs_protobuf_cpp.cc

@@ -19,6 +19,7 @@ extern "C"
 #include "nodes/plannodes.h"
 #include "nodes/value.h"
 #include "utils/datum.h"
+#include "miscadmin.h"
 }
 
 #define OUT_TYPE(typename, typename_c) pg_query::typename*
@@ -203,6 +204,8 @@ _outAConst(pg_query::A_Const* out_node, const A_Const *node)
 static void
 _outNode(pg_query::Node* out, const void *obj)
 {
+	check_stack_depth();
+
 	if (obj == NULL)
 		return; // Keep out as NULL
 
@@ -224,46 +227,71 @@ pg_query_nodes_to_protobuf(const void *obj)
 {
 	PgQueryProtobuf protobuf;
 	const ListCell *lc;
-	pg_query::ParseResult parse_result;
+
 	if (obj == NULL) {
 		protobuf.data = strdup("");
 		protobuf.len = 0;
 		return protobuf;
 	}
 
-	parse_result.set_version(PG_VERSION_NUM);
-	foreach(lc, (List*) obj)
+	pg_query::ParseResult *parse_result = new pg_query::ParseResult();
+
+	PG_TRY();
 	{
-		_outRawStmt(parse_result.add_stmts(), (const RawStmt*) lfirst(lc));
-	}
+		parse_result->set_version(PG_VERSION_NUM);
+		foreach(lc, (List*) obj)
+		{
+			_outRawStmt(parse_result->add_stmts(), (const RawStmt*) lfirst(lc));
+		}
 
-	std::string output;
-	parse_result.SerializeToString(&output);
+		std::string output;
+		parse_result->SerializeToString(&output);
 
-	protobuf.data = (char*) calloc(output.size(), sizeof(char));
-	memcpy(protobuf.data, output.data(), output.size());
-	protobuf.len = output.size();
+		protobuf.data = (char*) calloc(output.size(), sizeof(char));
+		memcpy(protobuf.data, output.data(), output.size());
+		protobuf.len = output.size();
+	}
+	PG_CATCH();
+	{
+		delete parse_result;
+		PG_RE_THROW();
+	}
+	PG_END_TRY();
 
+	delete parse_result;
 	return protobuf;
 }
 
 extern "C" char *
 pg_query_nodes_to_json(const void *obj)
 {
 	const ListCell *lc;
-	pg_query::ParseResult parse_result;
+	char	   *result = NULL;
 
 	if (obj == NULL)
 		return pstrdup("{}");
 
-	parse_result.set_version(PG_VERSION_NUM);
-	foreach(lc, (List*) obj)
+	pg_query::ParseResult *parse_result = new pg_query::ParseResult();
+
+	PG_TRY();
 	{
-		_outRawStmt(parse_result.add_stmts(), (const RawStmt*) lfirst(lc));
-	}
+		parse_result->set_version(PG_VERSION_NUM);
+		foreach(lc, (List*) obj)
+		{
+			_outRawStmt(parse_result->add_stmts(), (const RawStmt*) lfirst(lc));
+		}
 
-	std::string output;
-	google::protobuf::util::MessageToJsonString(parse_result, &output);
+		std::string output;
+		google::protobuf::util::MessageToJsonString(*parse_result, &output);
+		result = pstrdup(output.c_str());
+	}
+	PG_CATCH();
+	{
+		delete parse_result;
+		PG_RE_THROW();
+	}
+	PG_END_TRY();
 
-	return pstrdup(output.c_str());
+	delete parse_result;
+	return result;
 }