From eae6ab255b390404633709e2a150e3af0aa13b00 Mon Sep 17 00:00:00 2001
From: Jakub Wartak <jakub.wartak@enterprisedb.com>
Date: Thu, 28 Sep 2023 10:56:57 +0200
Subject: [PATCH v3] Introduce memory limit for BackendActivityBuffer size, and
 cast calculations to size_t.

As per email discussion put a limit on BackendActivityBuffer max size to 4GB.

In addition to this, cast related MemoryContextAllocHuge() calculations
in pg_stat_get_activity() SQL function to size_t to avoid errors with
with still buffer size but lower than 4GB. This prevents integer overflow
for MemoryContextAllocHuge() when it is being called on system with
high values of max_connections (3000) and high
pgstat_track_activity_query_size (e.g. 1MB):
     postgres=# select * from pg_stat_get_activity(NULL);
     ERROR:  invalid memory alloc request size 18446744072590721024
---
 src/backend/utils/activity/backend_status.c | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/src/backend/utils/activity/backend_status.c b/src/backend/utils/activity/backend_status.c
index 722c5acf38..64d07b15da 100644
--- a/src/backend/utils/activity/backend_status.c
+++ b/src/backend/utils/activity/backend_status.c
@@ -37,6 +37,8 @@
  */
 #define NumBackendStatSlots (MaxBackends + NUM_AUXPROCTYPES)
 
+/* Safety net to prevent requesting huge memory by each query to pg_stat_activity */
+#define PGSTAT_MAX_ACTIVITY_BUF_SIZE 4 * 1024 * 1024 * 1024L
 
 /* ----------
  * GUC parameters
@@ -84,6 +86,7 @@ Size
 BackendStatusShmemSize(void)
 {
 	Size		size;
+	Size		pgstat_track_size;
 
 	/* BackendStatusArray: */
 	size = mul_size(sizeof(PgBackendStatus), NumBackendStatSlots);
@@ -94,8 +97,12 @@ BackendStatusShmemSize(void)
 	size = add_size(size,
 					mul_size(NAMEDATALEN, NumBackendStatSlots));
 	/* BackendActivityBuffer: */
-	size = add_size(size,
-					mul_size(pgstat_track_activity_query_size, NumBackendStatSlots));
+	pgstat_track_size = mul_size(pgstat_track_activity_query_size,
+					NumBackendStatSlots);
+	if(pgstat_track_size >= PGSTAT_MAX_ACTIVITY_BUF_SIZE)
+		elog(FATAL, "too big Backend Activity Buffer allocation of %zu bytes", pgstat_track_size);
+	size = add_size(size, pgstat_track_size);
+
 #ifdef USE_SSL
 	/* BackendSslStatusBuffer: */
 	size = add_size(size,
@@ -765,7 +772,7 @@ pgstat_read_current_status(void)
 						   NAMEDATALEN * NumBackendStatSlots);
 	localactivity = (char *)
 		MemoryContextAllocHuge(backendStatusSnapContext,
-							   pgstat_track_activity_query_size * NumBackendStatSlots);
+							   (size_t)pgstat_track_activity_query_size * (size_t)NumBackendStatSlots);
 #ifdef USE_SSL
 	localsslstatus = (PgBackendSSLStatus *)
 		MemoryContextAlloc(backendStatusSnapContext,
-- 
2.30.2

