1 files changed, 47 insertions, 0 deletions

diff --git a/src/interfaces/libpq/fe-connect.c b/src/interfaces/libpq/fe-connect.c
index 8f211821eb2..c7943d549e8 100644
--- a/src/interfaces/libpq/fe-connect.c
+++ b/src/interfaces/libpq/fe-connect.c
@@ -22,6 +22,7 @@
 #include <time.h>
 #include <unistd.h>
 
+#include "common/base64.h"
 #include "common/ip.h"
 #include "common/link-canary.h"
 #include "common/scram-common.h"
@@ -366,6 +367,12 @@ static const internalPQconninfoOption PQconninfoOptions[] = {
 		"Load-Balance-Hosts", "", 8,	/* sizeof("disable") = 8 */
 	offsetof(struct pg_conn, load_balance_hosts)},
 
+	{"scram_client_key", NULL, NULL, NULL, "SCRAM-Client-Key", "D", SCRAM_MAX_KEY_LEN * 2,
+	offsetof(struct pg_conn, scram_client_key)},
+
+	{"scram_server_key", NULL, NULL, NULL, "SCRAM-Server-Key", "D", SCRAM_MAX_KEY_LEN * 2,
+	offsetof(struct pg_conn, scram_server_key)},
+
 	/* Terminating entry --- MUST BE LAST */
 	{NULL, NULL, NULL, NULL,
 	NULL, NULL, 0}
@@ -1793,6 +1800,44 @@ pqConnectOptions2(PGconn *conn)
 	else
 		conn->target_server_type = SERVER_TYPE_ANY;
 
+	if (conn->scram_client_key)
+	{
+		int			len;
+
+		len = pg_b64_dec_len(strlen(conn->scram_client_key));
+		/* Consider the zero-terminator */
+		if (len != SCRAM_MAX_KEY_LEN + 1)
+		{
+			libpq_append_conn_error(conn, "invalid SCRAM client key length: %d", len);
+			return false;
+		}
+		conn->scram_client_key_len = len;
+		conn->scram_client_key_binary = malloc(len);
+		if (!conn->scram_client_key_binary)
+			goto oom_error;
+		pg_b64_decode(conn->scram_client_key, strlen(conn->scram_client_key),
+					  conn->scram_client_key_binary, len);
+	}
+
+	if (conn->scram_server_key)
+	{
+		int			len;
+
+		len = pg_b64_dec_len(strlen(conn->scram_server_key));
+		/* Consider the zero-terminator */
+		if (len != SCRAM_MAX_KEY_LEN + 1)
+		{
+			libpq_append_conn_error(conn, "invalid SCRAM server key length: %d", len);
+			return false;
+		}
+		conn->scram_server_key_len = len;
+		conn->scram_server_key_binary = malloc(len);
+		if (!conn->scram_server_key_binary)
+			goto oom_error;
+		pg_b64_decode(conn->scram_server_key, strlen(conn->scram_server_key),
+					  conn->scram_server_key_binary, len);
+	}
+
 	/*
 	 * validate load_balance_hosts option, and set load_balance_type
 	 */
@@ -4704,6 +4749,8 @@ freePGconn(PGconn *conn)
 	free(conn->rowBuf);
 	free(conn->target_session_attrs);
 	free(conn->load_balance_hosts);
+	free(conn->scram_client_key);
+	free(conn->scram_server_key);
 	termPQExpBuffer(&conn->errorMessage);
 	termPQExpBuffer(&conn->workBuffer);