Prevent TDE SMGR from engaging with tde_heap_basic

1. Currently, tde smgr functions (like tde_mdwritev) will encrypt/decrypt
tde_heap_basic relations because the only rule of engagement is the
existence of the internal key. And any tde_heap_basic would have one.
This commit creates a rel_type for the internal key so tde smgr can
distinguish if that key (hence relation) was meant for the "tde_heap".

2. During ALTER ... TABLESPACE, tde_mdcreate gets called and
tde_smgr_get_key is trying to get the internal key. But there is no key
moved to the new table space since it is called during
RelationCreateStorage() and pg_tde.map file might not exist. And it's ok in
this case but it results in an error. This commit introduces an option not
to have a pg_tde.map (treated as no key).

src/access/pg_tde_slot.c

@@ -582,6 +582,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 = GetRelationKey(rel->rd_locator, false);
 	return bslot->cached_relation_key;
 }

src/access/pg_tde_tdemap.c

@@ -137,9 +137,9 @@ static void finalize_key_rotation(char *m_path_old, char *k_path_old, char *m_pa
  * Generate an encrypted key for the relation and store it in the keymap file.
  */
 RelKeyData*
-pg_tde_create_key_map_entry(const RelFileLocator *newrlocator)
+pg_tde_create_key_map_entry(const RelFileLocator *newrlocator, InternalKeyRelType rel_type)
 {
-	InternalKey int_key;
+	InternalKey int_key = {.rel_type = rel_type};
 	RelKeyData *rel_key_data;
 	RelKeyData *enc_rel_key_data;
 	TDEPrincipalKey *principal_key;
@@ -919,7 +919,7 @@ pg_tde_move_rel_key(const RelFileLocator *newrlocator, const RelFileLocator *old
  * reads the key data from the keydata file.
  */
 RelKeyData *
-pg_tde_get_key_from_file(const RelFileLocator *rlocator)
+pg_tde_get_key_from_file(const RelFileLocator *rlocator, bool no_map_ok)
 {
 	int32		key_index = 0;
 	TDEPrincipalKey	*principal_key;
@@ -957,6 +957,11 @@ pg_tde_get_key_from_file(const RelFileLocator *rlocator)
 	/* Get the file paths */
 	pg_tde_set_db_file_paths(rlocator->dbOid, rlocator->spcOid, db_map_path, db_keydata_path);
 
+	if (no_map_ok && access(db_map_path, F_OK) == -1)
+	{
+		LWLockRelease(lock_pk);
+		return NULL;
+	}
 	/* Read the map entry and get the index of the relation key */
 	key_index = pg_tde_process_map_entry(rlocator, db_map_path, &offset, false);
 
@@ -1178,7 +1183,7 @@ pg_tde_open_file_basic(char *tde_filename, int fileFlags, bool ignore_missing)
 	fd = BasicOpenFile(tde_filename, fileFlags | PG_BINARY);
 	if (fd < 0 && !(errno == ENOENT && ignore_missing == true))
 	{
-		ereport(ERROR,
+		ereport(PANIC,
 				(errcode_for_file_access(),
 				 errmsg("could not open tde file \"%s\": %m",
 						tde_filename)));
@@ -1350,9 +1355,12 @@ pg_tde_get_principal_key_info(Oid dbOid, Oid spcOid)
  * Returns TDE key for a given relation.
  * First it looks in a cache. If nothing found in the cache, it reads data from
  * the tde fork file and populates cache.
+ * 
+ * If no_map_ok is true - return NULL key if pg_tde.map dosn't exist. It's a
+ * temporary hack while we fixing smgr relations map.
  */
 RelKeyData *
-GetRelationKey(RelFileLocator rel)
+GetRelationKey(RelFileLocator rel, bool no_map_ok)
 {
 	RelKeyData *key;
 	Oid rel_id = rel.relNumber;
@@ -1363,7 +1371,7 @@ GetRelationKey(RelFileLocator rel)
 		return key;
 	}
 
-	key = pg_tde_get_key_from_file(&rel);
+	key = pg_tde_get_key_from_file(&rel, no_map_ok);
 
 	if (key != NULL)
 	{

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 = GetRelationKey(GLOBAL_SPACE_RLOCATOR(XLOG_TDE_OID), false);
 	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 = GetRelationKey(GLOBAL_SPACE_RLOCATOR(XLOG_TDE_OID), false);
 	size_t	page_size = XLOG_BLCKSZ - offset % XLOG_BLCKSZ;
 	off_t	dec_off;
 	uint32	iv_ctr = 0;

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), false);
 
 		/*
 		 * Internal Key should be in the TopMemmoryContext because of SSL
@@ -134,7 +134,7 @@ init_default_keyring(void)
 static void
 init_keys(void)
 {
-	InternalKey int_key;
+	InternalKey int_key = {.rel_type = TDE_IKEY_REL_GLOBAL};
 	RelKeyData *rel_key_data;
 	RelKeyData *enc_rel_key_data;
 	RelFileLocator *rlocator;

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 = GetRelationKey(rel, false);
 
 	ItemPointerSet(&ip, bn, off); 
 

src/include/access/pg_tde_tdemap.h

@@ -14,10 +14,20 @@
 #include "catalog/tde_principal_key.h"
 #include "storage/relfilelocator.h"
 
+/* TODO: get rid of that (see. commit message) */
+typedef enum InternalKeyRelType
+{
+	TDE_IKEY_REL_UNKNOWN,
+	TDE_IKEY_REL_GLOBAL,
+	TDE_IKEY_REL_SMGR,
+	TDE_IKEY_REL_BASIC
+} InternalKeyRelType;
+
 typedef struct InternalKey
 {
-    uint8   key[INTERNAL_KEY_LEN];
-	void*   ctx; // TODO: shouldn't be here / written to the disk
+    uint8			key[INTERNAL_KEY_LEN];
+	InternalKeyRelType	rel_type;	
+	void*			ctx; // TODO: shouldn't be here / written to the disk
 } InternalKey;
 
 typedef struct RelKeyData
@@ -34,12 +44,12 @@ typedef struct XLogRelKey
 	TDEPrincipalKeyInfo pkInfo;
 } XLogRelKey;
 
-extern RelKeyData* pg_tde_create_key_map_entry(const RelFileLocator *newrlocator);
+extern RelKeyData* pg_tde_create_key_map_entry(const RelFileLocator *newrlocator, InternalKeyRelType ktype);
 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 *GetRelationKey(RelFileLocator rel);
+extern RelKeyData *GetRelationKey(RelFileLocator rel, bool no_map_is_ok);
 
 extern void pg_tde_delete_tde_files(Oid dbOid, Oid spcOid);
 
@@ -50,7 +60,7 @@ extern bool pg_tde_write_map_keydata_files(off_t map_size, char *m_file_data, of
 extern RelKeyData* tde_create_rel_key(Oid rel_id, 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, bool no_map_ok);
 extern bool pg_tde_move_rel_key(const RelFileLocator *newrlocator, const RelFileLocator *oldrlocator);
 
 extern void pg_tde_set_db_file_paths(Oid dbOid, Oid spcOid, char *map_path, char *keydata_path);

src/smgr/pg_tde_smgr.c

@@ -34,25 +34,29 @@ 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 = GetRelationKey(reln->smgr_rlocator.locator, true);
 
 	if(rkd != NULL)
 	{
+		/* TODO: get rid of that (see. commit message) */
+		if (rkd->internal_key.rel_type != TDE_IKEY_REL_SMGR)
+			return NULL;	
+
 		return rkd;
 	}
 
 	// 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_key_map_entry(&reln->smgr_rlocator.locator, TDE_IKEY_REL_SMGR);
 	}
 
 	// 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_key_map_entry(&reln->smgr_rlocator.locator, TDE_IKEY_REL_SMGR);
 	}
 
 	return NULL;

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);
+	GetRelationKey(relation->rd_locator, false);
 
 	/* Any error while applying the changes is critical */
 	START_CRIT_SECTION();

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);
+	GetRelationKey(relation->rd_locator, false);
 
 	/* 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);
+		GetRelationKey(relation->rd_locator, false);
 
 		/* 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, GetRelationKey(relation->rd_locator), false);
 
 	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));
+							GetRelationKey(relation->rd_locator), false);
 
 	// 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);
+	GetRelationKey(relation->rd_locator, false);
 
 	/* NO EREPORT(ERROR) from here till changes are logged */
 	START_CRIT_SECTION();

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_key_map_entry(newrlocator, TDE_IKEY_REL_BASIC);
 	}
 }
 

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 = GetRelationKey(toastrel->rd_locator, false);
 	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, GetRelationKey(toastrel->rd_locator, false));
 
 	/*
 	 * Initialize constant parts of the tuple data

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);
+	GetRelationKey(relation->rd_locator, false);
 
 	/* Any error while applying the changes is critical */
 	START_CRIT_SECTION();

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);
+	GetRelationKey(relation->rd_locator, false);
 
 	/* 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);
+		GetRelationKey(relation->rd_locator, false);
 
 		/* 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, GetRelationKey(relation->rd_locator, false));
 
 	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));
+							GetRelationKey(relation->rd_locator, false));
 
 	// 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);
+	GetRelationKey(relation->rd_locator, false);
 
 	/* NO EREPORT(ERROR) from here till changes are logged */
 	START_CRIT_SECTION();

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_key_map_entry(newrlocator, TDE_IKEY_REL_BASIC);
 	}
 }
 
@@ -677,6 +677,8 @@ pg_tdeam_relation_copy_data(Relation rel, const RelFileLocator *newrlocator)
 	 */
 	dstrel = RelationCreateStorage(*newrlocator, rel->rd_rel->relpersistence, true);
 
+	pg_tde_move_rel_key(newrlocator, &rel->rd_locator);
+
 	/* copy main fork */
 	RelationCopyStorage(RelationGetSmgr(rel), dstrel, MAIN_FORKNUM,
 						rel->rd_rel->relpersistence);
@@ -702,8 +704,6 @@ pg_tdeam_relation_copy_data(Relation rel, const RelFileLocator *newrlocator)
 		}
 	}
 
-	pg_tde_move_rel_key(newrlocator, &rel->rd_locator);
-
 	/* drop old relation, and close new one */
 	RelationDropStorage(rel);
 	smgrclose(dstrel);

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 = GetRelationKey(toastrel->rd_locator, false);
 	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, GetRelationKey(toastrel->rd_locator, false));
 
 	/*
 	 * Initialize constant parts of the tuple data