Add key type to internal relation keys in pg_tde: (#317)

This commit introduces a key type to internal relation keys, allowing segregation of SMGR and heap_basic internal keys. It resolves the issue of double encryption when both tde_heap and tde_heap_basic access methods are in use. With this enhancement, pg_tde now supports three types of internal keys: SMGR keys, heap_basic keys, and global keys, improving flexibility and preventing redundant encryption operations.
percona · Oct 21, 2024 · 5ca8ede · 5ca8ede
1 parent 0995806
commit 5ca8ede
Show file tree

Hide file tree

Showing 19 changed files with 151 additions and 97 deletions.
diff --git a/src/access/pg_tde_ddl.c b/src/access/pg_tde_ddl.c
@@ -29,7 +29,7 @@ void SetupTdeDDLHooks(void)
 }
 
 static void
- tdeheap_object_access_hook(ObjectAccessType access, Oid classId, Oid objectId,
+tdeheap_object_access_hook(ObjectAccessType access, Oid classId, Oid objectId,
                              int subId, void *arg)
  {
     Relation    rel = NULL;
@@ -49,7 +49,7 @@ static void
             rel->rd_rel->relkind == RELKIND_MATVIEW) &&
             (subId == 0) && is_tdeheap_rel(rel))
             {
-                pg_tde_delete_key_map_entry(&rel->rd_locator);
+                pg_tde_delete_key_map_entry(&rel->rd_locator, MAP_ENTRY_VALID);
             }
         relation_close(rel, AccessShareLock);
     }

diff --git a/src/access/pg_tde_slot.c b/src/access/pg_tde_slot.c
@@ -579,6 +579,6 @@ get_current_slot_relation_key(TDEBufferHeapTupleTableSlot *bslot, Relation rel)
 {
 	Assert(bslot != NULL);
 	if (bslot->cached_relation_key == NULL)
-		bslot->cached_relation_key = GetRelationKey(rel->rd_locator);
+		bslot->cached_relation_key = GetHeapBaiscRelationKey(rel->rd_locator);
 	return bslot->cached_relation_key;
 }
diff --git a/src/access/pg_tde_tdemap.c b/src/access/pg_tde_tdemap.c
diff --git a/src/access/pg_tde_xlog.c b/src/access/pg_tde_xlog.c
@@ -40,7 +40,7 @@ tdeheap_rmgr_redo(XLogReaderState *record)
 		XLogRelKey *xlrec = (XLogRelKey *) XLogRecGetData(record);
 
 		LWLockAcquire(tde_lwlock_enc_keys(), LW_EXCLUSIVE);
-		pg_tde_write_key_map_entry(&xlrec->rlocator, &xlrec->relKey, NULL);
+		pg_tde_write_key_map_entry(&xlrec->rlocator, xlrec->entry_type, &xlrec->relKey, NULL);
 		LWLockRelease(tde_lwlock_enc_keys());
 	}
 	else if (info == XLOG_TDE_ADD_PRINCIPAL_KEY)

diff --git a/src/access/pg_tde_xlog_encrypt.c b/src/access/pg_tde_xlog_encrypt.c
@@ -124,7 +124,7 @@ TDEXLogWriteEncryptedPages(int fd, const void *buf, size_t count, off_t offset)
 	size_t	data_size = 0;
 	XLogPageHeader	curr_page_hdr = &EncryptCurrentPageHrd;
 	XLogPageHeader	enc_buf_page;
-	RelKeyData		*key = GetRelationKey(GLOBAL_SPACE_RLOCATOR(XLOG_TDE_OID));
+	RelKeyData *key = GetTdeGlobaleRelationKey(GLOBAL_SPACE_RLOCATOR(XLOG_TDE_OID));
 	off_t	enc_off;
 	size_t	page_size = XLOG_BLCKSZ - offset % XLOG_BLCKSZ;
 	uint32	iv_ctr = 0;
@@ -229,7 +229,7 @@ tdeheap_xlog_seg_read(int fd, void *buf, size_t count, off_t offset)
 	char	iv_prefix[16] = {0,};
 	size_t	data_size = 0;
 	XLogPageHeader	curr_page_hdr = &DecryptCurrentPageHrd;
