kms_v4_fe.patch

diff --git a/src/bin/pg_waldump/pg_waldump.c b/src/bin/pg_waldump/pg_waldump.c
index 83202b5b87..38308d2160 100644
--- a/src/bin/pg_waldump/pg_waldump.c
+++ b/src/bin/pg_waldump/pg_waldump.c
@@ -24,11 +24,16 @@
 #include "common/logging.h"
 #include "getopt_long.h"
 #include "rmgrdesc.h"
+#include "common/kmgr_utils.h"
+#include "common/cipher.h"
 
 static const char *progname;
 
 static int	WalSegSz;
 
+pg_cipher_ctx *encrypt_cipher_ctx = NULL;
+pg_cipher_ctx *decrypt_cipher_ctx = NULL;
+
 typedef struct XLogDumpPrivate
 {
 	TimeLineID	timeline;
@@ -714,6 +719,9 @@ usage(void)
 	printf(_("  %s [OPTION]... [STARTSEG [ENDSEG]]\n"), progname);
 	printf(_("\nOptions:\n"));
 	printf(_("  -b, --bkp-details      output detailed information about backup blocks\n"));
+	printf(_("  -c  --cluster-passphrase-command=COMMAND\n"
+			 "			set command to obtain passphrase for data encryption key\n"));
+	printf(_(" [-D, --pgdata=]DATADIR  data directory to get the Data encryption cipher\n"));
 	printf(_("  -e, --end=RECPTR       stop reading at WAL location RECPTR\n"));
 	printf(_("  -f, --follow           keep retrying after reaching end of WAL\n"));
 	printf(_("  -n, --limit=N          number of records to display\n"));
@@ -746,9 +754,14 @@ main(int argc, char **argv)
 	XLogRecPtr	first_record;
 	char	   *waldir = NULL;
 	char	   *errormsg;
+	char	   *pgdataDir = NULL;
+	char	   *cluster_passphrase = NULL;
+	bool		ret = false;
 
 	static struct option long_options[] = {
 		{"bkp-details", no_argument, NULL, 'b'},
+		{"pgdata", required_argument, NULL, 'D'},
+		{"cluster-passphrase-command", required_argument, NULL, 'c'},
 		{"end", required_argument, NULL, 'e'},
 		{"follow", no_argument, NULL, 'f'},
 		{"help", no_argument, NULL, '?'},
@@ -809,7 +822,7 @@ main(int argc, char **argv)
 		goto bad_argument;
 	}
 
-	while ((option = getopt_long(argc, argv, "be:fn:p:r:s:t:x:z",
+	while ((option = getopt_long(argc, argv, "D:c:be:fn:p:r:s:t:x:z",
 								 long_options, &optindex)) != -1)
 	{
 		switch (option)
@@ -817,6 +830,12 @@ main(int argc, char **argv)
 			case 'b':
 				config.bkp_details = true;
 				break;
+			case 'c':
+				cluster_passphrase = pg_strdup(optarg);
+				break;
+			case 'D':
+				pgdataDir = pg_strdup(optarg);
+				break;
 			case 'e':
 				if (sscanf(optarg, "%X/%X", &xlogid, &xrecoff) != 2)
 				{
@@ -912,6 +931,52 @@ main(int argc, char **argv)
 		}
 	}
 
+	if (!pgdataDir)
+		pgdataDir = getenv("PGDATA");
+
+	if (pgdataDir && cluster_passphrase)
+	{
+		uint8		wrapped_key[KMGR_WRAPPED_KEY_LEN];
+		uint32		cipherMethod = 0;
+
+		/* get wrapped key and cipher method from control file */
+		ret = kmgr_get_cipher_info(pgdataDir, &wrapped_key[0], &cipherMethod);
+		if (!ret)
+		{
+			pg_log_error("failed to get cipher info from control data file in cluster %s", pgdataDir);
+			goto bad_argument;
+		}
+
+		if (!cipherMethod)
+		{
+			pg_log_error("encryption is not enabled in cluster %s", pgdataDir);
+			goto bad_argument;
+		}
+
+		/*
+		 * Based on the passphrase, derive cipher contexts for encryption and
+		 * decryption which can be used to encrypt or decrypt WAL when TDE is
+		 * enabled.
+		 */
+		ret = kmgr_derive_cipher_ctx_from_passphrase_cmd(cluster_passphrase,
+														 &wrapped_key[0], &encrypt_cipher_ctx, &decrypt_cipher_ctx);
+		if (!ret)
+		{
+			pg_log_error("failed to derive cipher context for cluster encryption");
+			goto bad_argument;
+		}
+		else
+		{
+			pg_log_info("cipher context for cluster encryption derived successfully");
+		}
+
+		/* free the cipher contexts when not in used. */
+		if(encrypt_cipher_ctx)
+			pg_cipher_ctx_free(encrypt_cipher_ctx);
+		if(decrypt_cipher_ctx)
+			pg_cipher_ctx_free(decrypt_cipher_ctx);
+	}
+
 	if ((optind + 2) < argc)
 	{
 		pg_log_error("too many command-line arguments (first is \"%s\")",
diff --git a/src/bin/pg_waldump/t/002_tde.pl b/src/bin/pg_waldump/t/002_tde.pl
new file mode 100644
index 0000000000..aa93719e97
--- /dev/null
+++ b/src/bin/pg_waldump/t/002_tde.pl
@@ -0,0 +1,83 @@
+use strict;
+use warnings;
+use PostgresNode;
+use TestLib;
+use Test::More tests => 4;
+
+my $tempdir = TestLib::tempdir;
+my $tde_params_ok =
+  "echo 1234567890123456789012345678901234567890123456789012345678901234";
+my $tde_params_ko =
+  "echo 7777777766666666555555554444444433333333222222221111111100000000";
+
+my $node;
+my $startlsn;
+my $endlsn;
+my $xlogdir;
+my $datadir;
+
+sub init_and_start_tde_cluster
+{
+	$node = get_new_node('main');
+	$node->init(enable_kms => 1);
+
+	$xlogdir = $node->data_dir . "/pg_wal";
+	$datadir = $node->data_dir;
+
+	$node->start;
+	$node->psql('postgres', 'CREATE TABLE test_table(x integer)');
+	$node->psql('postgres',
+		'INSERT INTO test_table(x) SELECT y FROM generate_series(1, 10) a(y);'
+	);
+	$startlsn =
+	  $node->safe_psql('postgres', 'SELECT pg_current_wal_insert_lsn()');
+
+	$node->psql('postgres',
+		'INSERT INTO test_table(x) SELECT y FROM generate_series(1, 5000) a(y);'
+	);
+
+	$endlsn =
+	  $node->safe_psql('postgres', 'SELECT pg_current_wal_insert_lsn()');
+}
+
+init_and_start_tde_cluster();
+
+command_ok(
+	[
+		'pg_waldump', '-c', $tde_params_ok, '-p', $xlogdir, '-D',
+		$datadir,     '-s', $startlsn,      '-e', $endlsn,  '-n',
+		'100'
+	],
+	'pg_waldump with correct cluster passphrase');
+
+command_fails(
+	[
+		'pg_waldump', '-c', $tde_params_ko, '-p', $xlogdir, '-D',
+		$datadir,     '-s', $startlsn,      '-e', $endlsn,  '-n',
+		'100'
+	],
+	'pg_waldump with incorrect cluster passphrase');
+
+$node->append_conf('postgresql.conf',
+	qq[cluster_passphrase_command = '$tde_params_ko']);
+$node->reload;
+$node->psql('postgres', 'select pg_rotate_encryption_key();');
+
+command_ok(
+	[
+		'pg_waldump', '-c', $tde_params_ko, '-p', $xlogdir, '-D',
+		$datadir,     '-s', $startlsn,      '-e', $endlsn,  '-n',
+		'100'
+	],
+	'pg_waldump with correct cluster passphrase after key rotation');
+
+command_fails(
+	[
+		'pg_waldump', '-c', $tde_params_ok, '-p', $xlogdir, '-D',
+		$datadir,     '-s', $startlsn,      '-e', $endlsn,  '-n',
+		'100'
+	],
+	'pg_waldump with incorrect cluster passphrase after key rotation');
+
+$node->stop;
+$node->teardown_node;
diff --git a/src/common/kmgr_utils.c b/src/common/kmgr_utils.c
index e1936ab529..069259714e 100644
--- a/src/common/kmgr_utils.c
+++ b/src/common/kmgr_utils.c
@@ -22,6 +22,7 @@
 #endif
 #include "common/kmgr_utils.h"
 #include "common/sha2.h"
+#include "common/controldata_utils.h"
 #include "crypto/kmgr.h"
 #include "utils/elog.h"
 #include "storage/fd.h"
@@ -416,3 +417,145 @@ close_pipe_stream(FILE *file)
 #endif							/* WIN32 */
 }
 #endif							/* FRONTEND */
+
+/*
+ * Obtain cipher information from control data file. Return true and set
+ * *wrapped_key_out and *cipher_method_out accordingly. *wrapped_key_out must
+ * be pre-allocated to hold at least KMGR_WRAPPED_KEY_LEN bytes long. Return
+ * false if control data file fails to be accessed.
+ */
+bool
+kmgr_get_cipher_info(char *pgdata_dir,
+					 uint8 *wrapped_key_out,
+					 uint32 *cipher_method_out)
+{
+	ControlFileData *control_file = NULL;
+	bool		crc_ok = false;
+
+	Assert(pgdata_dir && wrapped_key_out && cipher_method_out);
+
+	control_file = get_controlfile(pgdata_dir, &crc_ok);
+	if (!control_file)
+	{
+#ifdef FRONTEND
+		pg_log_error("failed to get control_file from cluster %s", pgdata_dir);
+#else
+		ereport(ERROR,
+				(errcode(ERRCODE_INTERNAL_ERROR),
+				 errmsg("failed to get control_file from cluster %s", pgdata_dir)));
+#endif
+		return false;
+	}
+
+	if (!crc_ok)
+	{
+#ifdef FRONTEND
+		pg_log_error("control file failed CRC verification");
+#else
+		ereport(ERROR,
+				(errcode(ERRCODE_INTERNAL_ERROR),
+				 errmsg("control file failed CRC verification")));
+#endif
+		return false;
+	}
+
+	*cipher_method_out = control_file->key_management_enabled;
+	memcpy(wrapped_key_out, control_file->masterkey, KMGR_WRAPPED_KEY_LEN);
+	return true;
+}
+
+/*
+ * Evaluate the passphrase command and derive the raw master key if passphrase
+ * is correct. Return true if key derivation is successful and **encrypt_ctx_out
+ * and **decrypt_ctx_out will be allocated for subsequent encryption and decryption
+ * operations. Return false if passphrase is incorrect or key derivation is not
+ * successful. Caller is responsible for freeing **encrypt_ctx_out and **decrypt_ctx_out.
+ */
+bool
+kmgr_derive_cipher_ctx_from_passphrase_cmd(char *cluster_passphrase_command,
+										   uint8 *wrapped_key_in,
+										   pg_cipher_ctx * *encrypt_ctx_out,
+										   pg_cipher_ctx * *decrypt_ctx_out)
+{
+	int			len;
+	bool		ret = false;
+	char		passphrase[KMGR_MAX_PASSPHRASE_LEN];
+	uint8		rawkey[KMGR_KEY_AND_HMACKEY_LEN];
+
+	Assert(cluster_passphrase_command && wrapped_key_in);
+
+	/* derive passphrase from cluster_passphrase_cmd */
+	len = kmgr_run_cluster_passphrase_command(cluster_passphrase_command,
+											  &passphrase[0],
+											  KMGR_MAX_PASSPHRASE_LEN);
+
+	/* verify passphrase and obtain rawkey */
+	ret = kmgr_verify_passphrase(passphrase, len,
+								 wrapped_key_in,
+								 &rawkey[0]);
+
+	if (!ret)
+	{
+#ifdef FRONTEND
+		pg_log_error("encryption passphrase verification failed");
+#else
+		ereport(ERROR,
+				(errcode(ERRCODE_INTERNAL_ERROR),
+				 errmsg("encryption passphrase verification failed")));
+#endif
+		return false;
+	}
+
+	*encrypt_ctx_out = pg_cipher_ctx_create();
+	if (!(*encrypt_ctx_out))
+	{
+#ifdef FRONTEND
+		pg_log_error("failed to allocate cipher context for encryption");
+#else
+		ereport(ERROR,
+				(errcode(ERRCODE_INTERNAL_ERROR),
+				 errmsg("failed to allocate cipher context for encryption")));
+#endif
+		return false;
+	}
+
+	*decrypt_ctx_out = pg_cipher_ctx_create();
+	if (!(*decrypt_ctx_out))
+	{
+#ifdef FRONTEND
+		pg_log_error("failed to allocate cipher context for decryption");
+#else
+		ereport(ERROR,
+				(errcode(ERRCODE_INTERNAL_ERROR),
+				 errmsg("failed to allocate cipher context for decryption")));
+#endif
+		return false;
+	}
+
+	ret = pg_aes256_encrypt_init(*encrypt_ctx_out, &rawkey[0]);
+	if (!ret)
+	{
+#ifdef FRONTEND
+		pg_log_error("failed to initialize encryption cipher context with master key");
+#else
+		ereport(ERROR,
+				(errcode(ERRCODE_INTERNAL_ERROR),
+				 errmsg("failed to initialize encryption cipher context with master key")));
+#endif
+		return false;
+	}
+
+	ret = pg_aes256_decrypt_init(*decrypt_ctx_out, &rawkey[0]);
+	if (!ret)
+	{
+#ifdef FRONTEND
+		pg_log_error("failed to initialize decryption cipher context with master key");
+#else
+		ereport(ERROR,
+				(errcode(ERRCODE_INTERNAL_ERROR),
+				 errmsg("failed to initialize decryption cipher context with master key")));
+#endif
+		return false;
+	}
+	return true;
+}
diff --git a/src/include/common/kmgr_utils.h b/src/include/common/kmgr_utils.h
index 0d075bc2bb..3a152782a6 100644
--- a/src/include/common/kmgr_utils.h
+++ b/src/include/common/kmgr_utils.h
@@ -80,5 +80,12 @@ extern bool kmgr_compute_HMAC(KeyWrapCtx *ctx, const uint8 *in, int inlen,
 							  uint8 *out);
 extern int	kmgr_run_cluster_passphrase_command(char *passphrase_command,
 												char *buf, int size);
+extern bool kmgr_get_cipher_info(char *pgdata_dir,
+							     uint8 *wrapped_key_out,
+							     uint32 *cipher_method_out);
+extern bool kmgr_derive_cipher_ctx_from_passphrase_cmd(char *cluster_passphrase_command,
+													   uint8 *wrapped_key_in,
+													   pg_cipher_ctx * *encrypt_ctx_out,
+													   pg_cipher_ctx * *decrypt_ctx_out);
 
 #endif							/* KMGR_UTILS_H */