From eb3ce984d12a45ed073d5b2204d3cf6a22f9b84c Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
Date: Wed, 12 Mar 2025 08:25:06 +0100
Subject: Add additional memory context checks

Currently, a deleted aset MemoryContext in the freelist can't be
distinguished from a valid MemoryContext, and will be seen as valid by
checks like MemoryContextIsValid. To fix that, we set the node's type to
T_Invalid before putting it in the freelist, marking the MemoryContext
as deleted and not usable unless it goes through
AllocSetContextCreateInternal first. The T_Invalid node type will allow
to trigger the MemoryContextIsValid and AllocSetIsValid checks.

We also add additional MemoryContextIsValid checks:
- Before restoring the context stored by a
  transactionState->priorContext
- In MemoryContextCreate to make sure the provided parent is valid
---
 src/backend/access/transam/xact.c | 3 +++
 src/backend/utils/mmgr/aset.c     | 8 ++++++++
 src/backend/utils/mmgr/mcxt.c     | 2 ++
 3 files changed, 13 insertions(+)

diff --git a/src/backend/access/transam/xact.c b/src/backend/access/transam/xact.c
index 1b4f21a88d3..20777943081 100644
--- a/src/backend/access/transam/xact.c
+++ b/src/backend/access/transam/xact.c
@@ -1604,6 +1604,7 @@ AtCommit_Memory(void)
 	 * are about to delete.  If it somehow is, assertions in mcxt.c will
 	 * complain.)
 	 */
+	Assert(MemoryContextIsValid(s->priorContext));
 	MemoryContextSwitchTo(s->priorContext);
 
 	/*
@@ -1983,6 +1984,7 @@ AtCleanup_Memory(void)
 	 * are about to delete.  If it somehow is, assertions in mcxt.c will
 	 * complain.)
 	 */
+	Assert(MemoryContextIsValid(s->priorContext));
 	MemoryContextSwitchTo(s->priorContext);
 
 	/*
@@ -2030,6 +2032,7 @@ AtSubCleanup_Memory(void)
 	 * we are about to delete.  If it somehow is, assertions in mcxt.c will
 	 * complain.)
 	 */
+	Assert(MemoryContextIsValid(s->priorContext));
 	MemoryContextSwitchTo(s->priorContext);
 
 	/* Update CurTransactionContext (might not be same as priorContext) */
diff --git a/src/backend/utils/mmgr/aset.c b/src/backend/utils/mmgr/aset.c
index 666ecd8f78d..5c470561b47 100644
--- a/src/backend/utils/mmgr/aset.c
+++ b/src/backend/utils/mmgr/aset.c
@@ -654,6 +654,14 @@ AllocSetDelete(MemoryContext context)
 			Assert(freelist->num_free == 0);
 		}
 
+		/*
+		 * Set the node type to T_Invalid before putting it in the freelist.
+		 * This is used as a way to mark the context as deleted and shouldn't
+		 * be used as is. Attempting to use it will trip MemoryContextIsValid
+		 * and AllocSetIsValid checks.
+		 */
+		context->type = T_Invalid;
+
 		/* Now add the just-deleted context to the freelist. */
 		set->header.nextchild = (MemoryContext) freelist->first_free;
 		freelist->first_free = set;
diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index 91060de0ab7..db5c2bbfde7 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -1121,6 +1121,8 @@ MemoryContextCreate(MemoryContext node,
 	/* OK to link node into context tree */
 	if (parent)
 	{
+		/* Make sure the provided parent is valid */
+		Assert(MemoryContextIsValid(parent));
 		node->nextchild = parent->firstchild;
 		if (parent->firstchild != NULL)
 			parent->firstchild->prevchild = node;
-- 
2.39.5 (Apple Git-154)

