From 7943d11f983a8654978a1db868b4a5e3ca22ec28 Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@otacoo.com>
Date: Thu, 9 Jul 2015 12:49:31 +0000
Subject: [PATCH 1/2] Prevent COPY start from hangling in libpq on OOM

This commit improves the OOM error handling in getCopyStart of the
protocol 3 of libpq that could make process wait for caller if the
message is not consumed completely on client. When an OOM is found
out, register an error message indicating what has actually happened.

In order to allow the input parser to detect if a failure is happening
a new async flag called PGASYNC_FATAL is used to have a correct error
flow control.
---
 src/interfaces/libpq/fe-exec.c      |  3 ++
 src/interfaces/libpq/fe-protocol3.c | 82 ++++++++++++++++++++++++++++++-------
 src/interfaces/libpq/libpq-int.h    |  3 +-
 3 files changed, 72 insertions(+), 16 deletions(-)

diff --git a/src/interfaces/libpq/fe-exec.c b/src/interfaces/libpq/fe-exec.c
index 828f18e..6c6eeaf 100644
--- a/src/interfaces/libpq/fe-exec.c
+++ b/src/interfaces/libpq/fe-exec.c
@@ -1744,6 +1744,9 @@ PQgetResult(PGconn *conn)
 		case PGASYNC_COPY_BOTH:
 			res = getCopyResult(conn, PGRES_COPY_BOTH);
 			break;
+		case PGASYNC_FATAL:
+			res = NULL;			/* some trouble happened */
+			break;
 		default:
 			printfPQExpBuffer(&conn->errorMessage,
 							  libpq_gettext("unexpected asyncStatus: %d\n"),
diff --git a/src/interfaces/libpq/fe-protocol3.c b/src/interfaces/libpq/fe-protocol3.c
index dbc0d89..d0b22b7 100644
--- a/src/interfaces/libpq/fe-protocol3.c
+++ b/src/interfaces/libpq/fe-protocol3.c
@@ -49,7 +49,9 @@ static int	getParamDescriptions(PGconn *conn);
 static int	getAnotherTuple(PGconn *conn, int msgLength);
 static int	getParameterStatus(PGconn *conn);
 static int	getNotify(PGconn *conn);
-static int	getCopyStart(PGconn *conn, ExecStatusType copytype);
+static int	getCopyStart(PGconn *conn,
+						 ExecStatusType copytype,
+						 int msgLength);
 static int	getReadyForQuery(PGconn *conn);
 static void reportErrorPosition(PQExpBuffer msg, const char *query,
 					int loc, int encoding);
@@ -362,19 +364,28 @@ pqParseInput3(PGconn *conn)
 					}
 					break;
 				case 'G':		/* Start Copy In */
-					if (getCopyStart(conn, PGRES_COPY_IN))
+					if (getCopyStart(conn, PGRES_COPY_IN, msgLength))
+					{
+						conn->asyncStatus = PGASYNC_FATAL;
 						return;
+					}
 					conn->asyncStatus = PGASYNC_COPY_IN;
 					break;
 				case 'H':		/* Start Copy Out */
-					if (getCopyStart(conn, PGRES_COPY_OUT))
+					if (getCopyStart(conn, PGRES_COPY_OUT, msgLength))
+					{
+						conn->asyncStatus = PGASYNC_FATAL;
 						return;
+					}
 					conn->asyncStatus = PGASYNC_COPY_OUT;
 					conn->copy_already_done = 0;
 					break;
 				case 'W':		/* Start Copy Both */
-					if (getCopyStart(conn, PGRES_COPY_BOTH))
+					if (getCopyStart(conn, PGRES_COPY_BOTH, msgLength))
+					{
+						conn->asyncStatus = PGASYNC_FATAL;
 						return;
+					}
 					conn->asyncStatus = PGASYNC_COPY_BOTH;
 					conn->copy_already_done = 0;
 					break;
@@ -1338,31 +1349,45 @@ getNotify(PGconn *conn)
  * parseInput already read the message type and length.
  */
 static int
