v1-0001-Reinstate-MemoryContextContains-to-work-with-arbi.patch

From a9fcf078bc8a87742860abee527199f772f39f5b Mon Sep 17 00:00:00 2001
From: David Rowley <dgrowley@gmail.com>
Date: Thu, 8 Sep 2022 23:32:03 +1200
Subject: [PATCH v1 1/4] Reinstate MemoryContextContains to work with arbitrary
 pointers

---
 src/backend/utils/mmgr/aset.c            | 42 ++++++++++++++++++
 src/backend/utils/mmgr/generation.c      | 45 +++++++++++++++++++
 src/backend/utils/mmgr/mcxt.c            | 56 ++++++++++++++----------
 src/backend/utils/mmgr/slab.c            | 44 +++++++++++++++++++
 src/include/nodes/memnodes.h             |  4 +-
 src/include/utils/memutils_internal.h    |  3 ++
 src/include/utils/memutils_memorychunk.h | 21 +++++++++
 7 files changed, 191 insertions(+), 24 deletions(-)

diff --git a/src/backend/utils/mmgr/aset.c b/src/backend/utils/mmgr/aset.c
index ec423375ae..24ab7399a3 100644
--- a/src/backend/utils/mmgr/aset.c
+++ b/src/backend/utils/mmgr/aset.c
@@ -1333,6 +1333,48 @@ AllocSetGetChunkContext(void *pointer)
 	return &set->header;
 }
 
