make_memory_context_contains_more_robust_to_any_pointer.patch

diff --git a/src/backend/utils/mmgr/aset.c b/src/backend/utils/mmgr/aset.c
index ec423375ae..841d6cda5d 100644
--- a/src/backend/utils/mmgr/aset.c
+++ b/src/backend/utils/mmgr/aset.c
@@ -1333,6 +1333,46 @@ 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;
+
+	/*
+	 * 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.
+	 */
+	if (MemoryChunkIsExternal(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..653fa08064 100644
--- a/src/backend/utils/mmgr/generation.c
+++ b/src/backend/utils/mmgr/generation.c
@@ -848,6 +848,49 @@ 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;
+
+	/*
+	 * 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.
+	 */
+	if (MemoryChunkIsExternal(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 3d80abbfae..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,20 +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)
 {
-	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.
@@ -839,12 +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;
+	return context->methods->contains(context, pointer);
 }
 
 /*
@@ -891,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;
diff --git a/src/backend/utils/mmgr/slab.c b/src/backend/utils/mmgr/slab.c
index 9149aaafcb..208c034349 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..0e2097250c 100644
--- a/src/include/utils/memutils_memorychunk.h
+++ b/src/include/utils/memutils_memorychunk.h
@@ -192,8 +192,10 @@ MemoryChunkIsExternal(MemoryChunk *chunk)
 	 * External chunks should always store MEMORYCHUNK_MAGIC in the upper
 	 * portion of the hdrmask, check that nothing has stomped on that.
 	 */
+	/* XXX think of what to do here. These will fail if chunk is just a
+	 * arbitrary pointer passed in via MemoryContextContains()
 	Assert(!HdrMaskIsExternal(chunk->hdrmask) ||
-		   HdrMaskCheckMagic(chunk->hdrmask));
+		   HdrMaskCheckMagic(chunk->hdrmask));*/
 
 	return HdrMaskIsExternal(chunk->hdrmask);
 }