From 039a176da01f4a7bf6a28f434683b7da6c9f0965 Mon Sep 17 00:00:00 2001
From: Jakub Wartak <jakub.wartak@enterprisedb.com>
Date: Mon, 2 Oct 2023 10:19:46 +0200
Subject: [PATCH v4] 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 | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/src/backend/utils/activity/backend_status.c b/src/backend/utils/activity/backend_status.c
index 722c5acf38..f638058cc3 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,14 @@ 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)
+		ereport(ERROR,
+			(errcode(ERRCODE_CONFIG_FILE_ERROR),
+			 errmsg("size requested for backend status is out of range")));
+	size = add_size(size, pgstat_track_size);
+
 #ifdef USE_SSL
 	/* BackendSslStatusBuffer: */
 	size = add_size(size,
@@ -765,7 +774,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