+/*
+ * AllocSetContains
+ *		Attempt to determine if 'pointer' is memory which was allocated by
+ *		'context'.  Return true if it is, otherwise false.
+ */
+bool
+AllocSetContains(MemoryContext context, void *pointer)
+{
+	MemoryChunk *chunk = PointerGetMemoryChunk(pointer);
+	AllocBlock	block;
+	AllocSet	set = (AllocSet) context;
+
+	/*
+	 * We must use MemoryChunkIsExternalUnSafePointer() instead of
+	 * MemoryChunkIsExternal() here as 'chunk' may not be a MemoryChunk if
+	 * 'pointer' is not pointing to palloced memory.  Below we're careful
+	 * never to dereference 'block' as it could point to memory not owned by
+	 * this process.
+	 */
+	if (MemoryChunkIsExternalUnSafePointer(chunk))
+		block = ExternalChunkGetBlock(chunk);
+	else
+		block = (AllocBlock) MemoryChunkGetBlock(chunk);
+
+	for (AllocBlock blk = set->blocks; blk != NULL; blk = blk->next)
+	{
+		if (block == blk)
+		{
+			/*
+			 * The block matches. Now check if 'pointer' falls within the
+			 * block's memory.  We don't detect if the pointer is pointing to
+			 * an allocated chunk as that would require looping over the
+			 * freelist for this chunk's size.
+			 */
+			if ((char *) pointer >= (char *) blk + ALLOC_BLOCKHDRSZ + ALLOC_CHUNKHDRSZ &&
+				(char *) pointer < blk->endptr)
+				return true;
+		}
+	}
+	return false;
+}
+
 /*
  * AllocSetGetChunkSpace
  *		Given a currently-allocated chunk, determine the total space
diff --git a/src/backend/utils/mmgr/generation.c b/src/backend/utils/mmgr/generation.c
index c743b24fa7..212aba6bf3 100644
--- a/src/backend/utils/mmgr/generation.c
+++ b/src/backend/utils/mmgr/generation.c
@@ -848,6 +848,51 @@ GenerationGetChunkContext(void *pointer)
 	return &block->context->header;
 }
 
+/*
+ * GenerationContains
+ *		Attempt to determine if 'pointer' is memory which was allocated by
+ *		'context'.  Return true if it is, otherwise false.
+ */
+bool
+GenerationContains(MemoryContext context, void *pointer)
+{
+	MemoryChunk *chunk = PointerGetMemoryChunk(pointer);
+	GenerationBlock *block;
+	GenerationContext *set = (GenerationContext *) context;
+	dlist_iter	iter;
+
+	/*
+	 * We must use MemoryChunkIsExternalUnSafePointer() instead of
+	 * MemoryChunkIsExternal() here as 'chunk' may not be a MemoryChunk if
+	 * 'pointer' is not pointing to palloced memory.  Below we're careful
+	 * never to dereference 'block' as it could point to memory not owned by
+	 * this process.
+	 */
+	if (MemoryChunkIsExternalUnSafePointer(chunk))
+		block = ExternalChunkGetBlock(chunk);
+	else
+		block = (GenerationBlock *) MemoryChunkGetBlock(chunk);
+
+	dlist_foreach(iter, &set->blocks)
+	{
+		GenerationBlock *blk = dlist_container(GenerationBlock, node, iter.cur);
+
+		if (block == blk)
+		{
+			/*
+			 * The block matches. Now check if 'pointer' falls within the
+			 * block's memory.  We don't detect if the pointer is pointing to
+			 * an allocated chunk as that would require looping over the
+			 * freelist for this chunk's size.
+			 */
+			if ((char *) pointer >= (char *) blk + Generation_BLOCKHDRSZ + Generation_CHUNKHDRSZ &&
+				(char *) pointer < blk->endptr)
+				return true;
+		}
+	}
+	return false;
+}
+
 /*
  * GenerationGetChunkSpace
  *		Given a currently-allocated chunk, determine the total space
diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index 115a64cfe4..3615a83d53 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -44,6 +44,7 @@ static const MemoryContextMethods mcxt_methods[] = {
 	[MCTX_ASET_ID].reset = AllocSetReset,
 	[MCTX_ASET_ID].delete_context = AllocSetDelete,
 	[MCTX_ASET_ID].get_chunk_context = AllocSetGetChunkContext,
+	[MCTX_ASET_ID].contains = AllocSetContains,
 	[MCTX_ASET_ID].get_chunk_space = AllocSetGetChunkSpace,
 	[MCTX_ASET_ID].is_empty = AllocSetIsEmpty,
 	[MCTX_ASET_ID].stats = AllocSetStats,
@@ -58,6 +59,7 @@ static const MemoryContextMethods mcxt_methods[] = {
 	[MCTX_GENERATION_ID].reset = GenerationReset,
 	[MCTX_GENERATION_ID].delete_context = GenerationDelete,
 	[MCTX_GENERATION_ID].get_chunk_context = GenerationGetChunkContext,
+	[MCTX_GENERATION_ID].contains = GenerationContains,
 	[MCTX_GENERATION_ID].get_chunk_space = GenerationGetChunkSpace,
 	[MCTX_GENERATION_ID].is_empty = GenerationIsEmpty,
 	[MCTX_GENERATION_ID].stats = GenerationStats,
@@ -72,6 +74,7 @@ static const MemoryContextMethods mcxt_methods[] = {
 	[MCTX_SLAB_ID].reset = SlabReset,
 	[MCTX_SLAB_ID].delete_context = SlabDelete,
 	[MCTX_SLAB_ID].get_chunk_context = SlabGetChunkContext,
+	[MCTX_SLAB_ID].contains = SlabContains,
 	[MCTX_SLAB_ID].get_chunk_space = SlabGetChunkSpace,
 	[MCTX_SLAB_ID].is_empty = SlabIsEmpty,
 	[MCTX_SLAB_ID].stats = SlabStats
@@ -818,28 +821,16 @@ MemoryContextCheck(MemoryContext context)
  *		Detect whether an allocated chunk of memory belongs to a given
  *		context or not.
  *
- * Caution: 'pointer' must point to a pointer which was allocated by a
- * MemoryContext.  It's not safe or valid to use this function on arbitrary
- * pointers as obtaining the MemoryContext which 'pointer' belongs to requires
- * possibly several pointer dereferences.
+ * Caution: this test is reliable as long as 'pointer' does point to
+ * a chunk of memory allocated from *some* context.  If 'pointer' points
+ * at memory obtained in some other way, there is a chance of a segmentation
+ * violation due to accessing the chunk header, which look for in the 8 bytes
+ * prior to the pointer.
  */
 bool
 MemoryContextContains(MemoryContext context, void *pointer)
 {
 	/*
-	 * Temporarily make this always return false as we don't yet have a fully
-	 * baked idea on how to make it work correctly with the new MemoryChunk
-	 * code.
-	 */
-	return false;
-
-#ifdef NOT_USED
-	MemoryContext ptr_context;
-
-	/*
-	 * NB: We must perform run-time checks here which GetMemoryChunkContext()
-	 * does as assertions before calling GetMemoryChunkContext().
-	 *
 	 * Try to detect bogus pointers handed to us, poorly though we can.
 	 * Presumably, a pointer that isn't MAXALIGNED isn't pointing at an
 	 * allocated chunk.
@@ -847,13 +838,10 @@ MemoryContextContains(MemoryContext context, void *pointer)
 	if (pointer == NULL || pointer != (void *) MAXALIGN(pointer))
 		return false;
 
-	/*
-	 * OK, it's probably safe to look at the context.
-	 */
-	ptr_context = GetMemoryChunkContext(pointer);
+	if (GetMemoryChunkMethodID(pointer) != context->method_id)
+		return false;
 
-	return ptr_context == context;
-#endif
+	return context->methods->contains(context, pointer);
 }
 
 /*
@@ -900,6 +888,8 @@ MemoryContextCreate(MemoryContext node,
 
 	/* Initialize all standard fields of memory context header */
 	node->type = tag;
