PG-1605 take 2: only switch keys after recovery ended

Or alternatively, if we progressed at least 1 page after the previous
key. This makes sure that even replicas, which are always in recovery
start using a new key with a small delay, while also allowing a gap for
crash recovery.

Author: dutow

contrib/pg_tde/meson.build

@@ -127,6 +127,7 @@ tap_tests = [
   't/wal_archiving.pl',
   't/wal_encrypt.pl',
   't/wal_key_tli.pl',
+  't/059_tde_2pc_replication.pl',
 ]
 
 tests += {

contrib/pg_tde/src/access/pg_tde_xlog_smgr.c

@@ -48,6 +48,16 @@ static WalEncryptionKey EncryptionKey =
 	.wal_start = {.tli = 0,.lsn = InvalidXLogRecPtr},
 };
 
+static void
+iv_prefix_debug(const char *iv_prefix, char *out_hex)
+{
+	for (int i = 0; i < 16; ++i)
+	{
+		sprintf(out_hex + i * 2, "%02x", (int) *(iv_prefix + i));
+	}
+	out_hex[32] = 0;
+}
+
 /*
  * Must be the same as in replication/walsender.c
  *
@@ -248,6 +258,13 @@ TDEXLogSmgrInitWrite(bool encrypt_xlog)
 		TDEXLogSetEncKeyLocation(EncryptionKey.wal_start);
 	}
 
+	{
+		WalLocation start = {.tli = 1,.lsn = 0};
+
+		/* TODO cache is empty, try to preread keys from disk */
+		//pg_tde_fetch_wal_keys(start);
+	}
+
 	if (key)
 		pfree(key);
 }
@@ -265,6 +282,27 @@ 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
+
+	/* 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.
  */
@@ -280,12 +318,17 @@ TDEXLogWriteEncryptedPages(int fd, const void *buf, size_t count, off_t offset,
 	Assert(count <= TDEXLogEncryptBuffSize());
 #endif
 
+	CalcXLogPageIVPrefix(tli, segno, key->base_iv, iv_prefix);
+
 #ifdef TDE_XLOG_DEBUG
-	elog(DEBUG1, "write encrypted WAL, size: %lu, offset: %ld [%lX], seg: %X/%X, key_start_lsn: %u_%X/%X",
-		 count, offset, offset, LSN_FORMAT_ARGS(segno), key->wal_start.tli, LSN_FORMAT_ARGS(key->wal_start.lsn));
+	char		tbuf[33], tbuf2[33], tbuf3[33];
+	iv_prefix_debug(key->key, tbuf);
+	iv_prefix_debug(iv_prefix, tbuf2);
+	iv_prefix_debug(buf, tbuf3);
+	elog(DEBUG1, "write encrypted WAL, size: %lu, offset: %ld [%lX], seg: %X/%X, key_start_lsn: %u_%X/%X %s %s %s",
+		 count, offset, offset, LSN_FORMAT_ARGS(segno), key->wal_start.tli, LSN_FORMAT_ARGS(key->wal_start.lsn), tbuf, tbuf2, tbuf3);
 #endif
 
-	CalcXLogPageIVPrefix(tli, segno, key->base_iv, iv_prefix);
 	pg_tde_stream_crypt(iv_prefix,
 						offset,
 						(char *) buf,
@@ -301,26 +344,53 @@ static ssize_t
 tdeheap_xlog_seg_write(int fd, const void *buf, size_t count, off_t offset,
 					   TimeLineID tli, XLogSegNo segno, int segSize)
 {
+	#ifdef FRONTEND
+	bool recoveryInProgress = false;
+	#else
+	bool recoveryInProgress = RecoveryInProgress();
+	#endif
 	/*
 	 * 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)
+
+	bool lastKeyUsable = (TDEXLogGetEncKeyLsn() != 0);
+
+	if (EncryptionKey.type != WAL_KEY_TYPE_INVALID && !lastKeyUsable)
 	{
 		WalLocation loc = {.tli = tli};
 
 		XLogSegNoOffsetToRecPtr(segno, offset, segSize, loc.lsn);
 
+		if(!recoveryInProgress || EncryptionKey.type == WAL_KEY_TYPE_UNENCRYPTED) {
+			// TODO: detect if we have conflicting later keys in case of unencrypted, and error out
 		pg_tde_wal_last_key_set_location(loc);
 		EncryptionKey.wal_start = loc;
 		TDEXLogSetEncKeyLocation(EncryptionKey.wal_start);
+		lastKeyUsable = true;
+		} else {
+			WALKeyCacheRec *last_key = pg_tde_get_last_wal_key();
+			if(last_key) {
+				if (wal_location_cmp(last_key->start, loc) < -8192) {
+					// If we progressed at least one more page, it is safe to switch keys
+					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)
-		return TDEXLogWriteEncryptedPages(fd, buf, count, offset, tli, segno);
-	else
+	if(!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);
+	}
 }
 
 /*
@@ -342,7 +412,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;
 }
@@ -351,7 +421,7 @@ 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();
@@ -423,6 +493,7 @@ TDEXLogCryptBuffer(void *buf, size_t count, off_t offset,
 				off_t		dec_end = XLogSegmentOffset(minlsn, 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);
 
@@ -436,18 +507,23 @@ TDEXLogCryptBuffer(void *buf, size_t count, off_t offset,
 
 				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);
-			}
+#ifdef TDE_XLOG_DEBUG
+char		tbuf[33], tbuf2[33], tbuf3[33];
+	iv_prefix_debug(curr_key->key.key, tbuf);
+	iv_prefix_debug(iv_prefix, tbuf2);
+	iv_prefix_debug(o_buf, tbuf3);
+				elog(DEBUG1, "decrypt WAL, dec_off: %lu [buff_off %lu], sz: %lu | key %u_%X/%X %s %s %s",
+					 dec_off, dec_off - offset, dec_sz, curr_key->key.wal_start.tli, LSN_FORMAT_ARGS(curr_key->key.wal_start.lsn), tbuf, tbuf2, tbuf3);
+#endif
+								}
 		}
 	}
 }

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 */

contrib/pg_tde/src/include/pg_tde_defines.h

@@ -7,6 +7,6 @@
 
 /* #define ENCRYPTION_DEBUG 1 */
 /* #define KEYRING_DEBUG 1 */
-/* #define TDE_XLOG_DEBUG 1 */
+#define TDE_XLOG_DEBUG 1
 
 #endif							/* PG_TDE_DEFINES_H */