From 8b6b7b478478f213376d9c925571eb7418241579 Mon Sep 17 00:00:00 2001
From: Bertrand Drouvot <bertranddrouvot.pg@gmail.com>
Date: Wed, 1 Jul 2026 05:43:07 +0000
Subject: [PATCH v2] Prevent crash when calling pgstat functions with
 unregistered stats kind

If a custom statistics extension is loaded via CREATE EXTENSION without
being listed in shared_preload_libraries, its _PG_init() skips the call
to pgstat_register_kind(). The SQL functions are still created, and
calling them invokes pgstat functions with a kind that was never registered.

pgstat_get_kind_info() returns NULL in this case. The existing code only
checked this via Assert() in some paths, so non-assert builds would
dereference NULL and segfault.

Add runtime checks in all public-facing pgstat functions that accept a
PgStat_Kind and dereference the returned kind info:

  - pgstat_prep_pending_entry()
  - pgstat_fetch_entry()
  - pgstat_reset()
  - pgstat_reset_of_kind()
  - pgstat_have_entry()
  - pgstat_snapshot_fixed()
  - pgstat_init_entry()
  - pgstat_reset_entry()

Each now raises ERROR with ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE when
the kind is not known or registered.

This affects any extension using the custom cumulative statistics API
introduced in PG18.

Author: Bertrand Drouvot <bertranddrouvot.pg@gmail.com>
Reviewed-by: Ewan Young <kdbase.hack@gmail.com>
Discussion: https://postgr.es/m/akS/ldidWeqG1FWk%40bdtpg
---
 src/backend/utils/activity/pgstat.c       | 46 ++++++++++++++++++++---
 src/backend/utils/activity/pgstat_shmem.c | 13 ++++++-
 2 files changed, 52 insertions(+), 7 deletions(-)
 100.0% src/backend/utils/activity/

diff --git a/src/backend/utils/activity/pgstat.c b/src/backend/utils/activity/pgstat.c
index c4fa14f138f..540db1ef115 100644
--- a/src/backend/utils/activity/pgstat.c
+++ b/src/backend/utils/activity/pgstat.c
@@ -885,8 +885,13 @@ pgstat_reset(PgStat_Kind kind, Oid dboid, uint64 objid)
 	const PgStat_KindInfo *kind_info = pgstat_get_kind_info(kind);
 	TimestampTz ts = GetCurrentTimestamp();
 
+	if (unlikely(kind_info == NULL))
+		ereport(ERROR,
+				(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
+				 errmsg("statistics kind %u is not known or registered", kind)));
+
 	/* not needed atm, and doesn't make sense with the current signature */
-	Assert(!pgstat_get_kind_info(kind)->fixed_amount);
+	Assert(!kind_info->fixed_amount);
 
 	/* reset the "single counter" */
 	pgstat_reset_entry(kind, dboid, objid, ts);
@@ -907,6 +912,11 @@ pgstat_reset_of_kind(PgStat_Kind kind)
 	const PgStat_KindInfo *kind_info = pgstat_get_kind_info(kind);
 	TimestampTz ts = GetCurrentTimestamp();
 
+	if (unlikely(kind_info == NULL))
+		ereport(ERROR,
+				(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
+				 errmsg("statistics kind %u is not known or registered", kind)));
+
 	if (kind_info->fixed_amount)
 		kind_info->reset_all_cb(ts);
 	else
@@ -967,6 +977,11 @@ pgstat_fetch_entry(PgStat_Kind kind, Oid dboid, uint64 objid, bool *may_free)
 	void	   *stats_data;
 	const PgStat_KindInfo *kind_info = pgstat_get_kind_info(kind);
 
+	if (unlikely(kind_info == NULL))
+		ereport(ERROR,
+				(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
+				 errmsg("statistics kind %u is not known or registered", kind)));
+
 	/* should be called from backends */
 	Assert(IsUnderPostmaster || !IsPostmasterEnvironment);
 	Assert(!kind_info->fixed_amount);