+	/* Use uint8 to prevent increasing sizeof(MemoryContextData) */
+	node->method_id = (uint8) method_id;
 	node->isReset = true;
 	node->methods = &mcxt_methods[method_id];
 	node->parent = parent;
@@ -969,6 +959,8 @@ MemoryContextAlloc(MemoryContext context, Size size)
 
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
+	Assert(MemoryContextContains(context, ret));
+
 	return ret;
 }
 
@@ -1007,6 +999,8 @@ MemoryContextAllocZero(MemoryContext context, Size size)
 
 	MemSetAligned(ret, 0, size);
 
+	Assert(MemoryContextContains(context, ret));
+
 	return ret;
 }
 
@@ -1045,6 +1039,8 @@ MemoryContextAllocZeroAligned(MemoryContext context, Size size)
 
 	MemSetLoop(ret, 0, size);
 
+	Assert(MemoryContextContains(context, ret));
+
 	return ret;
 }
 
@@ -1086,6 +1082,8 @@ MemoryContextAllocExtended(MemoryContext context, Size size, int flags)
 	if ((flags & MCXT_ALLOC_ZERO) != 0)
 		MemSetAligned(ret, 0, size);
 
+	Assert(MemoryContextContains(context, ret));
+
 	return ret;
 }
 
@@ -1169,6 +1167,8 @@ palloc(Size size)
 
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
+	Assert(MemoryContextContains(context, ret));
+
 	return ret;
 }
 
@@ -1202,6 +1202,8 @@ palloc0(Size size)
 
 	MemSetAligned(ret, 0, size);
 
+	Assert(MemoryContextContains(context, ret));
+
 	return ret;
 }
 
@@ -1241,6 +1243,8 @@ palloc_extended(Size size, int flags)
 	if ((flags & MCXT_ALLOC_ZERO) != 0)
 		MemSetAligned(ret, 0, size);
 
+	Assert(MemoryContextContains(context, ret));
+
 	return ret;
 }
 
@@ -1294,6 +1298,8 @@ repalloc(void *pointer, Size size)
 
 	VALGRIND_MEMPOOL_CHANGE(context, pointer, ret, size);
 
+	Assert(MemoryContextContains(context, ret));
+
 	return ret;
 }
 
@@ -1329,6 +1335,8 @@ MemoryContextAllocHuge(MemoryContext context, Size size)
 
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
+	Assert(MemoryContextContains(context, ret));
+
 	return ret;
 }
 
@@ -1368,6 +1376,8 @@ repalloc_huge(void *pointer, Size size)
 
 	VALGRIND_MEMPOOL_CHANGE(context, pointer, ret, size);
 
+	Assert(MemoryContextContains(context, ret));
+
 	return ret;
 }
 
diff --git a/src/backend/utils/mmgr/slab.c b/src/backend/utils/mmgr/slab.c
index 9149aaafcb..b5c801eb01 100644
--- a/src/backend/utils/mmgr/slab.c
+++ b/src/backend/utils/mmgr/slab.c
@@ -567,6 +567,50 @@ SlabGetChunkContext(void *pointer)
 	return &slab->header;
 }
 
+/*
+ * SlabContains
+ *		Attempt to determine if 'pointer' is memory which was allocated by
+ *		'context'.  Return true if it is, otherwise false.
+ */
+bool
+SlabContains(MemoryContext context, void *pointer)
+{
+	MemoryChunk *chunk = PointerGetMemoryChunk(pointer);
+	SlabBlock  *block;
+	SlabContext *set = (SlabContext *) context;
+
+	/*
+	 * Careful not to dereference anything in 'block' as if 'pointer' is not a
+	 * pointer to an allocated chunk then 'block' could be pointing to about
+	 * anything.
+	 */
+	block = (SlabBlock *) MemoryChunkGetBlock(chunk);
+
+	for (int i = 0; i <= set->chunksPerBlock; i++)
+	{
+		dlist_iter	iter;
+
+		dlist_foreach(iter, &set->freelist[i])
+		{
+			SlabBlock  *blk = dlist_container(SlabBlock, node, iter.cur);
+
+			if (block == blk)
+			{
+				/*
+				 * The block matches. Now check if 'pointer' falls within the
+				 * block's memory.  We don't detect if the pointer is pointing
+				 * to an allocated chunk as that would require looping over
+				 * the freelist for this chunk's size.
+				 */
+				if ((char *) pointer >= (char *) blk + Slab_BLOCKHDRSZ + Slab_CHUNKHDRSZ &&
+					(char *) pointer < (char *) blk + set->blockSize)
+					return true;
+			}
+		}
+	}
+	return false;
+}
+
 /*
  * SlabGetChunkSpace
  *		Given a currently-allocated chunk, determine the total space
diff --git a/src/include/nodes/memnodes.h b/src/include/nodes/memnodes.h
index 63d07358cd..6a8d4f4e4e 100644
--- a/src/include/nodes/memnodes.h
+++ b/src/include/nodes/memnodes.h
@@ -64,6 +64,7 @@ typedef struct MemoryContextMethods
 	void		(*reset) (MemoryContext context);
 	void		(*delete_context) (MemoryContext context);
 	MemoryContext (*get_chunk_context) (void *pointer);
+	bool		(*contains) (MemoryContext context, void *pointer);
 	Size		(*get_chunk_space) (void *pointer);
 	bool		(*is_empty) (MemoryContext context);
 	void		(*stats) (MemoryContext context,
@@ -81,7 +82,8 @@ typedef struct MemoryContextData
 	pg_node_attr(abstract)		/* there are no nodes of this type */
 
 	NodeTag		type;			/* identifies exact kind of context */