-	RelKeyData		*key = GetRelationKey(GLOBAL_SPACE_RLOCATOR(XLOG_TDE_OID));
+	RelKeyData *key = GetTdeGlobaleRelationKey(GLOBAL_SPACE_RLOCATOR(XLOG_TDE_OID));
 	size_t	page_size = XLOG_BLCKSZ - offset % XLOG_BLCKSZ;
 	off_t	dec_off;
 	uint32	iv_ctr = 0;

diff --git a/src/catalog/tde_global_space.c b/src/catalog/tde_global_space.c
@@ -67,7 +67,7 @@ TDEInitGlobalKeys(const char *dir)
 		if (dir != NULL)
 			pg_tde_set_globalspace_dir(dir);
 
-		ikey = pg_tde_get_key_from_file(&GLOBAL_SPACE_RLOCATOR(XLOG_TDE_OID));
+		ikey = pg_tde_get_key_from_file(&GLOBAL_SPACE_RLOCATOR(XLOG_TDE_OID), TDE_KEY_TYPE_GLOBAL);
 
 		/*
 		 * Internal Key should be in the TopMemmoryContext because of SSL
@@ -78,7 +78,7 @@ TDEInitGlobalKeys(const char *dir)
 		 * backend. (see
 		 * https://github.com/percona-Lab/pg_tde/pull/214#discussion_r1648998317)
 		 */
-		pg_tde_put_key_into_cache(XLOG_TDE_OID, ikey);
+		pg_tde_put_key_into_cache(XLOG_TDE_OID, TDE_KEY_TYPE_GLOBAL, ikey);
 	}
 }
 
@@ -156,9 +156,9 @@ init_keys(void)
 	}
 
 	rlocator = &GLOBAL_SPACE_RLOCATOR(XLOG_TDE_OID);
-	rel_key_data = tde_create_rel_key(rlocator->relNumber, &int_key, &mkey->keyInfo);
+	rel_key_data = tde_create_rel_key(rlocator->relNumber, TDE_KEY_TYPE_GLOBAL, &int_key, &mkey->keyInfo);
 	enc_rel_key_data = tde_encrypt_rel_key(mkey, rel_key_data, rlocator);
-	pg_tde_write_key_map_entry(rlocator, enc_rel_key_data, &mkey->keyInfo);
+	pg_tde_write_key_map_entry(rlocator, TDE_KEY_TYPE_GLOBAL, enc_rel_key_data, &mkey->keyInfo);
 	pfree(enc_rel_key_data);
 	pfree(mkey);
 }

