Do not use principal key to check encryption status

In many cases it's completely unnecessary to have the principal key in order to know if a relation is encrypted or not. This simplifies cases where principal key provider is not available for any reason. The unencrypted key should most likely be cached, but this commit does not do that.
percona · AndersAstrand · May 14, 2025 · May 13, 2025 · May 13, 2025 · May 13, 2025
commit 5816caa61a793185f6d6b67eda54312960b03df1
@@ -1074,6 +1074,38 @@ pg_tde_get_principal_key_info(Oid dbOid)
 	return signed_key_info;
 }
 
+/*
+ * Figures out whether a relation is encrypted or not, but without trying to
+ * decrypt the key if it is. This also means that this function cannot push the
+ * key to cache.
+ */
+bool
+IsSMGRRelationEncrypted(RelFileLocatorBackend rel)
+{
+	bool		result;
+	TDEMapEntry map_entry;
+	char		db_map_path[MAXPGPATH];
+
+	Assert(rel.locator.relNumber != InvalidRelFileNumber);
+
+	if (RelFileLocatorBackendIsTemp(rel))
+		return pg_tde_get_key_from_cache(&rel.locator, TDE_KEY_TYPE_SMGR) != NULL;
+	else if (pg_tde_get_key_from_cache(&rel.locator, TDE_KEY_TYPE_SMGR))
+		return true;
+
+	pg_tde_set_db_file_path(rel.locator.dbOid, db_map_path);
+
+	if (access(db_map_path, F_OK) == -1)
+		return false;
+
+	LWLockAcquire(tde_lwlock_enc_keys(), LW_SHARED);
+
+	result = pg_tde_find_map_entry(&rel.locator, TDE_KEY_TYPE_SMGR, db_map_path, &map_entry);
+
+	LWLockRelease(tde_lwlock_enc_keys());
+	return result;
+}
+
 /*
  * Returns TDE key for a given relation.
  * First it looks in a cache. If nothing found in the cache, it reads data from

@@ -29,7 +29,7 @@ pg_tde_is_encrypted(PG_FUNCTION_ARGS)
 	LOCKMODE	lockmode = AccessShareLock;
 	Relation	rel = relation_open(relationOid, lockmode);
 	RelFileLocatorBackend rlocator = {.locator = rel->rd_locator,.backend = rel->rd_backend};
-	InternalKey *key;
+	bool		result;
 
 	if (!RELKIND_HAS_STORAGE(rel->rd_rel->relkind))
 	{
@@ -42,11 +42,11 @@ pg_tde_is_encrypted(PG_FUNCTION_ARGS)
 				errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
 				errmsg("we cannot check if temporary relations from other backends are encrypted"));
 
-	key = GetSMGRRelationKey(rlocator);
+	result = IsSMGRRelationEncrypted(rlocator);
 
 	relation_close(rel, lockmode);
 
-	PG_RETURN_BOOL(key != NULL);
+	PG_RETURN_BOOL(result);
 }
 
 #endif							/* !FRONTEND */

@@ -109,6 +109,7 @@ pg_tde_set_db_file_path(Oid dbOid, char *path)
 	join_path_components(path, pg_tde_get_data_dir(), psprintf(PG_TDE_MAP_FILENAME, dbOid));
 }
 
+extern bool IsSMGRRelationEncrypted(RelFileLocatorBackend rel);
 extern InternalKey *GetSMGRRelationKey(RelFileLocatorBackend rel);
 extern int	pg_tde_count_relations(Oid dbOid);
 

@@ -67,6 +67,7 @@ tde_smgr_should_encrypt(const RelFileLocatorBackend *smgr_rlocator, RelFileLocat
 					.backend = smgr_rlocator->backend,
 				};
 
+				/* Actually get the key here to ensure result is cached. */
 				return GetSMGRRelationKey(old_smgr_locator) != 0;
 			}
 	}
@@ -111,6 +112,11 @@ tde_mdwritev(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum,
 	}
 }
 
+/*
+ * The current transaction might already be commited when this function is
+ * called, so do not call any code that uses ereport(ERROR) or otherwise tries
+ * to abort the transaction.
+ */
 static void
 tde_mdunlink(RelFileLocatorBackend rlocator, ForkNumber forknum, bool isRedo)
 {
@@ -128,7 +134,7 @@ tde_mdunlink(RelFileLocatorBackend rlocator, ForkNumber forknum, bool isRedo)
 	 */
 	if (forknum == MAIN_FORKNUM || forknum == InvalidForkNumber)
 	{
-		if (!RelFileLocatorBackendIsTemp(rlocator) && GetSMGRRelationKey(rlocator))
+		if (!RelFileLocatorBackendIsTemp(rlocator) && IsSMGRRelationEncrypted(rlocator))
 			pg_tde_free_key_map_entry(&rlocator.locator);
 	}
 }

@@ -121,7 +121,11 @@ SELECT * FROM test_enc ORDER BY id;
 SELECT pg_tde_verify_key();
 psql:<stdin>:1: ERROR:  failed to retrieve principal key test-key from keyring with ID 1
 SELECT pg_tde_is_encrypted('test_enc');
-psql:<stdin>:1: ERROR:  failed to retrieve principal key test-key from keyring with ID 1
+ pg_tde_is_encrypted 
+---------------------
+ t
+(1 row)
+
 SELECT * FROM test_enc ORDER BY id;
 psql:<stdin>:1: ERROR:  failed to retrieve principal key test-key from keyring with ID 1
 -- mv /tmp/change_key_provider_2.per /tmp/change_key_provider_3.per
@@ -191,7 +195,11 @@ SELECT pg_tde_change_database_key_provider_file('file-vault', '/tmp/change_key_p
 SELECT pg_tde_verify_key();
 psql:<stdin>:1: ERROR:  Failed to verify principal key header for key test-key, incorrect principal key or corrupted key file
 SELECT pg_tde_is_encrypted('test_enc');
-psql:<stdin>:1: ERROR:  Failed to verify principal key header for key test-key, incorrect principal key or corrupted key file
+ pg_tde_is_encrypted 
+---------------------
+ t
+(1 row)
+
 SELECT * FROM test_enc ORDER BY id;
 psql:<stdin>:1: ERROR:  Failed to verify principal key header for key test-key, incorrect principal key or corrupted key file
 CREATE TABLE test_enc2 (id serial, k integer, PRIMARY KEY (id)) USING tde_heap;

@@ -141,7 +141,7 @@ is_pg_tde_encypted(Oid spcOid, Oid dbOid, RelFileNumber relNumber)
 	RelFileLocator locator = {.spcOid = spcOid, .dbOid = dbOid,.relNumber = relNumber};
 	RelFileLocatorBackend rlocator = {.locator = locator,.backend = INVALID_PROC_NUMBER};
 
-	return GetSMGRRelationKey(rlocator) != NULL;
+	return IsSMGRRelationEncrypted(rlocator);
 }
 #endif