PG-1604: Improve last key LSN calculation logic

Previosly we simply set the LSN for the new key to the first write
location.

This is however not correct, as there are many corner cases around this:
* recovery / replication might write old LSNs
* we can't handle multiple keys with the same TLI/LSN, which can happen
  with quick restarts without writes

To support this in this commit we modify the following:
* We only activate new keys outside crash recovery, or immediately if
  encryption is turned off
* We also take the already existing last key into account (if exists),
  and only activate a new key if we progressed past its start location

The remaining changes are just support infrastructure for this:

* Since we might rewrite old records, we use the already existing keys
  for those writes, not the active last keys
* We prefetch existing keys during initialization, so it doesn't
  accidentally happen in the critical section during a write

There is a remaining bug with stopping wal encryption, also mentioned in
a TODO message in the code. This will be addressed in a later PR as this
fix already took too long.

Author: dutow

contrib/pg_tde/meson.build

@@ -128,6 +128,8 @@ tap_tests = [
   't/wal_archiving.pl',
   't/wal_encrypt.pl',
   't/wal_key_tli.pl',
+  't/recovery/2pc_replication.pl',
+  't/recovery/stream_rep.pl',
 ]
 
 tests += {

contrib/pg_tde/src/access/pg_tde_xlog_smgr.c

@@ -246,6 +246,16 @@ TDEXLogSmgrInitWrite(bool encrypt_xlog)
 		TDEXLogSetEncKeyLocation(EncryptionKey.wal_start);
 	}
 
+	keys = pg_tde_get_wal_cache_keys();
+
+	if (keys == NULL)
+	{
+		WalLocation start = {.tli = 1,.lsn = 0};
+
+		/* cache is empty, prefetch keys from disk */
+		pg_tde_fetch_wal_keys(start);
+	}
+
 	if (key)
 		pfree(key);
 }
@@ -263,6 +273,32 @@ TDEXLogSmgrInitWriteReuseKey()
 	}
 }
 
+/*
+ * Encrypt XLog page(s) from the buf and write to the segment file.
+ */
+static ssize_t
+TDEXLogWriteEncryptedPagesOldKeys(int fd, const void *buf, size_t count, off_t offset,
+								  TimeLineID tli, XLogSegNo segno, int segSize)
+{
+	char	   *enc_buff = EncryptionBuf;
+
+#ifndef FRONTEND
+	Assert(count <= TDEXLogEncryptBuffSize());
+#endif
+
+	/* Copy the data as-is, as we might have unencrypted parts */
+	memcpy(enc_buff, buf, count);
+
+	/*
+	 * This method potentially allocates, but only in very early execution
+	 * Shouldn't happen in a write, where we are in a critical section
+	 */
+	TDEXLogCryptBuffer(buf, enc_buff, count, offset, tli, segno, segSize);
+
+	return pg_pwrite(fd, enc_buff, count, offset);
+}
+
+
 /*
  * Encrypt XLog page(s) from the buf and write to the segment file.
  */
