Ensure that we don't try to allocate connection in/out buffers over

MaxAllocSize.

We take this opportunity to rearrange the code to avoid duplicity in handling
in and out buffers. Also, add some other checks to ensure that we don't overrun
the limits.

Report, investigation and draft patch by Krzysztof Nienartowicz.

1 files changed, 85 insertions, 73 deletions

diff --git a/src/backend/pgxc/pool/pgxcnode.c b/src/backend/pgxc/pool/pgxcnode.c
index 4375e8813a..eafd9cbbe0 100644
--- a/src/backend/pgxc/pool/pgxcnode.c
+++ b/src/backend/pgxc/pool/pgxcnode.c
@@ -950,105 +950,117 @@ release_handles(void)
 }
 
 /*
- * Ensure specified amount of data can fit to the incoming buffer and
- * increase it if necessary
+ * Ensure that the supplied buffer has enough capacity and if not, it's
+ * extended to an appropriate size.
+ *
+ * currbuf is the currently used buffer of currsize. bytes_needed is the
+ * minimum size required. We shall return the new buffer, if allocated
+ * successfully and set newsize_p to contain the size of the repalloced buffer.
+ * If allocation fails, NULL is returned.
+ *
+ * The function checks for requests beyond MaxAllocSize and throw an error.
  */
-int
-ensure_in_buffer_capacity(size_t bytes_needed, PGXCNodeHandle *handle)
+static char *
+ensure_buffer_capacity(char *currbuf, size_t currsize, size_t bytes_needed, size_t *newsize_p)
 {
-	int			newsize = handle->inSize;
 	char	   *newbuf;
+	Size		newsize = (Size) currsize;
 
-	if (bytes_needed <= (size_t) newsize)
-		return 0;
+	if (((Size) bytes_needed) >= MaxAllocSize)
+		ereport(ERROR,
+				(ENOSPC,
+				 errmsg("out of memory"),
+				 errdetail("Cannot enlarge buffer containing %ld bytes by %ld more bytes.",
+						   currsize, bytes_needed)));
 
-	do
+	if (bytes_needed <= newsize)
 	{
-		newsize *= 2;
-	} while (newsize > 0 && bytes_needed > (size_t) newsize);
+		*newsize_p = currsize;
+		return currbuf;
+	}
+
+	/*
+	 * The current size of the buffer should never be zero (init_pgxc_handle
+	 * guarantees that.
+	 */
+	Assert(newsize > 0);
+
+	/*
+	 * Double the buffer size until we have enough space to hold bytes_needed
+	 */
+	while (bytes_needed > newsize)
+		newsize = 2 * newsize;
 
-	if (newsize > 0 && bytes_needed <= (size_t) newsize)
+	/*
+	 * Clamp to MaxAllocSize in case we went past it.  Note we are assuming
+	 * here that MaxAllocSize <= INT_MAX/2, else the above loop could
+	 * overflow.  We will still have newsize >= bytes_needed.
+	 */
+	if (newsize > (int) MaxAllocSize)
+		newsize = (int) MaxAllocSize;
+
+	newbuf = repalloc(currbuf, newsize);
+	if (newbuf)
 	{
-		newbuf = repalloc(handle->inBuffer, newsize);
-		if (newbuf)
-		{
-			/* repalloc succeeded */
-			handle->inBuffer = newbuf;
-			handle->inSize = newsize;
-			return 0;
-		}
+		/* repalloc succeeded, set new size and return the buffer */
+		*newsize_p = newsize;
+		return newbuf;
 	}
 
-	newsize = handle->inSize;
-	do
-	{
-		newsize += 8192;
-	} while (newsize > 0 && bytes_needed > (size_t) newsize);
+	/*
+	 * If we fail to double the buffer, try to repalloc a buffer of the given
+	 * size, rounded to the next multiple of 8192 and see if that works.
+	 */
+	newsize = bytes_needed;
+	newsize = ((bytes_needed / 8192) + 1) * 8192;
 
-	if (newsize > 0 && bytes_needed <= (size_t) newsize)
+	newbuf = repalloc(currbuf, newsize);
+	if (newbuf)
 	{
-		newbuf = repalloc(handle->inBuffer, newsize);
-		if (newbuf)
-		{
-			/* repalloc succeeded */
-			handle->inBuffer = newbuf;
-			handle->inSize = newsize;
-			return 0;
-		}
+		/* repalloc succeeded, set new size and return the buffer */
+		*newsize_p = newsize;
+		return newbuf;
 	}
 
-	return EOF;
+	/* repalloc failed */
+	return NULL;
 }
 
-
 /*
- * Ensure specified amount of data can fit to the outgoing buffer and
+ * Ensure specified amount of data can fit to the incoming buffer and
  * increase it if necessary
  */
 int
-ensure_out_buffer_capacity(size_t bytes_needed, PGXCNodeHandle *handle)
+ensure_in_buffer_capacity(size_t bytes_needed, PGXCNodeHandle *handle)
 {
-	int			newsize = handle->outSize;
-	char	   *newbuf;
-
-	if (bytes_needed <= (size_t) newsize)
-		return 0;
-
-	do
+	size_t newsize;
+	char *newbuf = ensure_buffer_capacity(handle->inBuffer, handle->inSize,
+			bytes_needed, &newsize);
+	if (newbuf)
 	{
-		newsize *= 2;
-	} while (newsize > 0 && bytes_needed > (size_t) newsize);
-
-	if (newsize > 0 && bytes_needed <= (size_t) newsize)
-	{
-		newbuf = repalloc(handle->outBuffer, newsize);
-		if (newbuf)
-		{
-			/* repalloc succeeded */
-			handle->outBuffer = newbuf;
-			handle->outSize = newsize;
-			return 0;
-		}
+		handle->inBuffer = newbuf;
+		handle->inSize = newsize;
+		return 0;
 	}
+	return EOF;
+}
 
-	newsize = handle->outSize;
-	do
-	{
-		newsize += 8192;
-	} while (newsize > 0 && bytes_needed > (size_t) newsize);
-
-	if (newsize > 0 && bytes_needed <= (size_t) newsize)
+/*
+ * Ensure specified amount of data can fit to the outgoing buffer and
+ * increase it if necessary
+ */
+int
+ensure_out_buffer_capacity(size_t bytes_needed, PGXCNodeHandle *handle)
+{
+	size_t newsize;
+	char *newbuf = ensure_buffer_capacity(handle->outBuffer, handle->outSize,
+			bytes_needed, &newsize);
+	if (newbuf)
 	{
-		newbuf = repalloc(handle->outBuffer, newsize);
-		if (newbuf)
-		{
-			/* repalloc succeeded */
-			handle->outBuffer = newbuf;
-			handle->outSize = newsize;
-			return 0;
-		}
+		handle->outBuffer = newbuf;
+		handle->outSize = newsize;
+		return 0;
 	}
-
 	return EOF;
 }