v11-0007-POC-libpq-encryption-protocol.patch

From 2a976115128b61efd9f85bbc045c48d2679d9f82 Mon Sep 17 00:00:00 2001
From: Masahiko Sawada <sawada.mshk@gmail.com>
Date: Tue, 19 May 2020 16:57:57 +0900
Subject: [PATCH v11 7/7] POC: libpq encryption protocol.

---
 src/bin/psql/command.c              |  46 +++++++++++
 src/interfaces/libpq/exports.txt    |   1 +
 src/interfaces/libpq/fe-exec.c      | 119 ++++++++++++++++++++++++++++
 src/interfaces/libpq/fe-protocol3.c |   6 +-
 src/interfaces/libpq/libpq-fe.h     |   1 +
 src/interfaces/libpq/libpq-int.h    |   2 +-
 6 files changed, 172 insertions(+), 3 deletions(-)

diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index a5160f91de..3039abbc42 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -129,6 +129,8 @@ static backslashResult exec_command_write(PsqlScanState scan_state, bool active_
 										  PQExpBuffer query_buf, PQExpBuffer previous_buf);
 static backslashResult exec_command_watch(PsqlScanState scan_state, bool active_branch,
 										  PQExpBuffer query_buf, PQExpBuffer previous_buf);
+static backslashResult exec_command_cipher(PsqlScanState scan_state, bool active_branch,
+										   bool for_enc);
 static backslashResult exec_command_x(PsqlScanState scan_state, bool active_branch);
 static backslashResult exec_command_z(PsqlScanState scan_state, bool active_branch);
 static backslashResult exec_command_shell_escape(PsqlScanState scan_state, bool active_branch);