@@ -284,6 +320,7 @@ TDEXLogWriteEncryptedPages(int fd, const void *buf, size_t count, off_t offset,
 #endif
 
 	CalcXLogPageIVPrefix(tli, segno, key->base_iv, iv_prefix);
+
 	pg_tde_stream_crypt(iv_prefix,
 						offset,
 						(char *) buf,
@@ -299,26 +336,66 @@ static ssize_t
 tdeheap_xlog_seg_write(int fd, const void *buf, size_t count, off_t offset,
 					   TimeLineID tli, XLogSegNo segno, int segSize)
 {
+	bool		lastKeyUsable;
+	bool		afterLastKey;
+#ifdef FRONTEND
+	bool		crashRecovery = false;
+#else
+	bool		crashRecovery = GetRecoveryState() == RECOVERY_STATE_CRASH;
+#endif
+
+	WalLocation loc = {.tli = tli};
+	WalLocation last_key_loc;
+
+	XLogSegNoOffsetToRecPtr(segno, offset, segSize, loc.lsn);
+
 	/*
 	 * Set the last (most recent) key's start LSN if not set.
 	 *
 	 * This func called with WALWriteLock held, so no need in any extra sync.
 	 */
-	if (EncryptionKey.type != WAL_KEY_TYPE_INVALID && TDEXLogGetEncKeyLsn() == 0)
-	{
-		WalLocation loc = {.tli = tli};
 
-		XLogSegNoOffsetToRecPtr(segno, offset, segSize, loc.lsn);
+	last_key_loc.lsn = TDEXLogGetEncKeyLsn();
+	pg_read_barrier();
+	last_key_loc.tli = TDEXLogGetEncKeyTli();
+
+	lastKeyUsable = (TDEXLogGetEncKeyLsn() != 0);
+	afterLastKey = wal_location_cmp(last_key_loc, loc) <= 0;
 
-		pg_tde_wal_last_key_set_location(loc);
-		EncryptionKey.wal_start = loc;
-		TDEXLogSetEncKeyLocation(EncryptionKey.wal_start);
+	Assert(EncryptionKey.type != WAL_KEY_TYPE_INVALID);
+
+	if (!lastKeyUsable)
+	{
+		WALKeyCacheRec *last_key = pg_tde_get_last_wal_key();
+
+		if (!crashRecovery || EncryptionKey.type == WAL_KEY_TYPE_UNENCRYPTED)
+		{
+			/*
+			 * TODO: the unencrypted case is still not perfect, we need to
+			 * report an error in some cornercases
+			 */
+			if (last_key == NULL || last_key->start.lsn < loc.lsn)
+			{
+				pg_tde_wal_last_key_set_location(loc);
+				EncryptionKey.wal_start = loc;
+				TDEXLogSetEncKeyLocation(EncryptionKey.wal_start);
+				lastKeyUsable = true;
+			}
+		}
 	}
 
-	if (EncryptionKey.type == WAL_KEY_TYPE_ENCRYPTED)
+	if ((!afterLastKey || !lastKeyUsable) && EncryptionKey.type == WAL_KEY_TYPE_ENCRYPTED)
+	{
+		return TDEXLogWriteEncryptedPagesOldKeys(fd, buf, count, offset, tli, segno, segSize);
+	}
+	else if (EncryptionKey.type == WAL_KEY_TYPE_ENCRYPTED)
+	{
 		return TDEXLogWriteEncryptedPages(fd, buf, count, offset, tli, segno);
+	}
 	else
+	{
 		return pg_pwrite(fd, buf, count, offset);
+	}
 }
 
 /*
@@ -340,7 +417,7 @@ tdeheap_xlog_seg_read(int fd, void *buf, size_t count, off_t offset,
 	if (readsz <= 0)
 		return readsz;
 
-	TDEXLogCryptBuffer(buf, count, offset, tli, segno, segSize);
+	TDEXLogCryptBuffer(buf, buf, count, offset, tli, segno, segSize);
 
 	return readsz;
 }
@@ -349,15 +426,15 @@ tdeheap_xlog_seg_read(int fd, void *buf, size_t count, off_t offset,
  * [De]Crypt buffer if needed based on provided segment offset, number and TLI
  */
 void
-TDEXLogCryptBuffer(void *buf, size_t count, off_t offset,
+TDEXLogCryptBuffer(const void *buf, void *out_buf, size_t count, off_t offset,
 				   TimeLineID tli, XLogSegNo segno, int segSize)
 {
 	WALKeyCacheRec *keys = pg_tde_get_wal_cache_keys();
 	XLogRecPtr	write_key_lsn;
 	WalLocation data_end = {.tli = tli};
 	WalLocation data_start = {.tli = tli};
 
-	if (!keys)
+	if (keys == NULL)
 	{
 		WalLocation start = {.tli = 1,.lsn = 0};
 
@@ -454,6 +531,7 @@ TDEXLogCryptBuffer(void *buf, size_t count, off_t offset,
 					XLogSegmentOffset(end_lsn, segSize);
 				size_t		dec_sz;
 				char	   *dec_buf = (char *) buf + (dec_off - offset);
+				char	   *o_buf = (char *) out_buf + (dec_off - offset);
 
 				Assert(dec_off >= offset);
 
@@ -468,17 +546,19 @@ TDEXLogCryptBuffer(void *buf, size_t count, off_t offset,
 					dec_end = offset + count;
 				}
 
+				Assert(dec_end > dec_off);
 				dec_sz = dec_end - dec_off;
 
 #ifdef TDE_XLOG_DEBUG
 				elog(DEBUG1, "decrypt WAL, dec_off: %lu [buff_off %lu], sz: %lu | key %u_%X/%X",
 					 dec_off, dec_off - offset, dec_sz, curr_key->key.wal_start.tli, LSN_FORMAT_ARGS(curr_key->key.wal_start.lsn));
 #endif
+
 				pg_tde_stream_crypt(iv_prefix,
 									dec_off,
 									dec_buf,
 									dec_sz,
-									dec_buf,
+									o_buf,
 									curr_key->key.key,
 									&curr_key->crypt_ctx);
 			}

contrib/pg_tde/src/include/access/pg_tde_xlog_smgr.h

@@ -13,7 +13,7 @@ extern void TDEXLogSmgrInit(void);
 extern void TDEXLogSmgrInitWrite(bool encrypt_xlog);
 extern void TDEXLogSmgrInitWriteReuseKey(void);
 
-extern void TDEXLogCryptBuffer(void *buf, size_t count, off_t offset,
+extern void TDEXLogCryptBuffer(const void *buf, void *out_buf, size_t count, off_t offset,
 							   TimeLineID tli, XLogSegNo segno, int segSize);
 
 #endif							/* PG_TDE_XLOGSMGR_H */