From e7dd925d8272f259d4e3946847612d969cca42b7 Mon Sep 17 00:00:00 2001
From: Bharath Rupireddy <bharath.rupireddy@enterprisedb.com>
Date: Mon, 17 May 2021 10:26:10 +0530
Subject: [PATCH v1] Provide TDE nonce size as an initdb option

Add an initdb option --tde-nonce-size with supported range from
0 to PG_UINT8_MAX in bytes. 0 being default value disables the
feature. If the size is specified in bytes, server reserves space
on each page of table, index etc. that's getting newly added.

See further patches for the main feature implementation.
---
 src/backend/access/transam/xlog.c | 15 +++++++++++++
 src/backend/bootstrap/bootstrap.c | 15 ++++++++++++-
 src/backend/utils/init/globals.c  |  2 ++
 src/backend/utils/misc/guc.c      | 13 +++++++++++
 src/bin/initdb/initdb.c           | 37 ++++++++++++++++++++++++++++++-
 src/include/catalog/pg_control.h  |  2 ++
 src/include/miscadmin.h           |  2 ++
 7 files changed, 84 insertions(+), 2 deletions(-)

diff --git a/src/backend/access/transam/xlog.c b/src/backend/access/transam/xlog.c
index 8d163f190f..218aad275b 100644
--- a/src/backend/access/transam/xlog.c
+++ b/src/backend/access/transam/xlog.c
@@ -4680,6 +4680,7 @@ WriteControlFile(void)
 	ControlFile->relseg_size = RELSEG_SIZE;
 	ControlFile->xlog_blcksz = XLOG_BLCKSZ;
 	ControlFile->xlog_seg_size = wal_segment_size;
+	ControlFile->tde_nonce_size = tde_nonce_size;
 
 	ControlFile->nameDataLen = NAMEDATALEN;
 	ControlFile->indexMaxKeys = INDEX_MAX_KEYS;
@@ -4749,6 +4750,7 @@ ReadControlFile(void)
 	pg_crc32c	crc;
 	int			fd;
 	static char wal_segsz_str[20];
+	static char tde_nonce_size_str[5];
 	int			r;
 
 	/*
@@ -4781,6 +4783,8 @@ ReadControlFile(void)
 
 	close(fd);
 
+	tde_nonce_size = ControlFile->tde_nonce_size;
+
 	/*
 	 * Check for expected pg_control format version.  If this is wrong, the
 	 * CRC check will likely fail because we'll be checking the wrong number
@@ -4928,6 +4932,17 @@ ReadControlFile(void)
 		ereport(ERROR, (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
 						errmsg("\"max_wal_size\" must be at least twice \"wal_segment_size\"")));
 
+	if (tde_nonce_size > PG_UINT8_MAX || tde_nonce_size < 0)
+		ereport(ERROR, (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+				errmsg("TDE nonce size %d is out of range",
+						ControlFile->tde_nonce_size),
+				errdetail("TDE nonce size must be between 0 and %d.", (int) PG_UINT8_MAX),
+				errhint("It looks like you need to initdb.")));
+
+	snprintf(tde_nonce_size_str, sizeof(tde_nonce_size_str), "%d", tde_nonce_size);
+	SetConfigOption("tde_nonce_size", tde_nonce_size_str, PGC_INTERNAL,
+					PGC_S_OVERRIDE);
+
 	UsableBytesInSegment =
 		(wal_segment_size / XLOG_BLCKSZ * UsableBytesInPage) -
 		(SizeOfXLogLongPHD - SizeOfXLogShortPHD);
diff --git a/src/backend/bootstrap/bootstrap.c b/src/backend/bootstrap/bootstrap.c
index b155237488..f74387017b 100644
--- a/src/backend/bootstrap/bootstrap.c
+++ b/src/backend/bootstrap/bootstrap.c
@@ -225,7 +225,7 @@ AuxiliaryProcessMain(int argc, char *argv[])
 	/* If no -x argument, we are a CheckerProcess */
 	MyAuxProcType = CheckerProcess;
 
-	while ((flag = getopt(argc, argv, "B:c:d:D:Fkr:x:X:-:")) != -1)
+	while ((flag = getopt(argc, argv, "B:c:d:D:Fkr:T:x:X:-:")) != -1)
 	{
 		switch (flag)
 		{
@@ -299,6 +299,19 @@ AuxiliaryProcessMain(int argc, char *argv[])
 						free(value);
 					break;
 				}
+			case 'T':
+				{
+					unsigned long val = strtoul(optarg, NULL, 0);
+
+					if (val > PG_UINT8_MAX || val < 0)
+						ereport(ERROR,
+								(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+								 errmsg("TDE nonce option -T requires a value between 0 and %d", PG_UINT8_MAX)));
+
+					SetConfigOption("tde_nonce_size", optarg, PGC_INTERNAL,
+									PGC_S_OVERRIDE);
+				}
+				break;
 			default:
 				write_stderr("Try \"%s --help\" for more information.\n",
 							 progname);
diff --git a/src/backend/utils/init/globals.c b/src/backend/utils/init/globals.c
index 381d9e548d..91763c95ad 100644
--- a/src/backend/utils/init/globals.c
+++ b/src/backend/utils/init/globals.c
@@ -150,3 +150,5 @@ int64		VacuumPageDirty = 0;
 
 int			VacuumCostBalance = 0;	/* working state for vacuum */
 bool		VacuumCostActive = false;
+
+int         tde_nonce_size = 0;
diff --git a/src/backend/utils/misc/guc.c b/src/backend/utils/misc/guc.c
index ee731044b6..118191d976 100644
--- a/src/backend/utils/misc/guc.c
+++ b/src/backend/utils/misc/guc.c
@@ -3221,6 +3221,19 @@ static struct config_int ConfigureNamesInt[] =
 		NULL, NULL, NULL
 	},
 
