From 082b9ddb57454776ffe30d6c11a53650ca859a7a Mon Sep 17 00:00:00 2001
From: Abhishek Chanda <abhishek.becs@gmail.com>
Date: Tue, 7 Jan 2025 21:54:08 -0600
Subject: [PATCH v2] Add support for dumping SSL keylog to a file

This patch adds a new connection parameter which is used by libpq
to write keys used in a SSL context
---
 doc/src/sgml/libpq.sgml                  | 22 ++++++++++++++++++++
 src/interfaces/libpq/fe-connect.c        |  5 +++++
 src/interfaces/libpq/fe-secure-openssl.c | 26 ++++++++++++++++++++++++
 src/interfaces/libpq/libpq-int.h         |  1 +
 src/test/perl/PostgreSQL/Test/Utils.pm   |  1 +
 src/test/regress/pg_regress.c            |  1 +
 6 files changed, 56 insertions(+)

diff --git a/doc/src/sgml/libpq.sgml b/doc/src/sgml/libpq.sgml
index 105b22b317..e2d36efe21 100644
--- a/doc/src/sgml/libpq.sgml
+++ b/doc/src/sgml/libpq.sgml
@@ -1909,6 +1909,18 @@ postgresql://%2Fvar%2Flib%2Fpostgresql/dbname
       </listitem>
      </varlistentry>
 
+     <varlistentry id="libpq-connect-sslkeylogfile" xreflabel="sslkeylogfile">
+      <term><literal>sslkeylogfile</literal></term>
+      <listitem>
+       <para>
+        This parameter specifies the location where libpq will log keys
+        used in this SSL context. This is useful for debugging postgres
+        protocol using tools like wireshark. This parameter is ignored if an
+        SSL connection is not made.
+       </para>
+      </listitem>
+     </varlistentry>
+
      <varlistentry id="libpq-connect-sslpassword" xreflabel="sslpassword">
       <term><literal>sslpassword</literal></term>
       <listitem>
@@ -9032,6 +9044,16 @@ myEventProc(PGEventId evtId, void *evtInfo, void *passThrough)
      </para>
     </listitem>
 
+    <listitem>
+     <para>
+      <indexterm>
+       <primary><envar>PGSSLKEYLOGFILE</envar></primary>
+      </indexterm>
+      <envar>PGSSLKEYLOGFILE</envar> behaves the same as the <xref
+      linkend="libpq-connect-sslkeylogfile"/> connection parameter.
+     </para>
+    </listitem>
+
     <listitem>
      <para>
       <indexterm>
diff --git a/src/interfaces/libpq/fe-connect.c b/src/interfaces/libpq/fe-connect.c
index 8f211821eb..806688abc0 100644
--- a/src/interfaces/libpq/fe-connect.c
+++ b/src/interfaces/libpq/fe-connect.c
@@ -366,6 +366,11 @@ static const internalPQconninfoOption PQconninfoOptions[] = {
 		"Load-Balance-Hosts", "", 8,	/* sizeof("disable") = 8 */
 	offsetof(struct pg_conn, load_balance_hosts)},
 
+	{"sslkeylogfile", "SSLKEYLOGFILE",
+		"", NULL,
+		"SSL-Key-Log-File", "", 0,	/* sizeof("") = 0 */
+	offsetof(struct pg_conn, sslkeylogfile)},
+
 	/* Terminating entry --- MUST BE LAST */
 	{NULL, NULL, NULL, NULL,
 	NULL, NULL, 0}
diff --git a/src/interfaces/libpq/fe-secure-openssl.c b/src/interfaces/libpq/fe-secure-openssl.c
index 5bb9d9779d..4c8f86cd31 100644
--- a/src/interfaces/libpq/fe-secure-openssl.c
+++ b/src/interfaces/libpq/fe-secure-openssl.c
@@ -57,6 +57,7 @@
  * include <wincrypt.h>, but some other Windows headers do.)
  */
 #include "common/openssl.h"
+#include <openssl/ssl.h>
 #include <openssl/conf.h>
 #ifdef USE_SSL_ENGINE
 #include <openssl/engine.h>