-getCopyStart(PGconn *conn, ExecStatusType copytype)
+getCopyStart(PGconn *conn, ExecStatusType copytype, int msgLength)
 {
 	PGresult   *result;
 	int			nfields;
 	int			i;
+	const char *errmsg;
 
 	result = PQmakeEmptyPGresult(conn, copytype);
 	if (!result)
-		goto failure;
+	{
+		errmsg = NULL;
+		goto advance_and_error;
+	}
 
 	if (pqGetc(&conn->copy_is_binary, conn))
-		goto failure;
+	{
+		errmsg = libpq_gettext("extraneous data in COPY start message");
+		goto advance_and_error;
+	}
 	result->binary = conn->copy_is_binary;
+
 	/* the next two bytes are the number of fields	*/
-	if (pqGetInt(&(result->numAttributes), 2, conn))
-		goto failure;
+	if (pqGetInt(&result->numAttributes, 2, conn))
+	{
+		errmsg = libpq_gettext("extraneous data in COPY start message");
+		goto advance_and_error;
+	}
 	nfields = result->numAttributes;
 
 	/* allocate space for the attribute descriptors */
-	if (nfields > 0)
+	if (result && nfields > 0)
 	{
 		result->attDescs = (PGresAttDesc *)
 			pqResultAlloc(result, nfields * sizeof(PGresAttDesc), TRUE);
 		if (!result->attDescs)
-			goto failure;
+		{
+			errmsg = NULL;
+			goto advance_and_error;
+		}
 		MemSet(result->attDescs, 0, nfields * sizeof(PGresAttDesc));
 	}
 
@@ -1371,22 +1396,49 @@ getCopyStart(PGconn *conn, ExecStatusType copytype)
 		int			format;
 
 		if (pqGetInt(&format, 2, conn))
-			goto failure;
+		{
+			errmsg = libpq_gettext("extraneous data in COPY start message");
+			goto advance_and_error;
+		}
 
 		/*
 		 * Since pqGetInt treats 2-byte integers as unsigned, we need to
 		 * coerce these results to signed form.
 		 */
 		format = (int) ((int16) format);
-		result->attDescs[i].format = format;
+		if (result && result->attDescs)
+			result->attDescs[i].format = format;
 	}
 
 	/* Success! */
 	conn->result = result;
 	return 0;
 
-failure:
-	PQclear(result);
+advance_and_error:
+	/* Discard unsaved result, if any */
+	if (result && result != conn->result)
+		PQclear(result);
+
+	/* Discard the failed message by pretending we read it */
+	conn->inCursor = conn->inStart + 5 + msgLength;
+
+	/*
+	 * Replace partially constructed result with an error result. First
+	 * discard the old result to try to win back some memory.
+	 */
+	pqClearAsyncResult(conn);
+
+	/*
+	 * If preceding code didn't provide an error message, assume "out of
+	 * memory" was meant.  The advantage of having this special case is that
+	 * freeing the old result first greatly improves the odds that gettext()
+	 * will succeed in providing a translation.
+	 */
+	if (!errmsg)
+		errmsg = libpq_gettext("out of memory");
+	printfPQExpBuffer(&conn->errorMessage, "%s\n", errmsg);
+	pqSaveErrorResult(conn);
+
 	return EOF;
 }
 
diff --git a/src/interfaces/libpq/libpq-int.h b/src/interfaces/libpq/libpq-int.h
index 2175957..5fe72eb 100644
--- a/src/interfaces/libpq/libpq-int.h
+++ b/src/interfaces/libpq/libpq-int.h
@@ -219,7 +219,8 @@ typedef enum
 	PGASYNC_READY,				/* result ready for PQgetResult */
 	PGASYNC_COPY_IN,			/* Copy In data transfer in progress */
 	PGASYNC_COPY_OUT,			/* Copy Out data transfer in progress */
-	PGASYNC_COPY_BOTH			/* Copy In/Out data transfer in progress */
+	PGASYNC_COPY_BOTH,			/* Copy In/Out data transfer in progress */
+	PGASYNC_FATAL				/* Something awry happended during execution */
 } PGAsyncStatusType;
 
 /* PGQueryClass tracks which query protocol we are now executing */
-- 
2.4.5

