Custom allocators in libpq

From be463a2da13734d5b104a5fc9c58e7bbb8908323 Mon Sep 17 00:00:00 2001
From: Aaron Patterson <aaron.patterson@gmail.com>
Date: Thu, 19 Nov 2015 14:44:49 -0800
Subject: [PATCH 1/2] add PQunescapeByteaConn for unescaping bytes given a
 connection

This commit adds an unescape method that takes a connection as a
parameter.  The connection may have settings on it (like allocators)
that are important to unescaping the string.
---
 src/interfaces/libpq/exports.txt |  1 +
 src/interfaces/libpq/fe-exec.c   | 21 +++++++++++++++++++--
 src/interfaces/libpq/libpq-fe.h  |  6 ++++--
 3 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/src/interfaces/libpq/exports.txt b/src/interfaces/libpq/exports.txt
index d6a38d0df8..ba718f4071 100644
--- a/src/interfaces/libpq/exports.txt
+++ b/src/interfaces/libpq/exports.txt
@@ -172,3 +172,4 @@ PQsslAttribute            169
 PQsetErrorContextVisibility 170
 PQresultVerboseErrorMessage 171
 PQencryptPasswordConn     172
+PQunescapeByteaConn       173
diff --git a/src/interfaces/libpq/fe-exec.c b/src/interfaces/libpq/fe-exec.c
index e1e2d18e3a..464bdc635f 100644
--- a/src/interfaces/libpq/fe-exec.c
+++ b/src/interfaces/libpq/fe-exec.c
@@ -3651,8 +3651,9 @@ PQescapeBytea(const unsigned char *from, size_t from_length, size_t *to_length)
  *		\ooo == a byte whose value = ooo (ooo is an octal number)
  *		\x	 == x (x is any character not matched by the above transformations)
  */
-unsigned char *
-PQunescapeBytea(const unsigned char *strtext, size_t *retbuflen)
+static unsigned char *
+PQunescapeByteaInternal(PGconn *conn,
+					  const unsigned char *strtext, size_t *retbuflen)
 {
 	size_t		strtextlen,
 				buflen;
@@ -3762,3 +3763,19 @@ PQunescapeBytea(const unsigned char *strtext, size_t *retbuflen)
 	*retbuflen = buflen;
 	return tmpbuf;
 }
+
+unsigned char *
+PQunescapeByteaConn(PGconn *conn,
+				  const unsigned char *strtext, size_t *retbuflen)
+{
+	if (!conn)
+	    return NULL;
+
+	return PQunescapeByteaInternal(conn, strtext, retbuflen);
+}
+
+unsigned char *
+PQunescapeBytea(const unsigned char *strtext, size_t *retbuflen)
+{
+	return PQunescapeByteaInternal(NULL, strtext, retbuflen);
+}
diff --git a/src/interfaces/libpq/libpq-fe.h b/src/interfaces/libpq/libpq-fe.h
index 1d915e7915..7fce4b387e 100644
--- a/src/interfaces/libpq/libpq-fe.h
+++ b/src/interfaces/libpq/libpq-fe.h
@@ -527,13 +527,15 @@ extern char *PQescapeIdentifier(PGconn *conn, const char *str, size_t len);
 extern unsigned char *PQescapeByteaConn(PGconn *conn,
 				  const unsigned char *from, size_t from_length,
 				  size_t *to_length);
-extern unsigned char *PQunescapeBytea(const unsigned char *strtext,
-				size_t *retbuflen);
+extern unsigned char *PQunescapeByteaConn(PGconn *conn,
+                                  const unsigned char *strtext, size_t *retbuflen);
 
 /* These forms are deprecated! */
 extern size_t PQescapeString(char *to, const char *from, size_t length);
 extern unsigned char *PQescapeBytea(const unsigned char *from, size_t from_length,
 			  size_t *to_length);
+extern unsigned char *PQunescapeBytea(const unsigned char *strtext,
+				size_t *retbuflen);
 
 
 
-- 
2.11.0

From b62fed4a7e38efeeee9e2057d5556312a5a76cdf Mon Sep 17 00:00:00 2001
From: Aaron Patterson <aaron.patterson@gmail.com>
Date: Thu, 19 Nov 2015 15:28:27 -0800
Subject: [PATCH 2/2] use configured malloc / free from the connection.