@@ -398,6 +400,8 @@ exec_command(const char *cmd,
 	else if (strcmp(cmd, "watch") == 0)
 		status = exec_command_watch(scan_state, active_branch,
 									query_buf, previous_buf);
+	else if (strcmp(cmd, "encrypt") == 0)
+		status = exec_command_cipher(scan_state, active_branch, true);
 	else if (strcmp(cmd, "x") == 0)
 		status = exec_command_x(scan_state, active_branch);
 	else if (strcmp(cmd, "z") == 0)
@@ -2656,6 +2660,48 @@ exec_command_watch(PsqlScanState scan_state, bool active_branch,
 	return success ? PSQL_CMD_SKIP_LINE : PSQL_CMD_ERROR;
 }
 
+static backslashResult
+exec_command_cipher(PsqlScanState scan_state, bool active_branch, bool for_enc)
+{
+	bool success = true;
+
+	if (active_branch)
+	{
+		char token1[128];
+		char token2[128];
+		char *res;
+
+		simple_prompt("Enter data: ", token1, sizeof(token1), false);
+		simple_prompt("Enter it again: ", token2, sizeof(token2), false);
+
+		if (strcmp(token1, token2) != 0)
+		{
+			pg_log_error("Data didn't match.");
+			success = false;
+		}
+		else
+		{
+			res = PQencrypt(pset.db, token1);
+
+			if (!res)
+			{
+				pg_log_info("%s", PQerrorMessage(pset.db));
+				success = false;
+			}
+			else
+			{
+				fprintf(stdout, "encrypted data: %s\n", res);
+				fflush(stdout);
+				free(res);
+			}
+		}
+	}
+	else
+		ignore_slash_options(scan_state);
+
+	return success ? PSQL_CMD_SKIP_LINE : PSQL_CMD_ERROR;
+}
+
 /*
  * \x -- set or toggle expanded table representation
  */
diff --git a/src/interfaces/libpq/exports.txt b/src/interfaces/libpq/exports.txt
index bbc1f90481..a3a84951dc 100644
--- a/src/interfaces/libpq/exports.txt
+++ b/src/interfaces/libpq/exports.txt
@@ -179,3 +179,4 @@ PQgetgssctx               176
 PQsetSSLKeyPassHook_OpenSSL         177
 PQgetSSLKeyPassHook_OpenSSL         178
 PQdefaultSSLKeyPassHook_OpenSSL     179
+PQencrypt                 180
diff --git a/src/interfaces/libpq/fe-exec.c b/src/interfaces/libpq/fe-exec.c
index eea0237c3a..67d291696d 100644
--- a/src/interfaces/libpq/fe-exec.c
+++ b/src/interfaces/libpq/fe-exec.c
@@ -27,6 +27,7 @@
 #include "libpq-fe.h"
 #include "libpq-int.h"
 #include "mb/pg_wchar.h"
+#include "common/aead.h"
 
 /* keep this in same order as ExecStatusType in libpq-fe.h */
 char	   *const pgresStatus[] = {
@@ -69,6 +70,13 @@ static bool PQexecStart(PGconn *conn);
 static PGresult *PQexecFinish(PGconn *conn);
 static int	PQsendDescribe(PGconn *conn, char desc_type,
 						   const char *desc_target);
+static PGresult *PQfnText(PGconn *conn,
+						  int fnid,
+						  int *result_buf,
+						  int *result_len,
+						  int result_is_int,
+						  const PQArgBlock *args,
+						  int nargs);
 static int	check_field_number(const PGresult *res, int field_num);
 
 
@@ -2677,6 +2685,7 @@ PQfn(PGconn *conn,
 		return pqFunctionCall3(conn, fnid,
 							   result_buf, result_len,
 							   result_is_int,
+							   1,	/* output format is binary */
 							   args, nargs);
 	else
 		return pqFunctionCall2(conn, fnid,
@@ -2685,6 +2694,45 @@ PQfn(PGconn *conn,
 							   args, nargs);
 }
 
+static PGresult *
+PQfnText(PGconn *conn,
+		 int fnid,
+		 int *result_buf,
+		 int *result_len,
+		 int result_is_int,
+		 const PQArgBlock *args,
+		 int nargs)
+{
+	*result_len = 0;
+
+	if (!conn)
+		return NULL;
+
+	/* This isn't gonna work on a 2.0 server */
+	if (PG_PROTOCOL_MAJOR(conn->pversion) < 3)
+	{
+		printfPQExpBuffer(&conn->errorMessage,
+						  libpq_gettext("function requires at least protocol version 3.0\n"));
+		return NULL;
+	}
+
+	/* clear the error string */
+	resetPQExpBuffer(&conn->errorMessage);
+
+	if (conn->sock == PGINVALID_SOCKET || conn->asyncStatus != PGASYNC_IDLE ||
+		conn->result != NULL)
+	{
+		printfPQExpBuffer(&conn->errorMessage,
+						  libpq_gettext("connection in wrong state\n"));
+		return NULL;
+	}
+
+	return pqFunctionCall3(conn, fnid,
+						   result_buf, result_len,
+						   result_is_int,
+						   0,	/* output format is text */
+						   args, nargs);
+}
 
 /* ====== accessor funcs for PGresult ======== */
 
@@ -3858,3 +3906,74 @@ PQunescapeBytea(const unsigned char *strtext, size_t *retbuflen)
 	*retbuflen = buflen;
 	return tmpbuf;
 }
+
+char *
+PQencrypt(PGconn *conn, const char *token)
+{
+	static Oid	foid = InvalidOid;
+	PQArgBlock	argv[1];
+	PGresult	*res;
+	int			reslen;
+	int			tokenlen;
+	int			*buf;
+
+	/* This isn't gonna work on a 2.0 server */
+	if (PG_PROTOCOL_MAJOR(conn->pversion) < 3)
+	{
+		printfPQExpBuffer(&conn->errorMessage,
+						  libpq_gettext("function requires at least protocol version 3.0\n"));
+		return 0;
+	}
+
+	if (!OidIsValid(foid))
+	{
+		const char *query;
+		char *fname;
+
+		query = "select proname, oid from pg_proc "
+			"where proname = 'pg_encrypt'";
+		res = PQexec(conn, query);
+
+		if (res == NULL)
+		{
+			PQclear(res);
+			printfPQExpBuffer(&conn->errorMessage,
+							  libpq_gettext("query to initialize encryption functions did not return data\n"));
+			return NULL;
+		}
+
+		fname = PQgetvalue(res, 0, 0);
+
+		if (PQntuples(res) > 1 || strcmp(fname, "pg_encrypt") != 0)
+		{
+			PQclear(res);
+			printfPQExpBuffer(&conn->errorMessage,
+							  libpq_gettext("query to initialize encryption functions returned invalid data\n"));
+			return NULL;
+		}
+
+		/* get OID of pg_encrypt function */
+		foid = (Oid) atoi(PQgetvalue(res, 0, 1));
+
+		PQclear(res);
+	}
+
+	tokenlen = strlen(token);
+	buf = (int *) malloc(AEADSizeOfCipherText(tokenlen));
+
+	argv[0].isint = 0;
+	argv[0].len = strlen(token);
+	argv[0].u.ptr = (int *) token;
+	res = PQfnText(conn, foid, buf, &reslen, 0, argv, 1);
+
+	if (PQresultStatus(res) == PGRES_COMMAND_OK)
+	{
+		PQclear(res);
+		return (char *) buf;
+	}
+	else
+	{
+		PQclear(res);
+		return NULL;
+	}
+}
diff --git a/src/interfaces/libpq/fe-protocol3.c b/src/interfaces/libpq/fe-protocol3.c
index 1696525475..2179143a11 100644
--- a/src/interfaces/libpq/fe-protocol3.c
+++ b/src/interfaces/libpq/fe-protocol3.c
@@ -1922,7 +1922,7 @@ pqEndcopy3(PGconn *conn)
 PGresult *
 pqFunctionCall3(PGconn *conn, Oid fnid,
 				int *result_buf, int *actual_result_len,
-				int result_is_int,
+				int result_is_int, int result_format,
 				const PQArgBlock *args, int nargs)
 {
 	bool		needInput = false;
@@ -1930,6 +1930,7 @@ pqFunctionCall3(PGconn *conn, Oid fnid,
 	char		id;
 	int			msgLength;
 	int			avail;
+	int			format;
 	int			i;
 
 	/* PQfn already validated connection state */
@@ -1963,7 +1964,8 @@ pqFunctionCall3(PGconn *conn, Oid fnid,
 		}
 	}
 
-	if (pqPutInt(1, 2, conn) < 0)	/* result format code: BINARY */
+	format = result_is_int ? 1 : result_format;
+	if (pqPutInt(format, 2, conn) < 0)	/* 0: TEXT, 1: BINARY */
 		return NULL;
 
 	if (pqPutMsgEnd(conn) < 0 ||
diff --git a/src/interfaces/libpq/libpq-fe.h b/src/interfaces/libpq/libpq-fe.h
index 3b6a9fbce3..2e5ae05ef8 100644
--- a/src/interfaces/libpq/libpq-fe.h
+++ b/src/interfaces/libpq/libpq-fe.h
@@ -457,6 +457,7 @@ extern int	PQisthreadsafe(void);
 extern PGPing PQping(const char *conninfo);
 extern PGPing PQpingParams(const char *const *keywords,
 						   const char *const *values, int expand_dbname);
+extern char *PQencrypt(PGconn *conn, const char *token);
 
 /* Force the write buffer to be written (or at least try) */
 extern int	PQflush(PGconn *conn);
diff --git a/src/interfaces/libpq/libpq-int.h b/src/interfaces/libpq/libpq-int.h
index 1de91ae295..f479aa8412 100644
--- a/src/interfaces/libpq/libpq-int.h
+++ b/src/interfaces/libpq/libpq-int.h
@@ -643,7 +643,7 @@ extern int	pqGetlineAsync3(PGconn *conn, char *buffer, int bufsize);
 extern int	pqEndcopy3(PGconn *conn);
 extern PGresult *pqFunctionCall3(PGconn *conn, Oid fnid,
 								 int *result_buf, int *actual_result_len,
-								 int result_is_int,
+								 int result_is_int, int result_format,
 								 const PQArgBlock *args, int nargs);
 
 /* === in fe-misc.c === */
-- 
2.23.0