+	{
+		{"tde_nonce_size", PGC_INTERNAL, PRESET_OPTIONS,
+			gettext_noop("Shows the size of TDE nonce."),
+			NULL,
+			GUC_NOT_IN_SAMPLE | GUC_DISALLOW_IN_FILE
+		},
+		&tde_nonce_size,
+		0,
+		0,
+		PG_UINT8_MAX,
+		NULL, NULL, NULL
+	},
+
 	{
 		{"autovacuum_naptime", PGC_SIGHUP, AUTOVACUUM,
 			gettext_noop("Time to sleep between autovacuum runs."),
diff --git a/src/bin/initdb/initdb.c b/src/bin/initdb/initdb.c
index 152d21e88b..1af215abe3 100644
--- a/src/bin/initdb/initdb.c
+++ b/src/bin/initdb/initdb.c
@@ -147,6 +147,8 @@ static bool data_checksums = false;
 static char *xlog_dir = NULL;
 static char *str_wal_segment_size_mb = NULL;
 static int	wal_segment_size_mb;
+static char *str_tde_nonce_size_b = NULL;
+static unsigned char tde_nonce_size_b;
 
 
 /* internal vars */
@@ -1403,9 +1405,10 @@ bootstrap_template1(void)
 	unsetenv("PGCLIENTENCODING");
 
 	snprintf(cmd, sizeof(cmd),
-			 "\"%s\" --boot -x1 -X %u %s %s %s",
+			 "\"%s\" --boot -x1 -X %u -T %d %s %s %s",
 			 backend_exec,
 			 wal_segment_size_mb * (1024 * 1024),
+			 tde_nonce_size_b,
 			 data_checksums ? "-k" : "",
 			 boot_options,
 			 debug ? "-d 5" : "");
@@ -2262,6 +2265,7 @@ usage(const char *progname)
 	printf(_("  -W, --pwprompt            prompt for a password for the new superuser\n"));
 	printf(_("  -X, --waldir=WALDIR       location for the write-ahead log directory\n"));
 	printf(_("      --wal-segsize=SIZE    size of WAL segments, in megabytes\n"));
+	printf(_("      --tde-nonce-size=SIZE    size of TDE nonce in bytes\n"));
 	printf(_("\nLess commonly used options:\n"));
 	printf(_("  -d, --debug               generate lots of debugging output\n"));
 	printf(_("  -L DIRECTORY              where to find the input files\n"));
@@ -2943,6 +2947,7 @@ main(int argc, char *argv[])
 		{"wal-segsize", required_argument, NULL, 12},
 		{"data-checksums", no_argument, NULL, 'k'},
 		{"allow-group-access", no_argument, NULL, 'g'},
+		{"tde-nonce-size", required_argument, NULL, 14},
 		{NULL, 0, NULL, 0}
 	};
 
@@ -3084,6 +3089,9 @@ main(int argc, char *argv[])
 			case 'g':
 				SetDataDirectoryCreatePerm(PG_DIR_MODE_GROUP);
 				break;
+			case 14:
+				str_tde_nonce_size_b =  pg_strdup(optarg);;
+				break;
 			default:
 				/* getopt_long already emitted a complaint */
 				fprintf(stderr, _("Try \"%s --help\" for more information.\n"),
@@ -3170,6 +3178,33 @@ main(int argc, char *argv[])
 		}
 	}
 
+	/* set TDE nonce size */
+	if (str_tde_nonce_size_b == NULL)
+		tde_nonce_size_b = 0;
+	else
+	{
+		char	   *endptr;
+		long		val;
+
+		/* check that the argument is a number */
+		val = strtol(str_tde_nonce_size_b, &endptr, 10);
+
+		/* verify that TDE nonce size is valid */
+		if (endptr == str_tde_nonce_size_b || *endptr != '\0')
+		{
+			pg_log_error("argument of --tde-nonce-size must be a number");
+			exit(1);
+		}
+
+		if (val > PG_UINT8_MAX || val < 0)
+		{
+			pg_log_error("argument of --tde-nonce-size must be between 0 and %d", PG_UINT8_MAX);
+			exit(1);
+		}
+
+		tde_nonce_size_b = val;
+	}
+
 	get_restricted_token();
 
 	setup_pgdata();
diff --git a/src/include/catalog/pg_control.h b/src/include/catalog/pg_control.h
index e3f48158ce..2e34bdc6ab 100644
--- a/src/include/catalog/pg_control.h
+++ b/src/include/catalog/pg_control.h
@@ -226,6 +226,8 @@ typedef struct ControlFileData
 	 */
 	char		mock_authentication_nonce[MOCK_AUTH_NONCE_LEN];
 
+	uint8		tde_nonce_size;		/* size of TDE nonce */
+
 	/* CRC of all above ... MUST BE LAST! */
 	pg_crc32c	crc;
 } ControlFileData;
diff --git a/src/include/miscadmin.h b/src/include/miscadmin.h
index 4dc343cbc5..dc9bee3e0c 100644
--- a/src/include/miscadmin.h
+++ b/src/include/miscadmin.h
@@ -275,6 +275,8 @@ extern int64 VacuumPageDirty;
 extern int	VacuumCostBalance;
 extern bool VacuumCostActive;
 
+extern int tde_nonce_size;
+
 
 /* in tcop/postgres.c */
 
-- 
2.25.1