---
 src/interfaces/libpq/fe-connect.c |  5 +++
 src/interfaces/libpq/fe-exec.c    | 87 ++++++++++++++++++++++++++++-----------
 src/interfaces/libpq/libpq-int.h  | 19 +++++++++
 3 files changed, 87 insertions(+), 24 deletions(-)

diff --git a/src/interfaces/libpq/fe-connect.c b/src/interfaces/libpq/fe-connect.c
index d0e97ecdd4..203bbbc310 100644
--- a/src/interfaces/libpq/fe-connect.c
+++ b/src/interfaces/libpq/fe-connect.c
@@ -3336,6 +3336,11 @@ makeEmptyPGconn(void)
 	conn->noticeHooks.noticeRec = defaultNoticeReceiver;
 	conn->noticeHooks.noticeProc = defaultNoticeProcessor;
 
+	/* set default allocators */
+	conn->malloc = malloc;
+	conn->realloc = realloc;
+	conn->free = free;
+
 	conn->status = CONNECTION_BAD;
 	conn->asyncStatus = PGASYNC_IDLE;
 	conn->xactStatus = PQTRANS_IDLE;
diff --git a/src/interfaces/libpq/fe-exec.c b/src/interfaces/libpq/fe-exec.c
index 464bdc635f..83a8bb627c 100644
--- a/src/interfaces/libpq/fe-exec.c
+++ b/src/interfaces/libpq/fe-exec.c
@@ -141,7 +141,7 @@ PQmakeEmptyPGresult(PGconn *conn, ExecStatusType status)
 {
 	PGresult   *result;
 
-	result = (PGresult *) malloc(sizeof(PGresult));
+	result = (PGresult *) conn->malloc(sizeof(PGresult));
 	if (!result)
 		return NULL;
 
@@ -164,6 +164,9 @@ PQmakeEmptyPGresult(PGconn *conn, ExecStatusType status)
 	result->curBlock = NULL;
 	result->curOffset = 0;
 	result->spaceLeft = 0;
+	result->malloc = conn->malloc;
+	result->realloc = conn->realloc;
+	result->free = conn->free;
 
 	if (conn)
 	{
@@ -546,7 +549,7 @@ pqResultAlloc(PGresult *res, size_t nBytes, bool isBinary)
 	 */
 	if (nBytes >= PGRESULT_SEP_ALLOC_THRESHOLD)
 	{
-		block = (PGresult_data *) malloc(nBytes + PGRESULT_BLOCK_OVERHEAD);
+		block = (PGresult_data *) res->malloc(nBytes + PGRESULT_BLOCK_OVERHEAD);
 		if (!block)
 			return NULL;
 		space = block->space + PGRESULT_BLOCK_OVERHEAD;
@@ -570,7 +573,7 @@ pqResultAlloc(PGresult *res, size_t nBytes, bool isBinary)
 	}
 
 	/* Otherwise, start a new block. */
-	block = (PGresult_data *) malloc(PGRESULT_DATA_BLOCKSIZE);
+	block = (PGresult_data *) res->malloc(PGRESULT_DATA_BLOCKSIZE);
 	if (!block)
 		return NULL;
 	block->next = res->curBlock;
@@ -666,22 +669,22 @@ PQclear(PGresult *res)
 			(void) res->events[i].proc(PGEVT_RESULTDESTROY, &evt,
 									   res->events[i].passThrough);
 		}
-		free(res->events[i].name);
+		res->free(res->events[i].name);
 	}
 
 	if (res->events)
-		free(res->events);
+		res->free(res->events);
 
 	/* Free all the subsidiary blocks */
 	while ((block = res->curBlock) != NULL)
 	{
 		res->curBlock = block->next;
-		free(block);
+		res->free(block);
 	}
 
 	/* Free the top-level tuple pointer array */
 	if (res->tuples)
-		free(res->tuples);
+		res->free(res->tuples);
 
 	/* zero out the pointer fields to catch programming errors */
 	res->attDescs = NULL;
@@ -693,7 +696,7 @@ PQclear(PGresult *res)
 	/* res->curBlock was zeroed out earlier */
 
 	/* Free the PGresult structure itself */
