TLI

Author: dAdAbird

contrib/pg_tde/src/access/pg_tde_tdemap.c

@@ -44,6 +44,7 @@
 #define TDE_FILE_HEADER_SIZE	sizeof(TDEFileHeader)
 
 #define MaxXLogRecPtr (~(XLogRecPtr)0)
+#define MaxTimeLineID (~(TimeLineID)0)
 
 typedef struct TDEFileHeader
 {
@@ -369,13 +370,19 @@ pg_tde_delete_principal_key(Oid dbOid)
  * needs keyfile_path
  */
 void
-pg_tde_wal_last_key_set_lsn(XLogRecPtr lsn, const char *keyfile_path)
+pg_tde_wal_last_key_set_lsn(XLogRecPtr lsn, TimeLineID tli, const char *keyfile_path)
 {
 	LWLock	   *lock_pk = tde_lwlock_enc_keys();
 	int			fd;
 	off_t		read_pos,
 				write_pos,
 				last_key_idx;
+	struct
+	{
+			XLogRecPtr	start_lsn;
+			TimeLineID	tli;
+	} lsn_tli;
+	
 
 	LWLockAcquire(lock_pk, LW_EXCLUSIVE);
 
@@ -384,7 +391,10 @@ pg_tde_wal_last_key_set_lsn(XLogRecPtr lsn, const char *keyfile_path)
 	last_key_idx = ((lseek(fd, 0, SEEK_END) - TDE_FILE_HEADER_SIZE) / MAP_ENTRY_SIZE) - 1;
 	write_pos = TDE_FILE_HEADER_SIZE + (last_key_idx * MAP_ENTRY_SIZE) + offsetof(TDEMapEntry, enc_key) + offsetof(InternalKey, start_lsn);
 
-	if (pg_pwrite(fd, &lsn, sizeof(XLogRecPtr), write_pos) != sizeof(XLogRecPtr))
+	lsn_tli.start_lsn = lsn;
+	lsn_tli.tli = tli;
+
+	if (pg_pwrite(fd, &lsn_tli, sizeof(lsn_tli), write_pos) != sizeof(lsn_tli))
 	{
 		ereport(ERROR,
 				errcode_for_file_access(),
@@ -408,7 +418,7 @@ pg_tde_wal_last_key_set_lsn(XLogRecPtr lsn, const char *keyfile_path)
 					errmsg("could not read previous WAL key: %m"));
 		}
 
-		if (prev_map_entry.enc_key.start_lsn >= lsn)
+		if (prev_map_entry.enc_key.start_lsn >= lsn && prev_map_entry.enc_key.tli >= tli)
 		{
 			prev_map_entry.enc_key.type = TDE_KEY_TYPE_WAL_INVALID;
 
@@ -1071,6 +1081,7 @@ pg_tde_fetch_wal_keys(XLogRecPtr start_lsn)
 		WALKeyCacheRec *wal_rec;
 		InternalKey stub_key = {
 			.start_lsn = InvalidXLogRecPtr,
+			.tli = 0,
 		};
 
 		wal_rec = pg_tde_add_wal_key_to_cache(&stub_key, InvalidXLogRecPtr);
@@ -1132,8 +1143,10 @@ pg_tde_add_wal_key_to_cache(InternalKey *key, XLogRecPtr start_lsn)
 	MemoryContextSwitchTo(oldCtx);
 #endif
 
+	wal_rec->start_tli = key->tli;
 	wal_rec->start_lsn = start_lsn;
 	wal_rec->end_lsn = MaxXLogRecPtr;
+	wal_rec->end_tli = MaxTimeLineID;
 	wal_rec->key = *key;
 	wal_rec->crypt_ctx = NULL;
 	if (!tde_wal_key_last_rec)
@@ -1145,6 +1158,7 @@ pg_tde_add_wal_key_to_cache(InternalKey *key, XLogRecPtr start_lsn)
 	{
 		tde_wal_key_last_rec->next = wal_rec;
 		tde_wal_key_last_rec->end_lsn = wal_rec->start_lsn;
+		tde_wal_key_last_rec->end_tli = wal_rec->start_tli;
 		tde_wal_key_last_rec = wal_rec;
 	}
 

contrib/pg_tde/src/access/pg_tde_xlog_smgr.c

@@ -67,6 +67,7 @@ typedef struct EncryptionStateData
 {
 	char		db_map_path[MAXPGPATH];
 	pg_atomic_uint64 enc_key_lsn;	/* to sync with readers */
+	pg_atomic_uint64 enc_key_tli;	/* to sync with readers */
 } EncryptionStateData;
 
 static EncryptionStateData *EncryptionState = NULL;
@@ -85,6 +86,18 @@ TDEXLogSetEncKeyLsn(XLogRecPtr start_lsn)
 	pg_atomic_write_u64(&EncryptionState->enc_key_lsn, start_lsn);
 }
 
+static TimeLineID
+TDEXLogGetEncKeyTli()
+{
+	return (TimeLineID) pg_atomic_read_u64(&EncryptionState->enc_key_tli);
+}
+
+static void
+TDEXLogSetEncKeyTli(TimeLineID tli)
+{
+	pg_atomic_write_u64(&EncryptionState->enc_key_tli, tli);
+}
+
 static Size TDEXLogEncryptBuffSize(void);
 
 static int	XLOGChooseNumBuffers(void);
@@ -159,6 +172,7 @@ TDEXLogShmemInit(void)
 	}
 
 	pg_atomic_init_u64(&EncryptionState->enc_key_lsn, 0);
+	pg_atomic_init_u64(&EncryptionState->enc_key_tli, 0);
 
 	elog(DEBUG1, "pg_tde: initialized encryption buffer %lu bytes", TDEXLogEncryptStateSize());
 }
