Add a interface to alter key provider of the principal key

pg_tde--1.0.sql

@@ -194,6 +194,11 @@ RETURNS boolean
 AS 'MODULE_PATHNAME'
 LANGUAGE C;
 
+CREATE FUNCTION pg_tde_alter_principal_key_keyring(new_provider_name VARCHAR(255))
+RETURNS boolean
+AS 'MODULE_PATHNAME'
+LANGUAGE C;
+
 CREATE FUNCTION pg_tde_extension_initialize()
 RETURNS VOID
 AS 'MODULE_PATHNAME'
@@ -342,6 +347,7 @@ BEGIN
     PERFORM pg_tde_grant_execute_privilege_on_function(target_user_or_role, 'pg_tde_add_key_provider_vault_v2', 'varchar, JSON, JSON,JSON,JSON');
 
     PERFORM pg_tde_grant_execute_privilege_on_function(target_user_or_role, 'pg_tde_set_principal_key', 'varchar, varchar, BOOLEAN');
+    PERFORM pg_tde_grant_execute_privilege_on_function(target_user_or_role, 'pg_tde_alter_principal_key_keyring', 'varchar');
 
     PERFORM pg_tde_grant_execute_privilege_on_function(target_user_or_role, 'pg_tde_rotate_principal_key', 'pg_tde_global, varchar, varchar');
     PERFORM pg_tde_grant_execute_privilege_on_function(target_user_or_role, 'pg_tde_rotate_principal_key', 'varchar, varchar');
@@ -410,6 +416,7 @@ BEGIN
     PERFORM pg_tde_revoke_execute_privilege_on_function(target_user_or_role, 'pg_tde_add_key_provider_vault_v2', 'varchar, JSON, JSON,JSON,JSON');
 
     PERFORM pg_tde_revoke_execute_privilege_on_function(target_user_or_role, 'pg_tde_set_principal_key', 'varchar, varchar, BOOLEAN');
+    PERFORM pg_tde_revoke_execute_privilege_on_function(target_user_or_role, 'pg_tde_alter_principal_key_keyring', 'varchar');
 
     PERFORM pg_tde_revoke_execute_privilege_on_function(target_user_or_role, 'pg_tde_rotate_principal_key', 'pg_tde_global, varchar, varchar');
     PERFORM pg_tde_revoke_execute_privilege_on_function(target_user_or_role, 'pg_tde_rotate_principal_key', 'varchar, varchar');

src/access/pg_tde_tdemap.c

@@ -118,7 +118,7 @@ static int pg_tde_open_file_basic(char *tde_filename, int fileFlags, bool ignore
 static int pg_tde_file_header_read(char *tde_filename, int fd, TDEFileHeader *fheader, bool *is_new_file, off_t *bytes_read);
 static bool pg_tde_read_one_map_entry(int fd, const RelFileLocator *rlocator, int flags, TDEMapEntry *map_entry, off_t *offset);
 static RelKeyData *pg_tde_read_one_keydata(int keydata_fd, int32 key_index, TDEPrincipalKey *principal_key);
-static int pg_tde_open_file(char *tde_filename, TDEPrincipalKeyInfo *principal_key_info, bool should_fill_info, int fileFlags, bool *is_new_file, off_t *offset);
+static int pg_tde_open_file(char *tde_filename, TDEPrincipalKeyInfo *principal_key_info, bool update_header, int fileFlags, bool *is_new_file, off_t *curr_pos);
 static RelKeyData *pg_tde_get_key_from_cache(RelFileNumber rel_number, uint32 key_type);
 
 #ifndef FRONTEND
@@ -282,7 +282,7 @@ pg_tde_delete_tde_files(Oid dbOid, Oid spcOid)
  * The caller must have an EXCLUSIVE LOCK on the files before calling this function.
  */
 bool
-pg_tde_save_principal_key(TDEPrincipalKeyInfo *principal_key_info)
+pg_tde_save_principal_key(TDEPrincipalKeyInfo *principal_key_info, bool truncate_existing, bool update_header)
 {
 	int map_fd = -1;
 	int keydata_fd = -1;
@@ -291,17 +291,24 @@ pg_tde_save_principal_key(TDEPrincipalKeyInfo *principal_key_info)
 	bool is_new_key_data = false;
 	char db_map_path[MAXPGPATH] = {0};
 	char db_keydata_path[MAXPGPATH] = {0};
+	int file_flags = O_RDWR | O_CREAT;
 
 	/* Set the file paths */
 	pg_tde_set_db_file_paths(principal_key_info->databaseId,
 							 principal_key_info->tablespaceId,
 							 db_map_path, db_keydata_path);
 
-	ereport(LOG, (errmsg("pg_tde_save_principal_key")));
+	ereport(DEBUG2,
+		(errmsg("pg_tde_save_principal_key"),
+			errdetail("truncate_existing:%s update_header:%s", truncate_existing?"YES":"NO", update_header?"YES":"NO")));
+	/*
+	 * Create or truncate these map and keydata files.
+	 */
+	if (truncate_existing)
+		file_flags |= O_TRUNC;
 
-	/* Create or truncate these map and keydata files. */
-	map_fd = pg_tde_open_file(db_map_path, principal_key_info, false, O_RDWR | O_CREAT | O_TRUNC, &is_new_map, &curr_pos);
-	keydata_fd = pg_tde_open_file(db_keydata_path, principal_key_info, false, O_RDWR | O_CREAT | O_TRUNC, &is_new_key_data, &curr_pos);
+	map_fd = pg_tde_open_file(db_map_path, principal_key_info, update_header, file_flags, &is_new_map, &curr_pos);
+	keydata_fd = pg_tde_open_file(db_keydata_path, principal_key_info, update_header, file_flags, &is_new_key_data, &curr_pos);
 
 	/* Closing files. */
 	close(map_fd);
@@ -345,6 +352,8 @@ pg_tde_file_header_write(char *tde_filename, int fd, TDEPrincipalKeyInfo *princi
 				(errcode_for_file_access(),
 				 errmsg("could not fsync file \"%s\": %m", tde_filename)));
 	}
+	ereport(DEBUG2,
+			(errmsg("Wrote the header to %s", tde_filename)));
 
 	return fd;
 }