@@ -86,6 +87,7 @@ static pthread_mutex_t ssl_config_mutex = PTHREAD_MUTEX_INITIALIZER;
 
 static PQsslKeyPassHook_OpenSSL_type PQsslKeyPassHook = NULL;
 static int	ssl_protocol_version_to_openssl(const char *protocol);
+static void SSL_CTX_keylog_cb(const SSL *ssl, const char *line);
 
 /* ------------------------------------------------------------ */
 /*			 Procedures common to all secure sessions			*/
@@ -684,6 +686,27 @@ pgtls_verify_peer_name_matches_certificate_guts(PGconn *conn,
 /* See pqcomm.h comments on OpenSSL implementation of ALPN (RFC 7301) */
 static unsigned char alpn_protos[] = PG_ALPN_PROTOCOL_VECTOR;
 
+/* This is a callback that writes to a given SSLKEYLOGFILE file */
+static void SSL_CTX_keylog_cb(const SSL *ssl, const char *line) {
+	FILE *log_file;
+	PGconn *conn = SSL_get_app_data(ssl);
+	if (conn == NULL)
+		return;
+
+	log_file = fopen(conn->sslkeylogfile, "a");
+	if (log_file == NULL) {
+		libpq_append_conn_error(conn, "could not open ssl key log file");
+		return;
+	}
+
+	if (chmod(conn->sslkeylogfile, 0600) == -1) {
+		libpq_append_conn_error(conn, "could not chmod ssl key log file");
+		return;
+	}
+	fprintf(log_file, "%s\n", line);
+	fclose(log_file);
+}
+
 /*
  *	Create per-connection SSL object, and load the client certificate,
  *	private key, and trusted CA certs.
@@ -1000,6 +1023,9 @@ initialize_SSL(PGconn *conn)
 	}
 	conn->ssl_in_use = true;
 
+	if (conn->sslkeylogfile && strlen(conn->sslkeylogfile) > 0)
+		SSL_CTX_set_keylog_callback(SSL_context, SSL_CTX_keylog_cb);
+
 	/*
 	 * SSL contexts are reference counted by OpenSSL. We can free it as soon
 	 * as we have created the SSL object, and it will stick around for as long
diff --git a/src/interfaces/libpq/libpq-int.h b/src/interfaces/libpq/libpq-int.h
index 4a5a7c8b5e..57b8f7e602 100644
--- a/src/interfaces/libpq/libpq-int.h
+++ b/src/interfaces/libpq/libpq-int.h
@@ -428,6 +428,7 @@ struct pg_conn
 	char	   *target_session_attrs;	/* desired session properties */
 	char	   *require_auth;	/* name of the expected auth method */
 	char	   *load_balance_hosts; /* load balance over hosts */
+	char	   *sslkeylogfile;	/* where should the client write ssl key logs */
 
 	bool		cancelRequest;	/* true if this connection is used to send a
 								 * cancel request, instead of being a normal
diff --git a/src/test/perl/PostgreSQL/Test/Utils.pm b/src/test/perl/PostgreSQL/Test/Utils.pm
index 9c83d93f79..eda8c094ab 100644
--- a/src/test/perl/PostgreSQL/Test/Utils.pm
+++ b/src/test/perl/PostgreSQL/Test/Utils.pm
@@ -129,6 +129,7 @@ BEGIN
 	  PGSSLCRL
 	  PGSSLCRLDIR
 	  PGSSLKEY
+	  PGSSLKEYLOGFILE
 	  PGSSLMAXPROTOCOLVERSION
 	  PGSSLMINPROTOCOLVERSION
 	  PGSSLMODE
diff --git a/src/test/regress/pg_regress.c b/src/test/regress/pg_regress.c
index cbef6d48d3..6e3dfc11af 100644
--- a/src/test/regress/pg_regress.c
+++ b/src/test/regress/pg_regress.c
@@ -830,6 +830,7 @@ initialize_environment(void)
 		unsetenv("PGSSLCRL");
 		unsetenv("PGSSLCRLDIR");
 		unsetenv("PGSSLKEY");
+		unsetenv("PGSSLKEYLOGFILE");
 		unsetenv("PGSSLMAXPROTOCOLVERSION");
 		unsetenv("PGSSLMINPROTOCOLVERSION");
 		unsetenv("PGSSLMODE");
-- 
2.47.1