@@ -169,6 +183,7 @@ typedef struct EncryptionStateData
 {
 	char		db_map_path[MAXPGPATH];
 	XLogRecPtr	enc_key_lsn;	/* to sync with reader */
+	XLogRecPtr	enc_key_tli;	/* to sync with reader */
 } EncryptionStateData;
 
 static EncryptionStateData EncryptionStateD = {0};
@@ -186,7 +201,19 @@ TDEXLogGetEncKeyLsn()
 static void
 TDEXLogSetEncKeyLsn(XLogRecPtr start_lsn)
 {
-	EncryptionState->enc_key_lsn = EncryptionKey.start_lsn;
+	EncryptionState->enc_key_lsn = start_lsn;
+}
+
+static TimeLineID
+TDEXLogGetEncKeyTli()
+{
+	return (TimeLineID) EncryptionState->enc_key_tli;
+}
+
+static void
+TDEXLogSetEncKeyTli(TimeLineID tli)
+{
+	EncryptionState->enc_key_lsn = tli;
 }
 
 #endif							/* FRONTEND */
@@ -221,6 +248,7 @@ TDEXLogSmgrInitWrite(bool encrypt_xlog)
 	{
 		EncryptionKey = *key;
 		TDEXLogSetEncKeyLsn(EncryptionKey.start_lsn);
+		TDEXLogSetEncKeyTli(EncryptionKey.tli);
 	}
 
 	if (key)