@@ -1176,7 +1185,7 @@ tde_decrypt_rel_key(TDEPrincipalKey *principal_key, RelKeyData *enc_rel_key_data
  * or an error is thrown if the file does not exist.
  */
 static int
-pg_tde_open_file(char *tde_filename, TDEPrincipalKeyInfo *principal_key_info, bool should_fill_info, int fileFlags, bool *is_new_file, off_t *curr_pos)
+pg_tde_open_file(char *tde_filename, TDEPrincipalKeyInfo *principal_key_info, bool update_header, int fileFlags, bool *is_new_file, off_t *curr_pos)
 {
 	int fd = -1;
 	TDEFileHeader fheader;
@@ -1193,7 +1202,7 @@ pg_tde_open_file(char *tde_filename, TDEPrincipalKeyInfo *principal_key_info, bo
 
 #ifndef FRONTEND
 	/* In case it's a new file, let's add the header now. */
-	if (*is_new_file && principal_key_info)
+	if ((*is_new_file || update_header) && principal_key_info)
 		pg_tde_file_header_write(tde_filename, fd, principal_key_info, &bytes_written);
 #endif							/* FRONTEND */
 

src/access/pg_tde_xlog.c

@@ -47,12 +47,16 @@ tdeheap_rmgr_redo(XLogReaderState *record)
 		pg_tde_write_key_map_entry(&xlrec->rlocator, &xlrec->relKey, pk);
 		LWLockRelease(tde_lwlock_enc_keys());
 	}
-	else if (info == XLOG_TDE_ADD_PRINCIPAL_KEY)
+	else if (info == XLOG_TDE_ADD_PRINCIPAL_KEY || info == XLOG_TDE_UPDATE_PRINCIPAL_KEY)
 	{
 		TDEPrincipalKeyInfo *mkey = (TDEPrincipalKeyInfo *) XLogRecGetData(record);
 
 		LWLockAcquire(tde_lwlock_enc_keys(), LW_EXCLUSIVE);
-		save_principal_key_info(mkey);
+		if (info == XLOG_TDE_ADD_PRINCIPAL_KEY)
+			save_principal_key_info(mkey);
+		else
+			update_principal_key_info(mkey);
+
 		LWLockRelease(tde_lwlock_enc_keys());
 	}
 	else if (info == XLOG_TDE_EXTENSION_INSTALL_KEY)
@@ -110,6 +114,12 @@ tdeheap_rmgr_desc(StringInfo buf, XLogReaderState *record)
 
 		appendStringInfo(buf, "add tde principal key for db %u/%u", xlrec->databaseId, xlrec->tablespaceId);
 	}
+	if (info == XLOG_TDE_UPDATE_PRINCIPAL_KEY)
+	{
+		TDEPrincipalKeyInfo *xlrec = (TDEPrincipalKeyInfo *) XLogRecGetData(record);
+
+		appendStringInfo(buf, "Alter key provider to:%d for tde principal key for db %u/%u", xlrec->keyringId, xlrec->databaseId, xlrec->tablespaceId);
+	}
 	if (info == XLOG_TDE_EXTENSION_INSTALL_KEY)
 	{
 		XLogExtensionInstall *xlrec = (XLogExtensionInstall *) XLogRecGetData(record);
@@ -139,6 +149,9 @@ tdeheap_rmgr_identify(uint8 info)
 	if ((info & ~XLR_INFO_MASK) == XLOG_TDE_ADD_PRINCIPAL_KEY)
 		return "XLOG_TDE_ADD_PRINCIPAL_KEY";
 
+	if ((info & ~XLR_INFO_MASK) == XLOG_TDE_UPDATE_PRINCIPAL_KEY)
+		return "XLOG_TDE_UPDATE_PRINCIPAL_KEY";
+
 	if ((info & ~XLR_INFO_MASK) == XLOG_TDE_EXTENSION_INSTALL_KEY)
 		return "XLOG_TDE_EXTENSION_INSTALL_KEY";
 

src/catalog/tde_principal_key.c

@@ -86,6 +86,7 @@ static TDEPrincipalKey *set_principal_key_with_keyring(const char *key_name,
 													   GenericKeyring *keyring,
 													   Oid dbOid, Oid spcOid,
 													   bool ensure_new_key);
+static TDEPrincipalKey *alter_keyprovider_for_principal_key(GenericKeyring *newKeyring,Oid dbOid, Oid spcOid);
 
 static const TDEShmemSetupRoutine principal_key_info_shmem_routine = {
 	.init_shared_state = initialize_shared_state,
@@ -213,7 +214,14 @@ save_principal_key_info(TDEPrincipalKeyInfo *principal_key_info)
 {
 	Assert(principal_key_info != NULL);
 
-	return pg_tde_save_principal_key(principal_key_info);
+	return pg_tde_save_principal_key(principal_key_info, true, true);
+}
+
+bool
+update_principal_key_info(TDEPrincipalKeyInfo *principal_key_info)
+{
+	Assert(principal_key_info != NULL);
+	return pg_tde_save_principal_key(principal_key_info, false, true);
 }
 
 /*
@@ -297,6 +305,61 @@ set_principal_key_with_keyring(const char *key_name, GenericKeyring *keyring,
 	return principalKey;
 }
 
+/*
+ * alter_keyprovider_for_principal_key:
+ */
+TDEPrincipalKey *
+alter_keyprovider_for_principal_key(GenericKeyring *newKeyring,
+                                    Oid dbOid, Oid spcOid)
+{
+    TDEPrincipalKeyInfo *principalKeyInfo = NULL;
+	TDEPrincipalKey *principal_key = NULL;
+
+	LWLock *lock_files = tde_lwlock_enc_keys();
+
+	Assert(newKeyring != NULL);
+	LWLockAcquire(lock_files, LW_EXCLUSIVE);
+
+	principalKeyInfo = pg_tde_get_principal_key_info(dbOid, spcOid);
+
+	if (principalKeyInfo == NULL)
+	{
+        LWLockRelease(lock_files);
+        ereport(ERROR,
+			(errmsg("Principal key not set for the database"),
+				errhint("Use set_principal_key interface to set the principal key")));
+	}
+
+	if (newKeyring->key_id == principalKeyInfo->keyringId)
+	{
+        LWLockRelease(lock_files);
+        ereport(ERROR,
+                (errmsg("New key provider is same as the current key provider")));
+    }
+    /* update the key provider in principal key info */
+
+	ereport(DEBUG2,
+			(errmsg("Changing keyprovider ID from :%d to %d", principalKeyInfo->keyringId, newKeyring->key_id)));
+
+	principalKeyInfo->keyringId = newKeyring->key_id;
+
+	update_principal_key_info(principalKeyInfo);
+
+	/* XLog the new key*/
+	XLogBeginInsert();
+	XLogRegisterData((char *)principalKeyInfo, sizeof(TDEPrincipalKeyInfo));
+	XLogInsert(RM_TDERMGR_ID, XLOG_TDE_UPDATE_PRINCIPAL_KEY);
+
+	/* clear the cache as well */
+	clear_principal_key_cache(dbOid);
+
+	principal_key = GetPrincipalKey(dbOid, spcOid, LW_EXCLUSIVE);
+
+	LWLockRelease(lock_files);
+
+	return principal_key;
+}
+
 bool
 SetPrincipalKey(const char *key_name, const char *provider_name, bool ensure_new_key)
 {
@@ -308,6 +371,15 @@ SetPrincipalKey(const char *key_name, const char *provider_name, bool ensure_new
 	return (principal_key != NULL);
 }
 
+bool
+AlterPrincipalKeyKeyring(const char *provider_name)
+{
+    TDEPrincipalKey *principal_key = alter_keyprovider_for_principal_key(GetKeyProviderByName(provider_name, MyDatabaseId, MyDatabaseTableSpace),
+                                                                    MyDatabaseId, MyDatabaseTableSpace);
+
+    return (principal_key != NULL);
+}
+
 bool
 RotatePrincipalKey(TDEPrincipalKey *current_key, const char *new_key_name, const char *new_provider_name, bool ensure_new_key)
 {
@@ -632,6 +704,19 @@ pg_tde_set_principal_key(PG_FUNCTION_ARGS)
 	PG_RETURN_BOOL(ret);
 }
 
+PG_FUNCTION_INFO_V1(pg_tde_alter_principal_key_keyring);
+Datum pg_tde_alter_principal_key_keyring(PG_FUNCTION_ARGS);
+
+Datum pg_tde_alter_principal_key_keyring(PG_FUNCTION_ARGS)
+{
+	char *provider_name = text_to_cstring(PG_GETARG_TEXT_PP(0));
+	bool ret;
+
+	ereport(LOG, (errmsg("Altering principal key provider to \"%s\" for the database", provider_name)));
+	ret = AlterPrincipalKeyKeyring(provider_name);
+	PG_RETURN_BOOL(ret);
+}
+
 /*
  * SQL interface for key rotation
  */

src/include/access/pg_tde_tdemap.h

@@ -65,7 +65,7 @@ extern RelKeyData *GetTdeGlobaleRelationKey(RelFileLocator rel);
 extern void pg_tde_delete_tde_files(Oid dbOid, Oid spcOid);
 
 extern TDEPrincipalKeyInfo *pg_tde_get_principal_key_info(Oid dbOid, Oid spcOid);
-extern bool pg_tde_save_principal_key(TDEPrincipalKeyInfo *principal_key_info);
+extern bool pg_tde_save_principal_key(TDEPrincipalKeyInfo *principal_key_info, bool truncate_existing, bool update_header);
 extern bool pg_tde_perform_rotate_key(TDEPrincipalKey *principal_key, TDEPrincipalKey *new_principal_key);
 extern bool pg_tde_write_map_keydata_files(off_t map_size, char *m_file_data, off_t keydata_size, char *k_file_data);
 extern RelKeyData *tde_create_rel_key(RelFileNumber rel_num, InternalKey *key, TDEPrincipalKeyInfo *principal_key_info);

src/include/access/pg_tde_xlog.h

@@ -22,6 +22,7 @@
 #define XLOG_TDE_ROTATE_KEY				0x30
 #define XLOG_TDE_ADD_KEY_PROVIDER_KEY 	0x40
 #define XLOG_TDE_FREE_MAP_ENTRY		 	0x50
+#define XLOG_TDE_UPDATE_PRINCIPAL_KEY	0x60
 
 /* TODO: ID has to be registedred and changed: https://wiki.postgresql.org/wiki/CustomWALResourceManagers */
 #define RM_TDERMGR_ID	RM_EXPERIMENTAL_ID

src/include/catalog/tde_principal_key.h

@@ -68,9 +68,11 @@ extern TDEPrincipalKey *GetPrincipalKey(Oid dbOid, Oid spcOid, void *lockMode);
 #endif
 
 extern bool save_principal_key_info(TDEPrincipalKeyInfo *principalKeyInfo);
+extern bool update_principal_key_info(TDEPrincipalKeyInfo *principal_key_info);
 
 extern Oid	GetPrincipalKeyProviderId(void);
 extern bool SetPrincipalKey(const char *key_name, const char *provider_name, bool ensure_new_key);
+extern bool AlterPrincipalKeyKeyring(const char *provider_name);
 extern bool RotatePrincipalKey(TDEPrincipalKey *current_key, const char *new_key_name, const char *new_provider_name, bool ensure_new_key);
 extern bool xl_tde_perform_rotate_key(XLogPrincipalKeyRotate *xlrec);