-	free(res);
+	res->free(res);
 }
 
 /*
@@ -870,10 +873,10 @@ pqAddTuple(PGresult *res, PGresAttValue *tup)
 
 		if (res->tuples == NULL)
 			newTuples = (PGresAttValue **)
-				malloc(newSize * sizeof(PGresAttValue *));
+				res->malloc(newSize * sizeof(PGresAttValue *));
 		else
 			newTuples = (PGresAttValue **)
-				realloc(res->tuples, newSize * sizeof(PGresAttValue *));
+				res->realloc(res->tuples, newSize * sizeof(PGresAttValue *));
 		if (!newTuples)
 			return FALSE;		/* malloc or realloc failed */
 		res->tupArrSize = newSize;
@@ -931,7 +934,7 @@ pqSaveParameterStatus(PGconn *conn, const char *name, const char *value)
 				prev->next = pstatus->next;
 			else
 				conn->pstatus = pstatus->next;
-			free(pstatus);		/* frees name and value strings too */
+			conn->free(pstatus);		/* frees name and value strings too */
 			break;
 		}
 	}
@@ -939,7 +942,7 @@ pqSaveParameterStatus(PGconn *conn, const char *name, const char *value)
 	/*
 	 * Store new info as a single malloc block
 	 */
-	pstatus = (pgParameterStatus *) malloc(sizeof(pgParameterStatus) +
+	pstatus = (pgParameterStatus *) conn->malloc(sizeof(pgParameterStatus) +
 										   strlen(name) + strlen(value) + 2);
 	if (pstatus)
 	{
@@ -1157,7 +1160,7 @@ PQsendQuery(PGconn *conn, const char *query)
 	/* and remember the query text too, if possible */
 	/* if insufficient memory, last_query just winds up NULL */
 	if (conn->last_query)
-		free(conn->last_query);
+		conn->free(conn->last_query);
 	conn->last_query = strdup(query);
 
 	/*
@@ -1297,7 +1300,7 @@ PQsendPrepare(PGconn *conn,
 	/* and remember the query text too, if possible */
 	/* if insufficient memory, last_query just winds up NULL */
 	if (conn->last_query)
-		free(conn->last_query);
+		conn->free(conn->last_query);
 	conn->last_query = strdup(query);
 
 	/*
@@ -1546,7 +1549,7 @@ PQsendQueryGuts(PGconn *conn,
 	/* and remember the query text too, if possible */
 	/* if insufficient memory, last_query just winds up NULL */
 	if (conn->last_query)
-		free(conn->last_query);
+		conn->free(conn->last_query);
 	if (command)
 		conn->last_query = strdup(command);
 	else
@@ -2159,7 +2162,7 @@ PQsendDescribe(PGconn *conn, char desc_type, const char *desc_target)
 	/* reset last-query string (not relevant now) */
 	if (conn->last_query)
 	{
-		free(conn->last_query);
+		conn->free(conn->last_query);
 		conn->last_query = NULL;
 	}
 
@@ -2866,11 +2869,11 @@ PQfnumber(const PGresult *res, const char *field_name)
 	{
 		if (strcmp(field_case, res->attDescs[i].name) == 0)
 		{
-			free(field_case);
+			res->free(field_case);
 			return i;
 		}
 	}
-	free(field_case);
+	res->free(field_case);
 	return -1;
 }
 
@@ -3391,7 +3394,7 @@ PQescapeInternal(PGconn *conn, const char *str, size_t len, bool as_ident)
 	result_size = input_len + num_quotes + 3;	/* two quotes, plus a NUL */
 	if (!as_ident && num_backslashes > 0)
 		result_size += num_backslashes + 2;
-	result = rp = (char *) malloc(result_size);
+	result = rp = (char *) conn->malloc(result_size);
 	if (rp == NULL)
 	{
 		printfPQExpBuffer(&conn->errorMessage,
@@ -3555,7 +3558,15 @@ PQescapeByteaInternal(PGconn *conn,
 	}
 
 	*to_length = len;
-	rp = result = (unsigned char *) malloc(len);
+	if (conn)
+	{
+		rp = result = (unsigned char *) conn->malloc(len);
+	}
+	else
+	{
+		rp = result = (unsigned char *) malloc(len);
+	}
+
 	if (rp == NULL)
 	{
 		if (conn)
@@ -3674,7 +3685,14 @@ PQunescapeByteaInternal(PGconn *conn,
 
 		buflen = (strtextlen - 2) / 2;
 		/* Avoid unportable malloc(0) */
-		buffer = (unsigned char *) malloc(buflen > 0 ? buflen : 1);
+		if (conn)
+		{
+			buffer = (unsigned char *) conn->malloc(buflen > 0 ? buflen : 1);
+		}
+		else
+		{
+			buffer = (unsigned char *) malloc(buflen > 0 ? buflen : 1);
+		}
 		if (buffer == NULL)
 			return NULL;
 
@@ -3705,7 +3723,14 @@ PQunescapeByteaInternal(PGconn *conn,
 		 * Length of input is max length of output, but add one to avoid
 		 * unportable malloc(0) if input is zero-length.
 		 */
-		buffer = (unsigned char *) malloc(strtextlen + 1);
+		if (conn)
+		{
+			buffer = (unsigned char *) conn->malloc(strtextlen + 1);
+		}
+		else
+		{
+			buffer = (unsigned char *) malloc(strtextlen + 1);
+		}
 		if (buffer == NULL)
 			return NULL;
 
@@ -3751,12 +3776,26 @@ PQunescapeByteaInternal(PGconn *conn,
 
 	/* Shrink the buffer to be no larger than necessary */
 	/* +1 avoids unportable behavior when buflen==0 */
-	tmpbuf = realloc(buffer, buflen + 1);
+	if (conn)
+	{
+		tmpbuf = conn->realloc(buffer, buflen + 1);
+	}
+	else
+	{
+		tmpbuf = realloc(buffer, buflen + 1);
+	}
 
 	/* It would only be a very brain-dead realloc that could fail, but... */
 	if (!tmpbuf)
 	{
-		free(buffer);
+		if (conn)
+		{
+			conn->free(buffer);
+		}
+		else
+		{
+			free(buffer);
+		}
 		return NULL;
 	}
 
diff --git a/src/interfaces/libpq/libpq-int.h b/src/interfaces/libpq/libpq-int.h
index 42913604e3..356be9ccbc 100644
--- a/src/interfaces/libpq/libpq-int.h
+++ b/src/interfaces/libpq/libpq-int.h
@@ -80,6 +80,14 @@ typedef struct
 #endif
 #endif							/* USE_OPENSSL */
 
+ /* ----------------
+ * malloc/realloc/free for PGResult is replasable for in-backend use
+ * ----------------
+ */
+typedef void *(*PQpgresult_malloc)(size_t size);
+typedef void *(*PQpgresult_realloc)(void *ptr, size_t size);
+typedef void (*PQpgresult_free)(void *ptr);
+
 /*
  * POSTGRES backend dependent Constants.
  */
@@ -162,6 +170,9 @@ typedef struct PGEvent
 	void	   *passThrough;	/* pointer supplied at registration time */
 	void	   *data;			/* optional state (instance) data */
 	bool		resultInitialized;	/* T if RESULTCREATE/COPY succeeded */
+	PQpgresult_malloc malloc;
+	PQpgresult_realloc realloc;
+	PQpgresult_free free;
 } PGEvent;
 
 struct pg_result
@@ -208,6 +219,10 @@ struct pg_result
 	PGresult_data *curBlock;	/* most recently allocated block */
 	int			curOffset;		/* start offset of free space in block */
 	int			spaceLeft;		/* number of free bytes remaining in block */
+
+	PQpgresult_malloc malloc;
+	PQpgresult_realloc realloc;
+	PQpgresult_free free;
 };
 
 /* PGAsyncStatusType defines the state of the query-execution state machine */
@@ -492,6 +507,10 @@ struct pg_conn
 
 	/* Buffer for receiving various parts of messages */
 	PQExpBufferData workBuffer; /* expansible string */
+
+	PQpgresult_malloc malloc;
+	PQpgresult_realloc realloc;
+	PQpgresult_free free;
 };
 
 /* PGcancel stores all data necessary to cancel a connection. A copy of this
-- 
2.11.0