PG-1446 Add missing lock to pg_tde_create_wal_key()

Make sure that we lock the key file before adding a WAL key. It should
be harmless since this happens very early in the start but if we ever
add the ability to rotate WAL keys while running this will become an
issue.

Additionally we add asserts to make sure we hold the right lock level
when we open the key map file.

Author: jeltz

contrib/pg_tde/src/access/pg_tde_tdemap.c

@@ -222,7 +222,9 @@ pg_tde_create_wal_key(InternalKey *rel_key_data, const RelFileLocator *newrlocat
 {
 	TDEPrincipalKey *principal_key;
 
-	principal_key = get_principal_key_from_keyring(newrlocator->dbOid);
+	LWLockAcquire(tde_lwlock_enc_keys(), LW_SHARED);
+
+	principal_key = GetPrincipalKey(newrlocator->dbOid, LW_SHARED);
 	if (principal_key == NULL)
 	{
 		ereport(ERROR,
@@ -237,6 +239,8 @@ pg_tde_create_wal_key(InternalKey *rel_key_data, const RelFileLocator *newrlocat
 	 * Add the encrypted key to the key map data file structure.
 	 */
 	pg_tde_write_key_map_entry(newrlocator, rel_key_data, principal_key, false);
+
+	LWLockRelease(tde_lwlock_enc_keys());
 }
 
 /*
@@ -872,6 +876,8 @@ pg_tde_open_file_write(const char *tde_filename, TDEPrincipalKeyInfo *principal_
 	off_t		bytes_written = 0;
 	int			file_flags = O_RDWR | O_CREAT | PG_BINARY | (truncate ? O_TRUNC : 0);
 
+	Assert(LWLockHeldByMeInMode(tde_lwlock_enc_keys(), LW_EXCLUSIVE));
+
 	fd = pg_tde_open_file_basic(tde_filename, file_flags, false);
 
 	pg_tde_file_header_read(tde_filename, fd, &fheader, &bytes_read);
@@ -1021,6 +1027,8 @@ pg_tde_open_file_read(const char *tde_filename, off_t *curr_pos)
 	TDEFileHeader fheader;
 	off_t		bytes_read = 0;
 
+	Assert(LWLockHeldByMeInMode(tde_lwlock_enc_keys(), LW_SHARED) || LWLockHeldByMeInMode(tde_lwlock_enc_keys(), LW_EXCLUSIVE));
+
 	fd = pg_tde_open_file_basic(tde_filename, O_RDONLY | PG_BINARY, false);
 
 	pg_tde_file_header_read(tde_filename, fd, &fheader, &bytes_read);

contrib/pg_tde/src/catalog/tde_principal_key.c

@@ -100,6 +100,7 @@ static bool pg_tde_is_same_principal_key(TDEPrincipalKey *a, TDEPrincipalKey *b)
 static void pg_tde_update_global_principal_key_everywhere(TDEPrincipalKey *oldKey, TDEPrincipalKey *newKey);
 static void push_principal_key_to_cache(TDEPrincipalKey *principalKey);
 static Datum pg_tde_get_key_info(PG_FUNCTION_ARGS, Oid dbOid);
+static TDEPrincipalKey *get_principal_key_from_keyring(Oid dbOid);
 static TDEPrincipalKey *GetPrincipalKeyNoDefault(Oid dbOid, LWLockMode lockMode);
 static void set_principal_key_with_keyring(const char *key_name,
 										   const char *provider_name,
@@ -708,7 +709,7 @@ pg_tde_get_key_info(PG_FUNCTION_ARGS, Oid dbOid)
  *
  * Caller should hold an exclusive tde_lwlock_enc_keys lock
  */
-TDEPrincipalKey *
+static TDEPrincipalKey *
 get_principal_key_from_keyring(Oid dbOid)
 {
 	TDEPrincipalKeyInfo *principalKeyInfo;
@@ -717,7 +718,7 @@ get_principal_key_from_keyring(Oid dbOid)
 	KeyringReturnCodes keyring_ret;
 	TDEPrincipalKey *principalKey;
 
-	/* Assert(LWLockHeldByMeInMode(tde_lwlock_enc_keys(), LW_EXCLUSIVE)); */
+	Assert(LWLockHeldByMeInMode(tde_lwlock_enc_keys(), LW_EXCLUSIVE));
 
 	principalKeyInfo = pg_tde_get_principal_key_info(dbOid);
 	if (principalKeyInfo == NULL)

contrib/pg_tde/src/include/catalog/tde_principal_key.h

@@ -54,6 +54,5 @@ extern TDEPrincipalKey *GetPrincipalKey(Oid dbOid, void *lockMode);
 #endif
 
 extern void xl_tde_perform_rotate_key(XLogPrincipalKeyRotate *xlrec);
-extern TDEPrincipalKey *get_principal_key_from_keyring(Oid dbOid);
 
 #endif							/* PG_TDE_PRINCIPAL_KEY_H */