Make MemoryContextContains work correctly again

c6e0fe1f2 recently changed the way we store headers for allocated chunks
of memory.  Prior to that commit, we stored a pointer to the owning
MemoryContext directly prior to the pointer to the allocated memory.
That's no longer true and c6e0fe1f2 neglected to update
MemoryContextContains() so that it correctly obtains the owning context
with the new method.

A side effect of this change and c6e0fe1f2, in general, is that it's even
less safe than it was previously to pass MemoryContextContains() an
arbitrary pointer which was not allocated by one of our MemoryContexts.
Previously some comments in MemoryContextContains() seemed to indicate
that the worst that could happen by passing an arbitrary pointer would be
a false positive return value.  It seems to me that this was a rather
wishful outlook as we subsequently proceeded to subtract sizeof(void *)
from the given pointer and then dereferenced that memory.  So it seems
quite likely that we could have segfaulted instead of returning a false
positive.  However, it's not impossible that the memory sizeof(void *)
bytes before the pointer could have been owned by the process, but it's
far less likely to work now as obtaining a pointer to the owning
MemoryContext is less direct than before c6e0fe1f2 and will access memory
that's possibly much further away to obtain the owning MemoryContext.
Because of this, I took the liberty of updating the comment to warn
against any future usages of the function and checked the existing core
usages to ensure that we only ever pass in a pointer to memory allocated
by a MemoryContext.

Extension authors updating their code for PG16 who are using
MemoryContextContains should check to ensure that only NULL pointers and
pointers to chunks allocated with a MemoryContext will ever be passed to
MemoryContextContains.

Reported-by: Andres Freund
Discussion: https://postgr.es/m/20220905230949.kb3x2fkpfwtngz43@awork3.anarazel.de

src/backend/utils/mmgr/mcxt.c

@@ -482,6 +482,15 @@ MemoryContextAllowInCriticalSection(MemoryContext context, bool allow)
 MemoryContext
 GetMemoryChunkContext(void *pointer)
 {
+	/*
+	 * 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.
+	 */
+	Assert(pointer != NULL);
+	Assert(pointer == (void *) MAXALIGN(pointer));
+	/* adding further Asserts here? See pre-checks in MemoryContextContains */
+
 	return MCXT_METHOD(pointer, get_chunk_context) (pointer);
 }
 
@@ -809,11 +818,10 @@ MemoryContextCheck(MemoryContext context)
  *		Detect whether an allocated chunk of memory belongs to a given
  *		context or not.
  *
- * Caution: this test is reliable as long as 'pointer' does point to
+ * Caution: 'pointer' must point to a pointer which was allocated by a
- * a chunk of memory allocated from *some* context.  If 'pointer' points
+ * MemoryContext.  It's not safe or valid to use this function on arbitrary
- * at memory obtained in some other way, there is a small chance of a
+ * pointers as obtaining the MemoryContext which 'pointer' belongs to requires
- * false-positive result, since the bits right before it might look like
+ * possibly several pointer dereferences.
- * a valid chunk header by chance.
  */
 bool
 MemoryContextContains(MemoryContext context, void *pointer)
@@ -821,9 +829,8 @@ MemoryContextContains(MemoryContext context, void *pointer)
 	MemoryContext ptr_context;
 
 	/*
-	 * NB: Can't use GetMemoryChunkContext() here - that performs assertions
+	 * NB: We must perform run-time checks here which GetMemoryChunkContext()
-	 * that aren't acceptable here since we might be passed memory not
+	 * does as assertions before calling GetMemoryChunkContext().
-	 * allocated by any memory context.
 	 *
 	 * Try to detect bogus pointers handed to us, poorly though we can.
 	 * Presumably, a pointer that isn't MAXALIGNED isn't pointing at an
@@ -835,7 +842,7 @@ MemoryContextContains(MemoryContext context, void *pointer)
 	/*
 	 * OK, it's probably safe to look at the context.
 	 */
-	ptr_context = *(MemoryContext *) (((char *) pointer) - sizeof(void *));
+	ptr_context = GetMemoryChunkContext(pointer);
 
 	return ptr_context == context;
 }
@@ -953,6 +960,8 @@ MemoryContextAlloc(MemoryContext context, Size size)
 
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
+	Assert(MemoryContextContains(context, ret));
+
 	return ret;
 }
 
@@ -991,6 +1000,8 @@ MemoryContextAllocZero(MemoryContext context, Size size)
 
 	MemSetAligned(ret, 0, size);
 
+	Assert(MemoryContextContains(context, ret));
+
 	return ret;
 }
 
@@ -1029,6 +1040,8 @@ MemoryContextAllocZeroAligned(MemoryContext context, Size size)
 
 	MemSetLoop(ret, 0, size);
 
+	Assert(MemoryContextContains(context, ret));
+
 	return ret;
 }
 
@@ -1070,6 +1083,8 @@ MemoryContextAllocExtended(MemoryContext context, Size size, int flags)
 	if ((flags & MCXT_ALLOC_ZERO) != 0)
 		MemSetAligned(ret, 0, size);
 
+	Assert(MemoryContextContains(context, ret));
+
 	return ret;
 }
 
@@ -1153,6 +1168,8 @@ palloc(Size size)
 
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
+	Assert(MemoryContextContains(context, ret));
+
 	return ret;
 }
 
@@ -1186,6 +1203,8 @@ palloc0(Size size)
 
 	MemSetAligned(ret, 0, size);
 
+	Assert(MemoryContextContains(context, ret));
+
 	return ret;
 }
 
@@ -1225,6 +1244,8 @@ palloc_extended(Size size, int flags)
 	if ((flags & MCXT_ALLOC_ZERO) != 0)
 		MemSetAligned(ret, 0, size);
 
+	Assert(MemoryContextContains(context, ret));
+
 	return ret;
 }
 
@@ -1278,6 +1299,8 @@ repalloc(void *pointer, Size size)
 
 	VALGRIND_MEMPOOL_CHANGE(context, pointer, ret, size);
 
+	Assert(MemoryContextContains(context, ret));
+
 	return ret;
 }
 
@@ -1313,6 +1336,8 @@ MemoryContextAllocHuge(MemoryContext context, Size size)
 
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
+	Assert(MemoryContextContains(context, ret));
+
 	return ret;
 }
 
@@ -1352,6 +1377,8 @@ repalloc_huge(void *pointer, Size size)
 
 	VALGRIND_MEMPOOL_CHANGE(context, pointer, ret, size);
 
+	Assert(MemoryContextContains(context, ret));
+
 	return ret;
 }