@@ -1088,8 +1103,15 @@ pgstat_get_stat_snapshot_timestamp(bool *have_snapshot)
 bool
 pgstat_have_entry(PgStat_Kind kind, Oid dboid, uint64 objid)
 {
+	const PgStat_KindInfo *kind_info = pgstat_get_kind_info(kind);
+
+	if (unlikely(kind_info == NULL))
+		ereport(ERROR,
+				(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
+				 errmsg("statistics kind %u is not known or registered", kind)));
+
 	/* fixed-numbered stats always exist */
-	if (pgstat_get_kind_info(kind)->fixed_amount)
+	if (kind_info->fixed_amount)
 		return true;
 
 	return pgstat_get_entry_ref(kind, dboid, objid, false, NULL) != NULL;
@@ -1104,8 +1126,14 @@ pgstat_have_entry(PgStat_Kind kind, Oid dboid, uint64 objid)
 void
 pgstat_snapshot_fixed(PgStat_Kind kind)
 {
-	Assert(pgstat_is_kind_valid(kind));
-	Assert(pgstat_get_kind_info(kind)->fixed_amount);
+	const PgStat_KindInfo *kind_info = pgstat_get_kind_info(kind);
+
+	if (unlikely(kind_info == NULL))
+		ereport(ERROR,
+				(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
+				 errmsg("statistics kind %u is not known or registered", kind)));
+
+	Assert(kind_info->fixed_amount);
 
 	if (force_stats_snapshot_clear)
 		pgstat_clear_snapshot();
@@ -1310,9 +1338,15 @@ PgStat_EntryRef *
 pgstat_prep_pending_entry(PgStat_Kind kind, Oid dboid, uint64 objid, bool *created_entry)
 {
 	PgStat_EntryRef *entry_ref;
+	const PgStat_KindInfo *kind_info = pgstat_get_kind_info(kind);
+
+	if (unlikely(kind_info == NULL))
+		ereport(ERROR,
+				(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
+				 errmsg("statistics kind %u is not known or registered", kind)));
 
 	/* need to be able to flush out */
-	Assert(pgstat_get_kind_info(kind)->flush_pending_cb != NULL);
+	Assert(kind_info->flush_pending_cb != NULL);
 
 	if (unlikely(!pgStatPendingContext))
 	{
@@ -1327,7 +1361,7 @@ pgstat_prep_pending_entry(PgStat_Kind kind, Oid dboid, uint64 objid, bool *creat
 
 	if (entry_ref->pending == NULL)
 	{
-		size_t		entrysize = pgstat_get_kind_info(kind)->pending_size;
+		size_t		entrysize = kind_info->pending_size;
 
 		Assert(entrysize != (size_t) -1);
 
diff --git a/src/backend/utils/activity/pgstat_shmem.c b/src/backend/utils/activity/pgstat_shmem.c
index 5ea3f1973f9..4e6a556af93 100644
--- a/src/backend/utils/activity/pgstat_shmem.c
+++ b/src/backend/utils/activity/pgstat_shmem.c
@@ -318,6 +318,11 @@ pgstat_init_entry(PgStat_Kind kind,
 	PgStatShared_Common *shheader;
 	const PgStat_KindInfo *kind_info = pgstat_get_kind_info(kind);
 
+	if (unlikely(kind_info == NULL))
+		ereport(ERROR,
+				(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
+				 errmsg("statistics kind %u is not known or registered", kind)));
+
 	/*
 	 * Initialize refcount to 1, marking it as valid / not dropped. The entry
 	 * can't be freed before the initialization because it can't be found as
@@ -1127,8 +1132,14 @@ void
 pgstat_reset_entry(PgStat_Kind kind, Oid dboid, uint64 objid, TimestampTz ts)
 {
 	PgStat_EntryRef *entry_ref;
+	const PgStat_KindInfo *kind_info = pgstat_get_kind_info(kind);
+
+	if (unlikely(kind_info == NULL))
+		ereport(ERROR,
+				(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
+				 errmsg("statistics kind %u is not known or registered", kind)));
 
-	Assert(!pgstat_get_kind_info(kind)->fixed_amount);
+	Assert(!kind_info->fixed_amount);
 
 	entry_ref = pgstat_get_entry_ref(kind, dboid, objid, false, NULL);
 	if (!entry_ref || entry_ref->shared_entry->dropped)
-- 
2.34.1