@@ -245,8 +273,8 @@ TDEXLogWriteEncryptedPages(int fd, const void *buf, size_t count, off_t offset,
 #endif
 
 #ifdef TDE_XLOG_DEBUG
-	elog(DEBUG1, "write encrypted WAL, size: %lu, offset: %ld [%lX], seg: %X/%X, key_start_lsn: %X/%X",
-		 count, offset, offset, LSN_FORMAT_ARGS(segno), LSN_FORMAT_ARGS(key->start_lsn));
+	elog(DEBUG1, "write encrypted WAL, size: %lu, offset: %ld [%lX] tli %u, seg: %X/%X, key_start: %u_%X/%X",
+		 count, offset, offset, tli, LSN_FORMAT_ARGS(segno), key->tli, LSN_FORMAT_ARGS(key->start_lsn));
 #endif
 
 	CalcXLogPageIVPrefix(tli, segno, key->base_iv, iv_prefix);
@@ -272,9 +300,11 @@ tdeheap_xlog_seg_write(int fd, const void *buf, size_t count, off_t offset,
 
 		XLogSegNoOffsetToRecPtr(segno, offset, segSize, lsn);
 
-		pg_tde_wal_last_key_set_lsn(lsn, EncryptionState->db_map_path);
+		pg_tde_wal_last_key_set_lsn(lsn, tli, EncryptionState->db_map_path);
 		EncryptionKey.start_lsn = lsn;
+		EncryptionKey.tli = tli;
 		TDEXLogSetEncKeyLsn(lsn);
+		TDEXLogSetEncKeyTli(tli);
 	}
 
 	if (EncryptionKey.type == TDE_KEY_TYPE_WAL_ENCRYPTED)
@@ -293,8 +323,8 @@ tdeheap_xlog_seg_read(int fd, void *buf, size_t count, off_t offset,
 	ssize_t		readsz;
 
 #ifdef TDE_XLOG_DEBUG
-	elog(DEBUG1, "read from a WAL segment, size: %lu offset: %ld [%lX], seg: %X/%X",
-		 count, offset, offset, LSN_FORMAT_ARGS(segno));
+	elog(DEBUG1, "read from a WAL segment, size: %lu offset: %ld [%lX], tli: %u, seg: %X/%X",
+		 count, offset, offset, tli, LSN_FORMAT_ARGS(segno));
 #endif
 
 	readsz = pg_pread(fd, buf, count, offset);
@@ -318,6 +348,8 @@ TDEXLogCryptBuffer(void *buf, size_t count, off_t offset,
 	XLogRecPtr	write_key_lsn;
 	XLogRecPtr	data_start;
 	XLogRecPtr	data_end;
+	KeyTliLsn	data_start_t = {.tli = tli};
+	KeyTliLsn	data_end_t = {.tli = tli};
 
 	if (!keys)
 	{
@@ -330,11 +362,14 @@ TDEXLogCryptBuffer(void *buf, size_t count, off_t offset,
 	if (!XLogRecPtrIsInvalid(write_key_lsn))
 	{
 		WALKeyCacheRec *last_key = pg_tde_get_last_wal_key();
+		KeyTliLsn	last_key_time = {.tli = last_key->start_tli, .lsn = last_key->start_lsn};
+		KeyTliLsn	write_key_time = {.tli = TDEXLogGetEncKeyTli(), .lsn = write_key_lsn};
 
 		Assert(last_key);
 
 		/* write has generated a new key, need to fetch it */
-		if (last_key->start_lsn < write_key_lsn)
+		if (key_tli_lsn_cmp(last_key_time, write_key_time) == -1)
+		// if (last_key->start_lsn < write_key_lsn)
 		{
 			pg_tde_fetch_wal_keys(write_key_lsn);
 
@@ -346,16 +381,22 @@ TDEXLogCryptBuffer(void *buf, size_t count, off_t offset,
 	XLogSegNoOffsetToRecPtr(segno, offset, segSize, data_start);
 	XLogSegNoOffsetToRecPtr(segno, offset + count, segSize, data_end);
 
+	data_start_t.lsn = data_start;
+	data_end_t.lsn = data_end;
+
 	/*
 	 * TODO: this is higly ineffective. We should get rid of linked list and
 	 * search from the last key as this is what the walsender is useing.
 	 */
 	for (WALKeyCacheRec *curr_key = keys; curr_key != NULL; curr_key = curr_key->next)
 	{
+		KeyTliLsn	key_start_t = {.lsn = curr_key->start_lsn, .tli = curr_key->start_tli};
+		KeyTliLsn	key_end_t = {.lsn = curr_key->end_lsn, .tli = curr_key->end_tli};
+
 #ifdef TDE_XLOG_DEBUG
-		elog(DEBUG1, "WAL key %X/%X-%X/%X, encrypted: %s",
-			 LSN_FORMAT_ARGS(curr_key->start_lsn),
-			 LSN_FORMAT_ARGS(curr_key->end_lsn),
+		elog(DEBUG1, "WAL key %u_%X/%X - %u_%X/%X, encrypted: %s",
+			 curr_key->start_tli, LSN_FORMAT_ARGS(curr_key->start_lsn),
+			 curr_key->end_tli, LSN_FORMAT_ARGS(curr_key->end_lsn),
 			 curr_key->key.type == TDE_KEY_TYPE_WAL_ENCRYPTED ? "yes" : "no");
 #endif
 
@@ -366,7 +407,7 @@ TDEXLogCryptBuffer(void *buf, size_t count, off_t offset,
 			 * Check if the key's range overlaps with the buffer's and decypt
 			 * the part that does.
 			 */
-			if (data_start < curr_key->end_lsn && data_end > curr_key->start_lsn)
+			if (key_tli_lsn_cmp(data_start_t, key_end_t) == -1 && key_tli_lsn_cmp(data_end_t, key_start_t) == 1)
 			{
 				char		iv_prefix[16];
 				off_t		dec_off = XLogSegmentOffset(Max(data_start, curr_key->start_lsn), segSize);
@@ -387,8 +428,8 @@ 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 %X/%X",
-					 dec_off, dec_off - offset, dec_sz, LSN_FORMAT_ARGS(curr_key->key->start_lsn));
+				elog(DEBUG1, "decrypt WAL, dec_off: %lu [buff_off %lu] tli %u, sz: %lu | key %u_%X/%X",
+					 dec_off, dec_off - offset, tli, dec_sz, curr_key->key.tli, LSN_FORMAT_ARGS(curr_key->start_lsn));
 #endif
 				pg_tde_stream_crypt(iv_prefix, dec_off, dec_buf, dec_sz, dec_buf,
 									&curr_key->key, &curr_key->crypt_ctx);

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

@@ -25,8 +25,33 @@ typedef struct InternalKey
 	uint32		type;
 
 	XLogRecPtr	start_lsn;
+	TimeLineID	tli;
 } InternalKey;
 
+typedef struct KeyTliLsn
+{
+	TimeLineID	tli;
+	XLogRecPtr	lsn;
+} KeyTliLsn;
+
+static inline int
+key_tli_lsn_cmp(KeyTliLsn t1, KeyTliLsn t2)
+{
+	if (t1.tli < t2.tli)
+		return -1;
+
+	if (t1.tli > t2.tli)
+		return 1;
+
+	if (t1.lsn < t2.lsn)
+		return -1;
+
+	if (t1.lsn > t2.lsn)
+		return 1;
+
+	return 0;
+}
+
 #define MAP_ENTRY_IV_SIZE 16
 #define MAP_ENTRY_AEAD_TAG_SIZE 16
 
@@ -62,6 +87,8 @@ typedef struct WALKeyCacheRec
 {
 	XLogRecPtr	start_lsn;
 	XLogRecPtr	end_lsn;
+	TimeLineID	start_tli;
+	TimeLineID	end_tli;
 
 	InternalKey key;
 	void	   *crypt_ctx;
@@ -73,7 +100,7 @@ extern InternalKey *pg_tde_read_last_wal_key(void);
 extern WALKeyCacheRec *pg_tde_get_last_wal_key(void);
 extern WALKeyCacheRec *pg_tde_fetch_wal_keys(XLogRecPtr start_lsn);
 extern WALKeyCacheRec *pg_tde_get_wal_cache_keys(void);
-extern void pg_tde_wal_last_key_set_lsn(XLogRecPtr lsn, const char *keyfile_path);
+extern void pg_tde_wal_last_key_set_lsn(XLogRecPtr lsn, TimeLineID tli, const char *keyfile_path);
 extern void pg_tde_create_wal_key(InternalKey *rel_key_data, const RelFileLocator *newrlocator, TDEMapEntryType entry_type);
 
 #define PG_TDE_MAP_FILENAME			"%d_keys"

contrib/pg_tde/t/RewindTest.pm

@@ -44,6 +44,7 @@ use PostgreSQL::Test::Cluster;
 use PostgreSQL::Test::RecursiveCopy;
 use PostgreSQL::Test::Utils;
 use Test::More;
+use pgtde;
 
 our @EXPORT = qw(
   $node_primary
@@ -199,7 +200,7 @@ sub create_standby
 	$node_standby =
 	  PostgreSQL::Test::Cluster->new(
 		'standby' . ($extra_name ? "_${extra_name}" : ''));
-	$node_primary->backup('my_backup');
+	PGTDE::backup($node_primary, 'my_backup');
 	$node_standby->init_from_backup($node_primary, 'my_backup');
 	my $connstr_primary = $node_primary->connstr();
 

contrib/pg_tde/t/pgtde.pm

@@ -108,4 +108,17 @@ sub compare_results
 	return compare($expected_filename_with_path, $out_filename_with_path);
 }
 
+sub backup
+{
+	my ($node, $backup_name, %params) = @_;
+	my $backup_dir = $node->backup_dir . '/'. $backup_name;
+
+	mkdir $backup_dir;
+
+	PostgreSQL::Test::RecursiveCopy::copypath($node->data_dir . '/pg_tde',
+		$backup_dir . '/pg_tde');
+
+	$node->backup($backup_name, %params);
+}
+
 1;

contrib/pg_tde/t/standby_source.pl

@@ -34,6 +34,7 @@
 use File::Copy;
 use PostgreSQL::Test::Cluster;
 use RewindTest;
+use pgtde;
 
 my $tmp_folder = PostgreSQL::Test::Utils::tempdir;
 
@@ -57,15 +58,19 @@
 # Set up node B and C, as cascaded standbys
 #
 # A (primary) <--- B (standby) <--- C (standby)
-$node_a->backup('my_backup');
+my $backup_name = 'my_backup';
+
+PGTDE::backup(node_a, $backup_name);
 $node_b = PostgreSQL::Test::Cluster->new('node_b');
-$node_b->init_from_backup($node_a, 'my_backup', has_streaming => 1);
+$node_b->init_from_backup($node_a, $backup_name, has_streaming => 1);
 $node_b->set_standby_mode();
 $node_b->start;
 
-$node_b->backup('my_backup');
+PostgreSQL::Test::RecursiveCopy::copypath($node_b->data_dir . '/pg_tde',
+	$node_b->backup_dir . '/'. $backup_name . '/pg_tde');
+PGTDE::backup(node_b, $backup_name);
 $node_c = PostgreSQL::Test::Cluster->new('node_c');
-$node_c->init_from_backup($node_b, 'my_backup', has_streaming => 1);
+$node_c->init_from_backup($node_b, $backup_name, has_streaming => 1);
 $node_c->set_standby_mode();
 $node_c->start;
 

src/bin/pg_basebackup/pg_basebackup.c

@@ -634,9 +634,6 @@ StartLogStreamer(char *startpos, uint32 timeline, char *sysidentifier,
 	uint32		hi,
 				lo;
 	char		statusdir[MAXPGPATH];
-#ifdef PERCONA_EXT
-	char		tdedir[MAXPGPATH];
-#endif
 
 	param = pg_malloc0(sizeof(logstreamer_param));
 	param->timeline = timeline;
@@ -671,9 +668,13 @@ StartLogStreamer(char *startpos, uint32 timeline, char *sysidentifier,
 			 "pg_xlog" : "pg_wal");
 
 #ifdef PERCONA_EXT
+{
+	char		tdedir[MAXPGPATH];
+
 	snprintf(tdedir, sizeof(tdedir), "%s/%s", basedir, PG_TDE_DATA_DIR);
 	pg_tde_fe_init(tdedir);
 	TDEXLogSmgrInit();
+}
 #endif
 
 	/* Temporary replication slots are only supported in 10 and newer */