From b205c48c0137333074b770a1e5568a984e3e29e1 Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@otacoo.com>
Date: Fri, 4 Sep 2015 16:22:03 +0900
Subject: [PATCH 1/2] Prevent COPY start and BIND to freeze in libpq on OOM

A new internal flag PGAYNC_FATAL is used to control the error flow.
---
 src/interfaces/libpq/fe-exec.c      |  13 ++++
 src/interfaces/libpq/fe-protocol3.c | 145 ++++++++++++++++++++++++++++++------
 src/interfaces/libpq/libpq-int.h    |   3 +-
 3 files changed, 136 insertions(+), 25 deletions(-)

diff --git a/src/interfaces/libpq/fe-exec.c b/src/interfaces/libpq/fe-exec.c
index 828f18e..d40a39f 100644
--- a/src/interfaces/libpq/fe-exec.c
+++ b/src/interfaces/libpq/fe-exec.c
@@ -1744,6 +1744,14 @@ PQgetResult(PGconn *conn)
 		case PGASYNC_COPY_BOTH:
 			res = getCopyResult(conn, PGRES_COPY_BOTH);
 			break;
+		case PGASYNC_FATAL:
+			res = NULL;
+			/*
+			 * Set the state back to BUSY, allowing parsing to proceed to
+			 * consume messages coming from the server.
+			 */
+			conn->asyncStatus = PGASYNC_BUSY;
+			break;
 		default:
 			printfPQExpBuffer(&conn->errorMessage,
 							  libpq_gettext("unexpected asyncStatus: %d\n"),
@@ -1991,6 +1999,8 @@ PQexecFinish(PGconn *conn)
 	 * We have to stop if we see copy in/out/both, however. We will resume
 	 * parsing after application performs the data transfer.
 	 *
+	 * Stop if a fatal error happened.
+	 *
 	 * Also stop if the connection is lost (else we'll loop infinitely).
 	 */
 	lastResult = NULL;
@@ -2015,6 +2025,9 @@ PQexecFinish(PGconn *conn)
 				PQclear(lastResult);
 		}
 		lastResult = result;
+		if (conn->asyncStatus == PGASYNC_FATAL &&
+			result->resultStatus == PGRES_FATAL_ERROR)
+			break;
 		if (result->resultStatus == PGRES_COPY_IN ||
 			result->resultStatus == PGRES_COPY_OUT ||
 			result->resultStatus == PGRES_COPY_BOTH ||
diff --git a/src/interfaces/libpq/fe-protocol3.c b/src/interfaces/libpq/fe-protocol3.c
index 641804c..0a360e4 100644
--- a/src/interfaces/libpq/fe-protocol3.c
+++ b/src/interfaces/libpq/fe-protocol3.c
@@ -45,11 +45,13 @@
 
 static void handleSyncLoss(PGconn *conn, char id, int msgLength);
 static int	getRowDescriptions(PGconn *conn, int msgLength);
-static int	getParamDescriptions(PGconn *conn);
+static int	getParamDescriptions(PGconn *conn, int msgLength);
 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);
@@ -329,8 +331,11 @@ pqParseInput3(PGconn *conn)
 					}
 					break;
 				case 't':		/* Parameter Description */
-					if (getParamDescriptions(conn))
+					if (getParamDescriptions(conn, msgLength))
+					{
+						conn->asyncStatus = PGASYNC_FATAL;
 						return;
+					}
 					break;
 				case 'D':		/* Data Row */
 					if (conn->result != NULL &&
@@ -362,19 +367,34 @@ 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;
+					}
+
+					/* getCopyStart() moves inStart itself */
 					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;
+					}
+
+					/* getCopyStart() moves inStart itself */
 					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;
+					}
+
+					/* getCopyStart() moves inStart itself */
 					conn->asyncStatus = PGASYNC_COPY_BOTH;
 					conn->copy_already_done = 0;
 					break;
@@ -630,27 +650,34 @@ advance_and_error:
 /*
  * parseInput subroutine to read a 't' (ParameterDescription) message.
  * We'll build a new PGresult structure containing the parameter data.
- * Returns: 0 if completed message, EOF if not enough data yet.
+ * Returns: 0 if completed message, EOF in case of error.
  *
  * Note that if we run out of data, we have to release the partially
  * constructed PGresult, and rebuild it again next time.  Fortunately,
  * that shouldn't happen often, since 't' messages usually fit in a packet.
  */
 static int
-getParamDescriptions(PGconn *conn)
+getParamDescriptions(PGconn *conn, int msgLength)
 {
 	PGresult   *result;
 	int			nparams;
 	int			i;
+	const char *errmsg;
 
 	result = PQmakeEmptyPGresult(conn, PGRES_COMMAND_OK);
 	if (!result)
-		goto failure;
+	{
+		errmsg = NULL;
+		goto advance_and_error;
+	}
 
 	/* parseInput already read the 't' label and message length. */
 	/* the next two bytes are the number of parameters */
-	if (pqGetInt(&(result->numParameters), 2, conn))
-		goto failure;
+	if (pqGetInt(&result->numParameters, 2, conn))
+	{
+		errmsg = libpq_gettext("extraneous data in \"t\" message");
+		goto advance_and_error;
+	}
 	nparams = result->numParameters;
 
 	/* allocate space for the parameter descriptors */
@@ -659,7 +686,10 @@ getParamDescriptions(PGconn *conn)
 		result->paramDescs = (PGresParamDesc *)
 			pqResultAlloc(result, nparams * sizeof(PGresParamDesc), TRUE);
 		if (!result->paramDescs)
-			goto failure;
+		{
+			errmsg = NULL;
+			goto advance_and_error;
+		}
 		MemSet(result->paramDescs, 0, nparams * sizeof(PGresParamDesc));
 	}
 
@@ -669,7 +699,10 @@ getParamDescriptions(PGconn *conn)
 		int			typid;
 
 		if (pqGetInt(&typid, 4, conn))
-			goto failure;
+		{
+			errmsg = libpq_gettext("extraneous data in \"t\" message");
+			goto advance_and_error;
+		}
 		result->paramDescs[i].typid = typid;
 	}
 
@@ -677,8 +710,31 @@ getParamDescriptions(PGconn *conn)
 	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->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;
 }
 
@@ -1343,22 +1399,33 @@ 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 */
@@ -1367,7 +1434,10 @@ getCopyStart(PGconn *conn, ExecStatusType copytype)
 		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));
 	}
 
@@ -1376,7 +1446,10 @@ 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
@@ -1390,8 +1463,32 @@ getCopyStart(PGconn *conn, ExecStatusType copytype)
 	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->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);
+
+	/* Build a error result that will be fetched by PQgetResult */
+	pqSaveErrorResult(conn);
 	return EOF;
 }
 
diff --git a/src/interfaces/libpq/libpq-int.h b/src/interfaces/libpq/libpq-int.h
index 04d056e..6eebcad 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 happened */
 } PGAsyncStatusType;
 
 /* PGQueryClass tracks which query protocol we are now executing */
-- 
2.5.1