-	/* these two fields are placed here to minimize alignment wastage: */
+	/* these three fields are placed here to minimize alignment wastage: */
+	uint8		method_id;		/* MemoryContextMethodID for this context */
 	bool		isReset;		/* T = no space alloced since last reset */
 	bool		allowInCritSection; /* allow palloc in critical section */
 	Size		mem_allocated;	/* track memory allocated for this context */
diff --git a/src/include/utils/memutils_internal.h b/src/include/utils/memutils_internal.h
index 377cee7a84..ebc1d60393 100644
--- a/src/include/utils/memutils_internal.h
+++ b/src/include/utils/memutils_internal.h
@@ -25,6 +25,7 @@ extern void *AllocSetRealloc(void *pointer, Size size);
 extern void AllocSetReset(MemoryContext context);
 extern void AllocSetDelete(MemoryContext context);
 extern MemoryContext AllocSetGetChunkContext(void *pointer);
+extern bool AllocSetContains(MemoryContext context, void *pointer);
 extern Size AllocSetGetChunkSpace(void *pointer);
 extern bool AllocSetIsEmpty(MemoryContext context);
 extern void AllocSetStats(MemoryContext context,
@@ -42,6 +43,7 @@ extern void *GenerationRealloc(void *pointer, Size size);
 extern void GenerationReset(MemoryContext context);
 extern void GenerationDelete(MemoryContext context);
 extern MemoryContext GenerationGetChunkContext(void *pointer);
+extern bool GenerationContains(MemoryContext context, void *pointer);
 extern Size GenerationGetChunkSpace(void *pointer);
 extern bool GenerationIsEmpty(MemoryContext context);
 extern void GenerationStats(MemoryContext context,
@@ -60,6 +62,7 @@ extern void *SlabRealloc(void *pointer, Size size);
 extern void SlabReset(MemoryContext context);
 extern void SlabDelete(MemoryContext context);
 extern MemoryContext SlabGetChunkContext(void *pointer);
+extern bool SlabContains(MemoryContext context, void *pointer);
 extern Size SlabGetChunkSpace(void *pointer);
 extern bool SlabIsEmpty(MemoryContext context);
 extern void SlabStats(MemoryContext context,
diff --git a/src/include/utils/memutils_memorychunk.h b/src/include/utils/memutils_memorychunk.h
index 2eefc138e3..721c7d004c 100644
--- a/src/include/utils/memutils_memorychunk.h
+++ b/src/include/utils/memutils_memorychunk.h
@@ -54,6 +54,12 @@
  *		Determine if the given MemoryChunk is externally managed, i.e.
  *		MemoryChunkSetHdrMaskExternal() was called on the chunk.
  *
+ * MemoryChunkIsExternalUnSafePointer:
+ *		As MemoryChunkIsExternal but safe to use on pointers that we're
+ *		uncertain are pointers to MemoryChunks.  Callers must be careful
+ *		not to put too much trust into the return value.  The primary usecase
+ *		for this is to implement MemoryContextContains.
+ *
  * MemoryChunkGetValue:
  *		For non-external chunks, return the stored 30-bit value as it was set
  *		in the call to MemoryChunkSetHdrMask().
@@ -198,6 +204,21 @@ MemoryChunkIsExternal(MemoryChunk *chunk)
 	return HdrMaskIsExternal(chunk->hdrmask);
 }
 
+/*
+ * MemoryChunkIsExternalUnSafePointer
+ *		As MemoryChunkIsExternal only without any Assert checking.  This
+ *		version should only be used when we're uncertain of 'chunk' is
+ *		actually a pointer to a MemoryChunk.
+ *
+ * Note: Callers must be careful not to put too much trust into the return
+ * value as 'chunk' may not be a MemoryChunk.
+ */
+static inline bool
+MemoryChunkIsExternalUnSafePointer(MemoryChunk *chunk)
+{
+	return HdrMaskIsExternal(chunk->hdrmask);
+}
+
 /*
  * MemoryChunkGetValue
  *		For non-external chunks, returns the value field as it was set in
-- 
2.34.1