diff --git a/src/encryption/enc_tde.c b/src/encryption/enc_tde.c
@@ -214,7 +214,7 @@ PGTdePageAddItemExtended(RelFileLocator rel,
 	uint32	data_len = size - header_size;
 	/* ctid stored in item is incorrect (not set) at this point */
 	ItemPointerData ip;
-	RelKeyData *key = GetRelationKey(rel);
+	RelKeyData *key = GetHeapBaiscRelationKey(rel);
 
 	ItemPointerSet(&ip, bn, off); 
 

diff --git a/src/include/access/pg_tde_tdemap.h b/src/include/access/pg_tde_tdemap.h
@@ -14,6 +14,13 @@
 #include "catalog/tde_principal_key.h"
 #include "storage/relfilelocator.h"
 
+/* Map entry flags */
+#define MAP_ENTRY_EMPTY         0x00
+#define TDE_KEY_TYPE_HEAP_BASIC 0x01
+#define TDE_KEY_TYPE_SMGR       0x02
+#define TDE_KEY_TYPE_GLOBAL 0x04
+#define MAP_ENTRY_VALID (TDE_KEY_TYPE_HEAP_BASIC | TDE_KEY_TYPE_SMGR | TDE_KEY_TYPE_GLOBAL)
+
 typedef struct InternalKey
 {
     uint8   key[INTERNAL_KEY_LEN];
@@ -30,31 +37,38 @@ typedef struct RelKeyData
 typedef struct XLogRelKey
 {
 	RelFileLocator  rlocator;
+    uint32          entry_type;
 	RelKeyData      relKey;
 } XLogRelKey;
 
-extern RelKeyData* pg_tde_create_key_map_entry(const RelFileLocator *newrlocator);
-extern void pg_tde_write_key_map_entry(const RelFileLocator *rlocator, RelKeyData *enc_rel_key_data, TDEPrincipalKeyInfo *principal_key_info);
-extern void pg_tde_delete_key_map_entry(const RelFileLocator *rlocator);
-extern void pg_tde_free_key_map_entry(const RelFileLocator *rlocator, off_t offset);
+extern RelKeyData *pg_tde_create_smgr_key(const RelFileLocator *newrlocator);
+extern RelKeyData *pg_tde_create_global_key(const RelFileLocator *newrlocator);
+extern RelKeyData *pg_tde_create_heap_basic_key(const RelFileLocator *newrlocator);
+extern RelKeyData *pg_tde_create_key_map_entry(const RelFileLocator *newrlocator, uint32 entry_type);
+extern void pg_tde_write_key_map_entry(const RelFileLocator *rlocator, uint32 entry_type, RelKeyData *enc_rel_key_data, TDEPrincipalKeyInfo *principal_key_info);
+extern void pg_tde_delete_key_map_entry(const RelFileLocator *rlocator, uint32 key_type);
+extern void pg_tde_free_key_map_entry(const RelFileLocator *rlocator, uint32 key_type, off_t offset);
 
-extern RelKeyData *GetRelationKey(RelFileLocator rel);
+extern RelKeyData *GetRelationKey(RelFileLocator rel, uint32 entry_type);
+extern RelKeyData *GetSMGRRelationKey(RelFileLocator rel);
+extern RelKeyData *GetHeapBaiscRelationKey(RelFileLocator rel);
+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_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(Oid rel_id, InternalKey *key, TDEPrincipalKeyInfo *principal_key_info);
+extern RelKeyData *tde_create_rel_key(RelFileNumber rel_num, uint32 key_type, InternalKey *key, TDEPrincipalKeyInfo *principal_key_info);
 extern RelKeyData *tde_encrypt_rel_key(TDEPrincipalKey *principal_key, RelKeyData *rel_key_data, const RelFileLocator *rlocator);
 extern RelKeyData *tde_decrypt_rel_key(TDEPrincipalKey *principal_key, RelKeyData *enc_rel_key_data, const RelFileLocator *rlocator);
-extern RelKeyData *pg_tde_get_key_from_file(const RelFileLocator *rlocator);
+extern RelKeyData *pg_tde_get_key_from_file(const RelFileLocator *rlocator, uint32 key_type);
 
 extern void pg_tde_set_db_file_paths(Oid dbOid, Oid spcOid, char *map_path, char *keydata_path);
 
 const char * tde_sprint_key(InternalKey *k);
 
-extern RelKeyData *pg_tde_put_key_into_cache(Oid rel_id, RelKeyData *key);
+extern RelKeyData *pg_tde_put_key_into_cache(RelFileNumber rel_num, uint32 key_type, RelKeyData *key);
 
 #endif /*PG_TDE_MAP_H*/
diff --git a/src/pg_tde_event_capture.c b/src/pg_tde_event_capture.c
@@ -69,7 +69,7 @@ pg_tde_ddl_command_start_capture(PG_FUNCTION_ARGS)
 	trigdata = (EventTriggerData *) fcinfo->context;
 	parsetree = trigdata->parsetree;
 
-	elog(DEBUG2, "EVENT TRIGGER (%s) %s", trigdata->event, nodeToString(parsetree));
+	elog(LOG, "EVENT TRIGGER (%s) %s", trigdata->event, nodeToString(parsetree));
 	reset_current_tde_create_event();
 
 	if (IsA(parsetree, IndexStmt))
@@ -184,7 +184,7 @@ pg_tde_ddl_command_end_capture(PG_FUNCTION_ARGS)
 		ereport(ERROR,
 				(errmsg("Function can only be fired by event trigger manager")));
 
-	elog(DEBUG1, "Type:%s EncryptMode:%s, Oid:%d, Relation:%s ",
+	elog(LOG, "Type:%s EncryptMode:%s, Oid:%d, Relation:%s ",
 		 (tdeCurrentCreateEvent.eventType == TDE_INDEX_CREATE_EVENT) ? "CREATE INDEX" :
 		 (tdeCurrentCreateEvent.eventType == TDE_TABLE_CREATE_EVENT) ? "CREATE TABLE" : "UNKNOWN",
 		 tdeCurrentCreateEvent.encryptMode ? "true" : "false",

diff --git a/src/smgr/pg_tde_smgr.c b/src/smgr/pg_tde_smgr.c
@@ -34,7 +34,7 @@ tde_smgr_get_key(SMgrRelation reln)
 	event = GetCurrentTdeCreateEvent();
 
 	// see if we have a key for the relation, and return if yes
-	rkd = GetRelationKey(reln->smgr_rlocator.locator);
+	rkd = GetSMGRRelationKey(reln->smgr_rlocator.locator);
 
 	if(rkd != NULL)
 	{
@@ -44,15 +44,15 @@ tde_smgr_get_key(SMgrRelation reln)
 	// if this is a CREATE TABLE, we have to generate the key
 	if(event->encryptMode == true && event->eventType == TDE_TABLE_CREATE_EVENT)
 	{
-		return pg_tde_create_key_map_entry(&reln->smgr_rlocator.locator);
+		return pg_tde_create_smgr_key(&reln->smgr_rlocator.locator);
 	}
 
 	// if this is a CREATE INDEX, we have to load the key based on the table
 	if(event->encryptMode == true && event->eventType == TDE_INDEX_CREATE_EVENT)
 	{
 		// For now keep it simple and create separate key for indexes
 		// Later we might modify the map infrastructure to support the same keys
-		return pg_tde_create_key_map_entry(&reln->smgr_rlocator.locator);
+		return pg_tde_create_smgr_key(&reln->smgr_rlocator.locator);
 	}
 
 	return NULL;

diff --git a/src/transam/pg_tde_xact_handler.c b/src/transam/pg_tde_xact_handler.c
@@ -123,7 +123,7 @@ do_pending_deletes(bool isCommit)
             ereport(LOG,
                     (errmsg("pg_tde_xact_callback: deleting entry at offset %d",
                             (int)(pending->map_entry_offset))));
-            pg_tde_free_key_map_entry(&pending->rlocator, pending->map_entry_offset);
+            pg_tde_free_key_map_entry(&pending->rlocator, MAP_ENTRY_VALID, pending->map_entry_offset);
         }
         pfree(pending);
         /* prev does not change */

diff --git a/src16/access/pg_tde_prune.c b/src16/access/pg_tde_prune.c
@@ -390,7 +390,7 @@ tdeheap_page_prune(Relation relation, Buffer buffer,
 	 * We need it here as there is `pgtde_compactify_tuples()` down in
 	 * the call stack wich reencrypt tuples.
 	*/
-	GetRelationKey(relation->rd_locator);
+	GetHeapBaiscRelationKey(relation->rd_locator);
 
 	/* Any error while applying the changes is critical */
 	START_CRIT_SECTION();

diff --git a/src16/access/pg_tdeam.c b/src16/access/pg_tdeam.c
@@ -1892,7 +1892,7 @@ tdeheap_insert(Relation relation, HeapTuple tup, CommandId cid,
 	 * Make sure relation keys in the cahce to avoid pallocs in
 	 * the critical section. 
 	*/
-	GetRelationKey(relation->rd_locator);
+	GetHeapBaiscRelationKey(relation->rd_locator);
 
 	/* NO EREPORT(ERROR) from here till changes are logged */
 	START_CRIT_SECTION();
@@ -2233,7 +2233,7 @@ tdeheap_multi_insert(Relation relation, TupleTableSlot **slots, int ntuples,
 		 * Make sure relation keys in the cahce to avoid pallocs in
 		 * the critical section. 
 		*/
-		GetRelationKey(relation->rd_locator);
+		GetHeapBaiscRelationKey(relation->rd_locator);
 
 		/* NO EREPORT(ERROR) from here till changes are logged */
 		START_CRIT_SECTION();
@@ -2799,7 +2799,7 @@ tdeheap_delete(Relation relation, ItemPointer tid,
 	 * won't be able to extract varlen attributes.
 	 */
 	decrypted_tuple = tdeheap_copytuple(&tp);
-	PG_TDE_DECRYPT_TUPLE(&tp, decrypted_tuple, GetRelationKey(relation->rd_locator));
+	PG_TDE_DECRYPT_TUPLE(&tp, decrypted_tuple, GetHeapBaiscRelationKey(relation->rd_locator));
 
 	old_key_tuple = ExtractReplicaIdentity(relation, decrypted_tuple, true, &old_key_copied);
 
@@ -3157,7 +3157,7 @@ tdeheap_update(Relation relation, ItemPointer otid, HeapTuple newtup,
 		oldtup_decrypted.t_data = (HeapTupleHeader)new_ptr;
 	}
 	PG_TDE_DECRYPT_TUPLE(&oldtup, &oldtup_decrypted,
-							GetRelationKey(relation->rd_locator));
+						 GetHeapBaiscRelationKey(relation->rd_locator));
 
 	// change field in oldtup now.
 	// We can't do it before, as PG_TDE_DECRYPT_TUPLE uses t_data address in 
@@ -3809,7 +3809,7 @@ tdeheap_update(Relation relation, ItemPointer otid, HeapTuple newtup,
 	 * Make sure relation keys in the cahce to avoid pallocs in
 	 * the critical section. 
 	*/
-	GetRelationKey(relation->rd_locator);
+	GetHeapBaiscRelationKey(relation->rd_locator);
 
 	/* NO EREPORT(ERROR) from here till changes are logged */
 	START_CRIT_SECTION();

diff --git a/src16/access/pg_tdeam_handler.c b/src16/access/pg_tdeam_handler.c
@@ -648,7 +648,7 @@ pg_tdeam_relation_set_new_filelocator(Relation rel,
 		ereport(DEBUG1,
 			(errmsg("creating key file for relation %s", RelationGetRelationName(rel))));
 
-		pg_tde_create_key_map_entry(newrlocator);
+		pg_tde_create_heap_basic_key(newrlocator);
 	}
 }
 

diff --git a/src16/access/pg_tdetoast.c b/src16/access/pg_tdetoast.c
@@ -657,7 +657,7 @@ tdeheap_fetch_toast_slice(Relation toastrel, Oid valueid, int32 attrsize,
 	int			validIndex;
 	SnapshotData SnapshotToast;
 	char		decrypted_data[TOAST_MAX_CHUNK_SIZE];
-	RelKeyData 	*key = GetRelationKey(toastrel->rd_locator);
+	RelKeyData *key = GetHeapBaiscRelationKey(toastrel->rd_locator);
 	char		iv_prefix[16] = {0,};
 
 
@@ -1093,7 +1093,7 @@ tdeheap_toast_save_datum(Relation rel, Datum value,
 	/*
 	* Encrypt toast data.
 	*/
-	tdeheap_toast_encrypt(dval, toast_pointer.va_valueid, GetRelationKey(toastrel->rd_locator));
+	tdeheap_toast_encrypt(dval, toast_pointer.va_valueid, GetHeapBaiscRelationKey(toastrel->rd_locator));
 
 	/*
 	 * Initialize constant parts of the tuple data

diff --git a/src17/access/pg_tde_prune.c b/src17/access/pg_tde_prune.c
@@ -761,7 +761,7 @@ tdeheap_page_prune_and_freeze(Relation relation, Buffer buffer,
 	 * We need it here as there is `pgtde_compactify_tuples()` down in
 	 * the call stack wich reencrypt tuples.
 	*/
-	GetRelationKey(relation->rd_locator);
+	GetHeapBaiscRelationKey(relation->rd_locator);
 
 	/* Any error while applying the changes is critical */
 	START_CRIT_SECTION();

diff --git a/src17/access/pg_tdeam.c b/src17/access/pg_tdeam.c
@@ -2050,7 +2050,7 @@ tdeheap_insert(Relation relation, HeapTuple tup, CommandId cid,
 	 * Make sure relation keys in the cahce to avoid pallocs in
 	 * the critical section. 
 	*/
-	GetRelationKey(relation->rd_locator);
+	GetHeapBaiscRelationKey(relation->rd_locator);
 
 	/* NO EREPORT(ERROR) from here till changes are logged */
 	START_CRIT_SECTION();
@@ -2391,7 +2391,7 @@ tdeheap_multi_insert(Relation relation, TupleTableSlot **slots, int ntuples,
 		 * Make sure relation keys in the cahce to avoid pallocs in
 		 * the critical section. 
 		*/
-		GetRelationKey(relation->rd_locator);
+		GetHeapBaiscRelationKey(relation->rd_locator);
 
 		/* NO EREPORT(ERROR) from here till changes are logged */
 		START_CRIT_SECTION();
@@ -2957,7 +2957,7 @@ tdeheap_delete(Relation relation, ItemPointer tid,
 	 * won't be able to extract varlen attributes.
 	 */
 	decrypted_tuple = tdeheap_copytuple(&tp);
-	PG_TDE_DECRYPT_TUPLE(&tp, decrypted_tuple, GetRelationKey(relation->rd_locator));
+	PG_TDE_DECRYPT_TUPLE(&tp, decrypted_tuple, GetHeapBaiscRelationKey(relation->rd_locator));
 
 	old_key_tuple = ExtractReplicaIdentity(relation, decrypted_tuple, true, &old_key_copied);
 
@@ -3315,7 +3315,7 @@ tdeheap_update(Relation relation, ItemPointer otid, HeapTuple newtup,
 		oldtup_decrypted.t_data = (HeapTupleHeader)new_ptr;
 	}
 	PG_TDE_DECRYPT_TUPLE(&oldtup, &oldtup_decrypted,
-							GetRelationKey(relation->rd_locator));
+						 GetHeapBaiscRelationKey(relation->rd_locator));
 
 	// change field in oldtup now.
 	// We can't do it before, as PG_TDE_DECRYPT_TUPLE uses t_data address in 
@@ -3967,7 +3967,7 @@ tdeheap_update(Relation relation, ItemPointer otid, HeapTuple newtup,
 	 * Make sure relation keys in the cahce to avoid pallocs in
 	 * the critical section. 
 	*/
-	GetRelationKey(relation->rd_locator);
+	GetHeapBaiscRelationKey(relation->rd_locator);
 
 	/* NO EREPORT(ERROR) from here till changes are logged */
 	START_CRIT_SECTION();

diff --git a/src17/access/pg_tdeam_handler.c b/src17/access/pg_tdeam_handler.c
@@ -645,7 +645,7 @@ pg_tdeam_relation_set_new_filelocator(Relation rel,
 		ereport(DEBUG1,
 			(errmsg("creating key file for relation %s", RelationGetRelationName(rel))));
 
-		pg_tde_create_key_map_entry(newrlocator);
+		pg_tde_create_heap_basic_key(newrlocator);
 	}
 }
 

diff --git a/src17/access/pg_tdetoast.c b/src17/access/pg_tdetoast.c
@@ -657,7 +657,7 @@ tdeheap_fetch_toast_slice(Relation toastrel, Oid valueid, int32 attrsize,
 	int			validIndex;
 	SnapshotData SnapshotToast;
 	char		decrypted_data[TOAST_MAX_CHUNK_SIZE];
-	RelKeyData 	*key = GetRelationKey(toastrel->rd_locator);
+	RelKeyData *key = GetHeapBaiscRelationKey(toastrel->rd_locator);
 	char		iv_prefix[16] = {0,};
 
 
@@ -1093,7 +1093,7 @@ tdeheap_toast_save_datum(Relation rel, Datum value,
 	/*
 	* Encrypt toast data.
 	*/
-	tdeheap_toast_encrypt(dval, toast_pointer.va_valueid, GetRelationKey(toastrel->rd_locator));
+	tdeheap_toast_encrypt(dval, toast_pointer.va_valueid, GetHeapBaiscRelationKey(toastrel->rd_locator));
 
 	/*
 	 * Initialize constant parts of the tuple data