From c7d1a03ad7d1e90b18049534860bdd87956b9952 Mon Sep 17 00:00:00 2001
From: Thomas Munro <thomas.munro@gmail.com>
Date: Mon, 16 Mar 2026 15:15:18 +1300
Subject: [PATCH v4 02/21] pg_stack_alloc: Defend against arithmetic overflow.

Add checks that pg_stack_alloc_array(T, n) and similar can't overflow
size_t and allow the stack to be overrun.

XXX Separated from initial patch for review.

Reviewed-by:
Discussion:
---
 src/include/utils/pg_stack_alloc.h | 52 ++++++++++++++++++++++++++++--
 src/test/regress/regress.c         | 30 +++++++++++++++++
 2 files changed, 79 insertions(+), 3 deletions(-)

diff --git a/src/include/utils/pg_stack_alloc.h b/src/include/utils/pg_stack_alloc.h
index bb8c58597b9..cd3c3838013 100644
--- a/src/include/utils/pg_stack_alloc.h
+++ b/src/include/utils/pg_stack_alloc.h
@@ -181,12 +181,58 @@ pg_stack_strdup_with_len_impl(char *dst, const char *data, size_t size)
 	return dst;
 }
 
-/* Compute sizeof(T) * n. */
+/* Is it impossible for sizeof(T) * maximum possible n to overflow size_t? */
+static inline bool
+pg_stack_T_mul_n_cannot_overflow_p(size_t sizeof_T, size_t sizeof_n)
+{
+	/* See static assertion that n is not wider than size_t. */
+	if (sizeof_T == 1)
+		return true;
+
+	/* Can't overflow if both factors fit in the lower half of size_t. */
+	if (sizeof_n <= sizeof(size_t) / 2 &&
+		(sizeof_T <= pg_stack_max_for_uint_size(sizeof(size_t) / 2)))
+		return true;
+
+	return false;
+}
+
+/* Would sizeof(T) * n overflow? */
+static inline bool
+pg_stack_T_mul_n_overflows_p(size_t sizeof_T, size_t n)
+{
+	return n > SIZE_MAX / sizeof_T;
+}
+
+/* Compute sizeof(T) * n or raise an error if that would overflow size_t. */
 static inline size_t
 pg_stack_T_mul_n(size_t sizeof_T, size_t sizeof_n, size_t n)
 {
-	/* XXX Could overflow, which might allow the stack to be corrupted. */
-	return sizeof_T * n;
+	size_t		result;
+
+	/*
+	 * These functions are split up so that we can sanity-check them
+	 * individually on 32-bit CI.  For the common case of sizeof(n) == 4 and
+	 * sizeof(size_t) == 8, this should reduce to simple multiplication.
+	 * 32-bit systems can only skip the runtime test for sizeof(T) == 1,
+	 * sizeof(n) == 2 or constexpr n <= UINT16_MAX.
+	 */
+	if (pg_stack_T_mul_n_cannot_overflow_p(sizeof_T, sizeof_n) ||
+		!pg_stack_T_mul_n_overflows_p(sizeof_T, n))
+		result = sizeof_T * n;
+	else
+		elog(ERROR, "pg_stack_alloc: %zu * %zu would overflow size_t",
+			 sizeof_T, n);
+
+	/*
+	 * Explain this theory in terms that GCC's -Werror=stringop-overflow
+	 * understands, so it doesn't complain when the return value is passed to
+	 * memset().  This also serves as an assertion.
+	 */
+	pg_assume(result == n ||
+			  result <= pg_stack_max_for_uint_size(sizeof_n));
+
+	return result;
 }
 
 /*
diff --git a/src/test/regress/regress.c b/src/test/regress/regress.c
index 8f9a7ac4221..f82a0a12e8f 100644
--- a/src/test/regress/regress.c
+++ b/src/test/regress/regress.c
@@ -1393,6 +1393,7 @@ test_pg_stack_alloc(PG_FUNCTION_ARGS)
 {
 	char	   *p;
 	char	   *p2 PG_USED_FOR_ASSERTS_ONLY;
+	bool		raised_error PG_USED_FOR_ASSERTS_ONLY;
 
 	/* Need extra space to run under ASAN. */
 	DECLARE_PG_STACK_SIZE(1024 * 8);
@@ -1415,6 +1416,35 @@ test_pg_stack_alloc(PG_FUNCTION_ARGS)
 	p = pg_stack_alloc(0);
 	Assert(pg_stack_ptr_p(p));
 
+	/* Overflow safety in pg_stack_alloc_array(T, n). */
+	Assert(pg_stack_T_mul_n_cannot_overflow_p(1, sizeof(size_t)));
+	Assert(!pg_stack_T_mul_n_cannot_overflow_p(2, sizeof(size_t)));
+	Assert(pg_stack_T_mul_n_cannot_overflow_p(2, sizeof(size_t) / 2));
+	Assert(!pg_stack_T_mul_n_cannot_overflow_p(2, sizeof(size_t)));
+	Assert(!pg_stack_T_mul_n_overflows_p(1, SIZE_MAX));
+	Assert(pg_stack_T_mul_n_overflows_p(2, SIZE_MAX / 2 + 1));
+	Assert(!pg_stack_T_mul_n_overflows_p(2, SIZE_MAX / 2));
+
+	/* When sizeof_T == 1, always returns n. */
+	Assert(pg_stack_T_mul_n(1, sizeof(uint32), SIZE_MAX) == SIZE_MAX);
+	Assert(pg_stack_T_mul_n(1, sizeof(uint64), SIZE_MAX) == SIZE_MAX);
+
+	/* The largest successful value of size_t n. */
+	Assert(pg_stack_T_mul_n(2, sizeof(size_t), SIZE_MAX / 2) == SIZE_MAX - 1);
+
+	/* Exceed value by one for size_t n. */
+	raised_error = false;
+	PG_TRY();
+	{
+		pg_stack_T_mul_n(2, sizeof(size_t), SIZE_MAX / 2 + 1);
+	}
+	PG_CATCH();
+	{
+		raised_error = true;
+	}
+	PG_END_TRY();
+	Assert(raised_error);
+
 	/* Test a range of alignments. */
 #define TEST_ALIGN MAXIMUM_ALIGNOF
 	for (int i = 1; i <= TEST_ALIGN * 8; i *= 2)
-- 
2.50.1 (Apple Git-155)

