open_pg_tde_tdemap.c

#include "postgres.h"

#include <openssl/err.h>
#include <openssl/rand.h>
#include <unistd.h>

#include "access/xlog.h"
#include "access/xlog_internal.h"
#include "access/xloginsert.h"
#include "common/file_perm.h"
#include "miscadmin.h"
#include "storage/fd.h"
#include "utils/builtins.h"
#include "utils/memutils.h"
#include "utils/wait_event.h"

#include "access/open_pg_tde_tdemap.h"
#include "access/open_pg_tde_xlog.h"
#include "catalog/tde_global_space.h"
#include "catalog/tde_principal_key.h"
#include "common/open_pg_tde_utils.h"
#include "encryption/enc_aes.h"
#include "encryption/enc_tde.h"
#include "keyring/keyring_api.h"
#include "open_pg_tde.h"

#ifdef FRONTEND
#include "open_pg_tde_fe.h"
#endif

/* A useful macro when debugging key encryption/decryption */
#ifdef DEBUG
#define ELOG_KEY(_msg, _key)												\
{																			\
	int i;																	\
	char buf[1024];															\
	for (i = 0; i < sizeof(_key.key); i++)					\
		sprintf(buf+i, "%02X", _key.key[i]);				\
	buf[i] = '\0';															\
	elog(INFO, "[%s] INTERNAL KEY => %s", _msg, buf);						\
}
#endif

#define OPEN_PG_TDE_MAP_FILENAME			"%d_keys"

typedef enum
{
	MAP_ENTRY_TYPE_EMPTY = 0,
	MAP_ENTRY_TYPE_KEY = 1,
}			TDEMapEntryType;

typedef struct TDEFileHeader
{
	int32		file_version;
	TDESignedPrincipalKeyInfo signed_key_info;
} TDEFileHeader;

typedef struct TDEMapEntry
{
	uint32		cipher;			/* Part of AAD. Cipher type. We support only
								 * AES_128 and AES_256 for now. */
	Oid			spcOid;			/* Part of AAD */
	RelFileNumber relNumber;	/* Part of AAD */
	uint32		type;			/* Part of AAD */

	/*
	 * key_base_iv is the per-relation IV/tweak base for the data files. It is
	 * placed here, before entry_iv, so that it falls within the AEAD
	 * additional authenticated data (offsetof(TDEMapEntry, entry_iv)). In the
	 * v5 format it lived after aead_tag and was therefore unauthenticated,
	 * which let an actor with write access to the key file alter every block
	 * IV undetected. See the v5 migration reader below.
	 */
	uint8		key_base_iv[INTERNAL_KEY_IV_LEN];

	/*
	 * IV and tag used when encrypting the key itself
	 *
	 * TODO: should we extend MAP_ENTRY_IV_SIZE to 192(?) bit and add an
	 * iv_size filed?
	 */
	unsigned char entry_iv[MAP_ENTRY_IV_SIZE];
	unsigned char aead_tag[MAP_ENTRY_AEAD_TAG_SIZE];

	uint8		encrypted_key_data[INTERNAL_KEY_MAX_LEN];
} TDEMapEntry;

static void open_pg_tde_set_db_file_path(Oid dbOid, char *path);
static bool open_pg_tde_find_map_entry(const RelFileLocator *rlocator, char *db_map_path, TDEMapEntry *map_entry);
static InternalKey *tde_decrypt_rel_key(const TDEPrincipalKey *principal_key, TDEMapEntry *map_entry);
static int	open_pg_tde_open_file_basic(const char *tde_filename, int fileFlags, bool ignore_missing);
static int	open_pg_tde_open_file_read(const char *tde_filename, bool ignore_missing, off_t *curr_pos);
static void open_pg_tde_file_header_read(const char *tde_filename, int fd, TDEFileHeader *fheader, off_t *bytes_read);
static bool open_pg_tde_read_one_map_entry(int fd, TDEMapEntry *map_entry, off_t *offset);
static void open_pg_tde_write_one_map_entry(int fd, const TDEMapEntry *map_entry, off_t *offset, const char *db_map_path);
static int	open_pg_tde_file_header_write(const char *tde_filename, int fd, const TDESignedPrincipalKeyInfo *signed_key_info, off_t *bytes_written);
static void open_pg_tde_initialize_map_entry(TDEMapEntry *map_entry, const TDEPrincipalKey *principal_key, const RelFileLocator *rlocator, const InternalKey *rel_key_data);
static int	open_pg_tde_open_file_write(const char *tde_filename, const TDESignedPrincipalKeyInfo *signed_key_info, bool truncate, off_t *curr_pos);

/*
 * Saves an internal key for the given relation. If replace_existing is false,
 * the function will not overwrite an existing key for the relation, but will
 * instead do nothing.
 */
void
open_pg_tde_save_smgr_key(RelFileLocator rel,
						  const InternalKey *rel_key_data,
						  bool replace_existing)
{
	char		file_path[MAXPGPATH];
	bool		entry_already_exists;
	int			fd;
	LWLock	   *lock_pk = tde_lwlock_enc_keys();
	TDEPrincipalKey *principal_key;
	off_t		scan_offset;
	TDESignedPrincipalKeyInfo signed_key_info;
	off_t		write_offset;

	LWLockAcquire(lock_pk, LW_EXCLUSIVE);
	principal_key = GetPrincipalKey(rel.dbOid, LW_EXCLUSIVE);
	if (principal_key == NULL)
	{
		ereport(ERROR,
				errmsg("principal key not configured"),
				errhint("Use open_pg_tde_set_key_using_database_key_provider() or open_pg_tde_set_key_using_global_key_provider() to configure one."));
	}

	open_pg_tde_set_db_file_path(rel.dbOid, file_path);
	open_pg_tde_sign_principal_key_info(&signed_key_info, principal_key);

	fd = open_pg_tde_open_file_write(file_path, &signed_key_info, false, &scan_offset);

	/*
	 * Look for an existing entry for the relation, also keep track of the
	 * first free entry we find to be reused for the new key we're saving.
	 */
	entry_already_exists = false;
	write_offset = 0;
	while (1)
	{
		TDEMapEntry entry;
		off_t		entry_offset = scan_offset;

		if (!open_pg_tde_read_one_map_entry(fd, &entry, &scan_offset))
		{
			/*
			 * If we're through the whole file without having found an empty
			 * entry to overwrite, write at the end.
			 */
			if (write_offset == 0)
				write_offset = entry_offset;
			break;
		}

		if (entry.spcOid == rel.spcOid &&
			entry.relNumber == rel.relNumber)
		{
			entry_already_exists = true;
			write_offset = entry_offset;
			break;
		}

		if (write_offset == 0 && entry.type == MAP_ENTRY_TYPE_EMPTY)
			write_offset = entry_offset;
	}

	if (!entry_already_exists || replace_existing)
	{
		TDEMapEntry write_entry;

		open_pg_tde_initialize_map_entry(&write_entry, principal_key, &rel, rel_key_data);
		open_pg_tde_write_one_map_entry(fd, &write_entry, &write_offset, file_path);
	}

#ifdef FRONTEND
	pfree(principal_key);
#endif

	CloseTransientFile(fd);
	LWLockRelease(lock_pk);
}

#ifndef FRONTEND
const char *
tde_sprint_key(InternalKey *k)
{
	static char buf[256];
	int			i;

	for (i = 0; i < sizeof(k->key); i++)
		sprintf(buf + i, "%02X", k->key[i]);

	return buf;
}

/*
 * Deletes the key file for a given database.
 */
void
open_pg_tde_delete_tde_files(Oid dbOid)
{
	char		db_map_path[MAXPGPATH];

	open_pg_tde_set_db_file_path(dbOid, db_map_path);

	/* Remove file without emitting any error */
	PathNameDeleteTemporaryFile(db_map_path, false);
}

void
open_pg_tde_save_principal_key_redo(const TDESignedPrincipalKeyInfo *signed_key_info)
{
	int			map_fd;
	off_t		curr_pos;
	char		db_map_path[MAXPGPATH];

	open_pg_tde_set_db_file_path(signed_key_info->data.databaseId, db_map_path);

	LWLockAcquire(tde_lwlock_enc_keys(), LW_EXCLUSIVE);

	map_fd = open_pg_tde_open_file_write(db_map_path, signed_key_info, false, &curr_pos);
	CloseTransientFile(map_fd);

	LWLockRelease(tde_lwlock_enc_keys());
}

/*
 * Creates the key file and saves the principal key information.
 *
 * If the file pre-exist, it truncates the file before adding principal key
 * information.
 *
 * The caller must have an EXCLUSIVE LOCK on the files before calling this function.
 *
 * write_xlog: if true, the function will write an XLOG record about the
 * principal key addition. We may want to skip this during server recovery/startup
 * or in some other cases when WAL writes are not allowed.
 */
void
open_pg_tde_save_principal_key(const TDEPrincipalKey *principal_key, bool write_xlog)
{
	int			map_fd;
	off_t		curr_pos = 0;
	char		db_map_path[MAXPGPATH];
	TDESignedPrincipalKeyInfo signed_key_Info;

	open_pg_tde_set_db_file_path(principal_key->keyInfo.databaseId, db_map_path);

	open_pg_tde_sign_principal_key_info(&signed_key_Info, principal_key);

	if (write_xlog)
	{
		XLogBeginInsert();
		XLogRegisterData((char *) &signed_key_Info, sizeof(TDESignedPrincipalKeyInfo));
		XLogInsert(RM_TDERMGR_ID, XLOG_TDE_ADD_PRINCIPAL_KEY);
	}

	map_fd = open_pg_tde_open_file_write(db_map_path, &signed_key_Info, true, &curr_pos);
	CloseTransientFile(map_fd);
}

/*
 * Mark relation map entry as free and overwrite the key
 *
 * This fucntion is called by the open_pg_tde SMGR when storage is unlinked on
 * transaction commit/abort.
 */
void
open_pg_tde_free_key_map_entry(const RelFileLocator rlocator)
{
	char		db_map_path[MAXPGPATH];
	File		map_fd;
	off_t		curr_pos = 0;

	open_pg_tde_set_db_file_path(rlocator.dbOid, db_map_path);

	LWLockAcquire(tde_lwlock_enc_keys(), LW_EXCLUSIVE);

	/* Open and validate file for basic correctness. */
	map_fd = open_pg_tde_open_file_write(db_map_path, NULL, false, &curr_pos);

	while (1)
	{
		TDEMapEntry map_entry;
		off_t		prev_pos = curr_pos;

		if (!open_pg_tde_read_one_map_entry(map_fd, &map_entry, &curr_pos))
			break;

		if (map_entry.type == MAP_ENTRY_TYPE_KEY && map_entry.spcOid == rlocator.spcOid && map_entry.relNumber == rlocator.relNumber)
		{
			TDEMapEntry empty_map_entry = {
				.type = MAP_ENTRY_TYPE_EMPTY,
			};

			open_pg_tde_write_one_map_entry(map_fd, &empty_map_entry, &prev_pos, db_map_path);
			break;
		}
	}

	CloseTransientFile(map_fd);

	LWLockRelease(tde_lwlock_enc_keys());
}

/*
 * Rotate keys and generates the WAL record for it.
 */
void
open_pg_tde_perform_rotate_key(const TDEPrincipalKey *principal_key, const TDEPrincipalKey *new_principal_key, bool write_xlog)
{
	TDESignedPrincipalKeyInfo new_signed_key_info;
	off_t		old_curr_pos,
				new_curr_pos;
	int			old_fd,
				new_fd;
	char		old_path[MAXPGPATH],
				new_path[MAXPGPATH];

	/* This function cannot be used to rotate the server key. */
	Assert(principal_key);
	Assert(principal_key->keyInfo.databaseId != GLOBAL_DATA_TDE_OID);

	open_pg_tde_sign_principal_key_info(&new_signed_key_info, new_principal_key);

	open_pg_tde_set_db_file_path(principal_key->keyInfo.databaseId, old_path);
	snprintf(new_path, MAXPGPATH, "%s.r", old_path);

	old_fd = open_pg_tde_open_file_read(old_path, false, &old_curr_pos);
	new_fd = open_pg_tde_open_file_write(new_path, &new_signed_key_info, true, &new_curr_pos);

	/* Read all entries until EOF */
	while (1)
	{
		InternalKey *rel_key_data;
		TDEMapEntry read_map_entry,
					write_map_entry;
		RelFileLocator rloc;

		if (!open_pg_tde_read_one_map_entry(old_fd, &read_map_entry, &old_curr_pos))
			break;

		if (read_map_entry.type == MAP_ENTRY_TYPE_EMPTY)
			continue;

		rloc.spcOid = read_map_entry.spcOid;
		rloc.dbOid = principal_key->keyInfo.databaseId;
		rloc.relNumber = read_map_entry.relNumber;

		/* Decrypt and re-encrypt key */
		rel_key_data = tde_decrypt_rel_key(principal_key, &read_map_entry);
		open_pg_tde_initialize_map_entry(&write_map_entry, new_principal_key, &rloc, rel_key_data);

		open_pg_tde_write_one_map_entry(new_fd, &write_map_entry, &new_curr_pos, new_path);

		pfree(rel_key_data);
	}

	CloseTransientFile(old_fd);
	CloseTransientFile(new_fd);

	/*
	 * Do the final step - replace the current _keys with the file with new
	 * data
	 */
	durable_rename(new_path, old_path, ERROR);

	/*
	 * We do WAL writes past the event ("the write behind logging") rather
	 * than before ("the write ahead") because we need logging here only for
	 * replication purposes. The rotation results in data written and fsynced
	 * to disk. Which in most cases would happen way before it's written to
	 * the WAL disk file. As WAL will be flushed at the end of the
	 * transaction, on its commit, hence after this function returns (there is
	 * also a bg writer, but the commit is what is guaranteed). And it makes
	 * sense to replicate the event only after its effect has been
	 * successfully applied to the source.
	 */
	if (write_xlog)
	{
		XLogPrincipalKeyRotate xlrec;

		xlrec.databaseId = new_principal_key->keyInfo.databaseId;
		xlrec.keyringId = new_principal_key->keyInfo.keyringId;
		memcpy(xlrec.keyName, new_principal_key->keyInfo.name, sizeof(new_principal_key->keyInfo.name));

		XLogBeginInsert();
		XLogRegisterData((char *) &xlrec, sizeof(XLogPrincipalKeyRotate));
		XLogInsert(RM_TDERMGR_ID, XLOG_TDE_ROTATE_PRINCIPAL_KEY);
	}
}

void
open_pg_tde_delete_principal_key_redo(Oid dbOid)
{
	char		path[MAXPGPATH];

	open_pg_tde_set_db_file_path(dbOid, path);

	LWLockAcquire(tde_lwlock_enc_keys(), LW_EXCLUSIVE);
	durable_unlink(path, WARNING);
	LWLockRelease(tde_lwlock_enc_keys());
}

/*
 * Deletes the principal key for the database. This fucntion checks if key map
 * file has any entries, and if not, it removes the file. Otherwise raises an error.
 */
void
open_pg_tde_delete_principal_key(Oid dbOid)
{
	char		path[MAXPGPATH];

	Assert(LWLockHeldByMeInMode(tde_lwlock_enc_keys(), LW_EXCLUSIVE));
	Assert(open_pg_tde_count_encryption_keys(dbOid, InvalidOid) == 0);

	open_pg_tde_set_db_file_path(dbOid, path);

	XLogBeginInsert();
	XLogRegisterData((char *) &dbOid, sizeof(Oid));
	XLogInsert(RM_TDERMGR_ID, XLOG_TDE_DELETE_PRINCIPAL_KEY);

	/* Remove whole key map file */
	durable_unlink(path, ERROR);
}

#endif							/* !FRONTEND */

static void
open_pg_tde_set_db_file_path(Oid dbOid, char *path)
{
	snprintf(path, MAXPGPATH, "%s/" OPEN_PG_TDE_MAP_FILENAME, open_pg_tde_get_data_dir(), dbOid);
}

void
open_pg_tde_sign_principal_key_info(TDESignedPrincipalKeyInfo *signed_key_info, const TDEPrincipalKey *principal_key)
{
	signed_key_info->data = principal_key->keyInfo;

	if (!RAND_bytes(signed_key_info->sign_iv, MAP_ENTRY_IV_SIZE))
		ereport(ERROR,
				errcode(ERRCODE_INTERNAL_ERROR),
				errmsg("could not generate iv for key map: %s", ERR_error_string(ERR_get_error(), NULL)));

	AesGcmEncrypt(principal_key->keyData, principal_key->keyLength,
				  signed_key_info->sign_iv, MAP_ENTRY_IV_SIZE,
				  (unsigned char *) &signed_key_info->data, sizeof(signed_key_info->data),
				  (unsigned char *) "", 0,
				  (unsigned char *) "",
				  signed_key_info->aead_tag, MAP_ENTRY_AEAD_TAG_SIZE);
}

static void
open_pg_tde_initialize_map_entry(TDEMapEntry *map_entry, const TDEPrincipalKey *principal_key, const RelFileLocator *rlocator, const InternalKey *rel_key_data)
{
	map_entry->spcOid = rlocator->spcOid;
	map_entry->relNumber = rlocator->relNumber;
	map_entry->type = MAP_ENTRY_TYPE_KEY;
	memcpy(map_entry->key_base_iv, rel_key_data->base_iv, INTERNAL_KEY_IV_LEN);

	Assert(rel_key_data->key_len == 16 || rel_key_data->key_len == 32 ||
		   rel_key_data->key_len == 64);
	map_entry->cipher = rel_key_data->cipher;	/* recorded so reads dispatch
												 * on the cipher chosen at
												 * creation */

	if (!RAND_bytes(map_entry->entry_iv, MAP_ENTRY_IV_SIZE))
		ereport(ERROR,
				errcode(ERRCODE_INTERNAL_ERROR),
				errmsg("could not generate iv for key map: %s", ERR_error_string(ERR_get_error(), NULL)));

	AesGcmEncrypt(principal_key->keyData, principal_key->keyLength,
				  map_entry->entry_iv, MAP_ENTRY_IV_SIZE,
				  (unsigned char *) map_entry, offsetof(TDEMapEntry, entry_iv),
				  rel_key_data->key, rel_key_data->key_len,
				  map_entry->encrypted_key_data,
				  map_entry->aead_tag, MAP_ENTRY_AEAD_TAG_SIZE);
}

static void
open_pg_tde_write_one_map_entry(int fd, const TDEMapEntry *map_entry, off_t *offset, const char *db_map_path)
{
	int			bytes_written = 0;

	bytes_written = pg_pwrite(fd, map_entry, sizeof(TDEMapEntry), *offset);

	if (bytes_written != sizeof(TDEMapEntry))
	{
		ereport(ERROR,
				errcode_for_file_access(),
				errmsg("could not write tde map file \"%s\": %m", db_map_path));
	}
	if (pg_fsync(fd) != 0)
	{
		ereport(data_sync_elevel(ERROR),
				errcode_for_file_access(),
				errmsg("could not fsync file \"%s\": %m", db_map_path));
	}

	*offset += bytes_written;
}

/*
 * Returns true if we find a valid match; e.g. type is not set to
 * MAP_ENTRY_TYPE_EMPTY and the relNumber and spcOid matches the one provided
 * in rlocator.
 */
static bool
open_pg_tde_find_map_entry(const RelFileLocator *rlocator, char *db_map_path, TDEMapEntry *map_entry)
{
	File		map_fd;
	off_t		curr_pos = 0;
	bool		found = false;

	Assert(rlocator != NULL);

	map_fd = open_pg_tde_open_file_read(db_map_path, false, &curr_pos);

	while (open_pg_tde_read_one_map_entry(map_fd, map_entry, &curr_pos))
	{
		if (map_entry->type == MAP_ENTRY_TYPE_KEY && map_entry->spcOid == rlocator->spcOid && map_entry->relNumber == rlocator->relNumber)
		{
			found = true;
			break;
		}
	}

	CloseTransientFile(map_fd);

	return found;
}

/*
 * Counts number of encryption keys in a key file.
 *
 * Does not check if objects actually exist but just that they have keys in
 * the key file.
 *
 * Works even if the database has no key file.
 */
int
open_pg_tde_count_encryption_keys(Oid dbOid, Oid spcOid)
{
	char		db_map_path[MAXPGPATH];
	File		map_fd;
	off_t		curr_pos = 0;
	TDEMapEntry map_entry;
	int			count = 0;

	Assert(LWLockHeldByMeInMode(tde_lwlock_enc_keys(), LW_SHARED) || LWLockHeldByMeInMode(tde_lwlock_enc_keys(), LW_EXCLUSIVE));

	open_pg_tde_set_db_file_path(dbOid, db_map_path);

	map_fd = open_pg_tde_open_file_read(db_map_path, true, &curr_pos);
	if (map_fd < 0)
		return count;

	while (open_pg_tde_read_one_map_entry(map_fd, &map_entry, &curr_pos))
	{
		if (map_entry.type == MAP_ENTRY_TYPE_KEY &&
			(spcOid == InvalidOid || map_entry.spcOid == spcOid))
			count++;
	}

	CloseTransientFile(map_fd);

	return count;
}

bool
open_pg_tde_verify_principal_key_info(TDESignedPrincipalKeyInfo *signed_key_info, const KeyData *principal_key_data)
{
	return AesGcmDecrypt(principal_key_data->data, principal_key_data->len,
						 signed_key_info->sign_iv, MAP_ENTRY_IV_SIZE,
						 (unsigned char *) &signed_key_info->data, sizeof(signed_key_info->data),
						 (unsigned char *) "", 0,
						 (unsigned char *) "",
						 signed_key_info->aead_tag, MAP_ENTRY_AEAD_TAG_SIZE);
}

static InternalKey *
tde_decrypt_rel_key(const TDEPrincipalKey *principal_key, TDEMapEntry *map_entry)
{
	InternalKey *key = palloc_object(InternalKey);
	uint32		key_len = open_pg_tde_cipher_key_length(map_entry->cipher);

	Assert(principal_key);

	if (!AesGcmDecrypt(principal_key->keyData, principal_key->keyLength,
					   map_entry->entry_iv, MAP_ENTRY_IV_SIZE,
					   (unsigned char *) map_entry, offsetof(TDEMapEntry, entry_iv),
					   map_entry->encrypted_key_data, key_len,
					   key->key,
					   map_entry->aead_tag, MAP_ENTRY_AEAD_TAG_SIZE))
		ereport(ERROR,
				errmsg("Failed to decrypt key, incorrect principal key or corrupted key file"));

	memcpy(key->base_iv, map_entry->key_base_iv, INTERNAL_KEY_IV_LEN);
	key->key_len = key_len;
	key->cipher = map_entry->cipher;

	return key;
}

/*
 * Open a TDE file:
 *
 * Returns the file descriptor in case of a success. Otherwise, error
 * is raised except when ignore_missing is true and the file does not exit.
 */
static int
open_pg_tde_open_file_basic(const char *tde_filename, int fileFlags, bool ignore_missing)
{
	int			fd;

	fd = OpenTransientFile(tde_filename, fileFlags);
	if (fd < 0 && !(errno == ENOENT && ignore_missing == true))
	{
		ereport(ERROR,
				errcode_for_file_access(),
				errmsg("could not open tde file \"%s\": %m", tde_filename));
	}

	return fd;
}

/*
 * Open for read and Validate File Header:
 * 		header: {Format Version, Principal Key Name}
 *
 * Returns the file descriptor in case of a success. Otherwise, error
 * is raised.
 */
static int
open_pg_tde_open_file_read(const char *tde_filename, bool ignore_missing, off_t *curr_pos)
{
	int			fd;
	TDEFileHeader fheader;
	off_t		bytes_read = 0;

	Assert(LWLockHeldByMeInMode(tde_lwlock_enc_keys(), LW_SHARED) || LWLockHeldByMeInMode(tde_lwlock_enc_keys(), LW_EXCLUSIVE));

	fd = open_pg_tde_open_file_basic(tde_filename, O_RDONLY | PG_BINARY, ignore_missing);
	if (ignore_missing && fd < 0)
		return fd;

	open_pg_tde_file_header_read(tde_filename, fd, &fheader, &bytes_read);
	if (bytes_read > 0 && fheader.file_version != OPEN_PG_TDE_SMGR_FILE_MAGIC)
		ereport(FATAL,
				errcode_for_file_access(),
				errmsg("key file \"%s\" has wrong version: %m", tde_filename));

	*curr_pos = bytes_read;

	return fd;
}

/*
 * Open for write and Validate File Header:
 * 		header: {Format Version, Principal Key Name}
 *
 * Returns the file descriptor in case of a success. Otherwise, error
 * is raised.
 */
static int
open_pg_tde_open_file_write(const char *tde_filename, const TDESignedPrincipalKeyInfo *signed_key_info, bool truncate, off_t *curr_pos)
{
	int			fd;
	TDEFileHeader fheader;
	off_t		bytes_read = 0;
	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 = open_pg_tde_open_file_basic(tde_filename, file_flags, false);

	open_pg_tde_file_header_read(tde_filename, fd, &fheader, &bytes_read);
	if (bytes_read > 0 && fheader.file_version != OPEN_PG_TDE_SMGR_FILE_MAGIC)
		ereport(FATAL,
				errcode_for_file_access(),
				errmsg("key file \"%s\" has wrong version: %m", tde_filename));

	/* In case it's a new file, let's add the header now. */
	if (bytes_read == 0 && signed_key_info)
		open_pg_tde_file_header_write(tde_filename, fd, signed_key_info, &bytes_written);

	*curr_pos = bytes_read + bytes_written;
	return fd;
}

/*
 * Read TDE file header from a TDE file and fill in the fheader data structure.
 */
static void
open_pg_tde_file_header_read(const char *tde_filename, int fd, TDEFileHeader *fheader, off_t *bytes_read)
{
	Assert(fheader);

	*bytes_read = pg_pread(fd, fheader, sizeof(TDEFileHeader), 0);

	/* File is empty */
	if (*bytes_read == 0)
		return;

	if (*bytes_read != sizeof(TDEFileHeader))
	{
		ereport(FATAL,
				errcode_for_file_access(),
				errmsg("TDE map file \"%s\" is corrupted: %m", tde_filename));
	}
}

/*
 * Write TDE file header to a TDE file.
 */
static int
open_pg_tde_file_header_write(const char *tde_filename, int fd, const TDESignedPrincipalKeyInfo *signed_key_info, off_t *bytes_written)
{
	TDEFileHeader fheader;

	Assert(signed_key_info);

	fheader.file_version = OPEN_PG_TDE_SMGR_FILE_MAGIC;
	fheader.signed_key_info = *signed_key_info;
	*bytes_written = pg_pwrite(fd, &fheader, sizeof(TDEFileHeader), 0);

	if (*bytes_written != sizeof(TDEFileHeader))
	{
		ereport(ERROR,
				errcode_for_file_access(),
				errmsg("could not write tde file \"%s\": %m", tde_filename));
	}

	if (pg_fsync(fd) != 0)
	{
		ereport(data_sync_elevel(ERROR),
				errcode_for_file_access(),
				errmsg("could not fsync file \"%s\": %m", tde_filename));
	}

	ereport(DEBUG2, errmsg("Wrote the header to %s", tde_filename));

	return fd;
}

/*
 * Returns true if a map entry if found or false if we have reached the end of
 * the file.
 */
static bool
open_pg_tde_read_one_map_entry(int map_file, TDEMapEntry *map_entry, off_t *offset)
{
	off_t		bytes_read = 0;

	Assert(map_entry);
	Assert(offset);

	bytes_read = pg_pread(map_file, map_entry, sizeof(TDEMapEntry), *offset);

	/* We've reached the end of the file. */
	if (bytes_read != sizeof(TDEMapEntry))
		return false;

	*offset += bytes_read;

	return true;
}

/*
 * Get the principal key from the key file. The caller must hold
 * a LW_SHARED or higher lock on files before calling this function.
 */
TDESignedPrincipalKeyInfo *
open_pg_tde_get_principal_key_info(Oid dbOid)
{
	char		db_map_path[MAXPGPATH];
	int			fd;
	TDEFileHeader fheader;
	TDESignedPrincipalKeyInfo *signed_key_info = NULL;
	off_t		bytes_read = 0;

	open_pg_tde_set_db_file_path(dbOid, db_map_path);

	/*
	 * Ensuring that we always open the file in binary mode. The caller must
	 * specify other flags for reading, writing or creating the file.
	 */
	fd = open_pg_tde_open_file_basic(db_map_path, O_RDONLY, true);

	/* The file does not exist. */
	if (fd < 0)
		return NULL;

	open_pg_tde_file_header_read(db_map_path, fd, &fheader, &bytes_read);

	if (bytes_read > 0 &&
		FILEMAGIC_TYPE(fheader.file_version) != FILEMAGIC_TYPE(OPEN_PG_TDE_SMGR_FILE_MAGIC))
	{
		ereport(FATAL,
				errcode_for_file_access(),
				errmsg("key file \"%s\" is corrupted or has wrong version: %m", db_map_path),
				errdetail("Getting principal key from the file."));
	}

	CloseTransientFile(fd);

	/*
	 * It's not a new file. So we can copy the principal key info from the
	 * header
	 */
	if (bytes_read > 0)
	{
		signed_key_info = palloc_object(TDESignedPrincipalKeyInfo);
		*signed_key_info = fheader.signed_key_info;
	}

	return signed_key_info;
}

/*
 * Figures out whether a relation is encrypted or not, but without trying to
 * decrypt the key if it is.
 */
bool
open_pg_tde_has_smgr_key(RelFileLocator rel)
{
	bool		result;
	TDEMapEntry map_entry;
	char		db_map_path[MAXPGPATH];

	Assert(rel.relNumber != InvalidRelFileNumber);

	open_pg_tde_set_db_file_path(rel.dbOid, db_map_path);

	if (access(db_map_path, F_OK) == -1)
		return false;

	LWLockAcquire(tde_lwlock_enc_keys(), LW_SHARED);

	result = open_pg_tde_find_map_entry(&rel, db_map_path, &map_entry);

	LWLockRelease(tde_lwlock_enc_keys());
	return result;
}

/*
 * Reads the map entry of the relation and decrypts the key.
 */
InternalKey *
open_pg_tde_get_smgr_key(RelFileLocator rel)
{
	TDEMapEntry map_entry;
	TDEPrincipalKey *principal_key;
	LWLock	   *lock_pk = tde_lwlock_enc_keys();
	char		db_map_path[MAXPGPATH];
	InternalKey *rel_key;

	Assert(rel.relNumber != InvalidRelFileNumber);

	open_pg_tde_set_db_file_path(rel.dbOid, db_map_path);

	if (access(db_map_path, F_OK) == -1)
		return NULL;

	LWLockAcquire(lock_pk, LW_SHARED);

	if (!open_pg_tde_find_map_entry(&rel, db_map_path, &map_entry))
	{
		LWLockRelease(lock_pk);
		return NULL;
	}

	/*
	 * Get/generate a principal key, create the key for relation and get the
	 * encrypted key with bytes to write
	 *
	 * We should hold the lock until the internal key is loaded to be sure the
	 * retrieved key was encrypted with the obtained principal key. Otherwise,
	 * the next may happen: - GetPrincipalKey returns key "PKey_1". - Some
	 * other process rotates the Principal key and re-encrypt an Internal key
	 * with "PKey_2". - We read the Internal key and decrypt it with "PKey_1"
	 * (that's what we've got). As the result we return an invalid Internal
	 * key.
	 */
	principal_key = GetPrincipalKey(rel.dbOid, LW_SHARED);
	if (principal_key == NULL)
		ereport(ERROR,
				errmsg("principal key not configured"),
				errhint("Use open_pg_tde_set_key_using_database_key_provider() or open_pg_tde_set_key_using_global_key_provider() to configure one."));
	rel_key = tde_decrypt_rel_key(principal_key, &map_entry);

	LWLockRelease(lock_pk);

	if (rel_key->key_len <= MAX_KEY_DATA_SIZE &&
		principal_key->keyLength != rel_key->key_len)
	{
		ereport(LOG,
				errmsg("length \"%u\" of principal key \"%s\" does not match the length \"%d\" of the internal key", principal_key->keyLength, principal_key->keyInfo.name, rel_key->key_len),
				errhint("Create a new principal key and set it instead of the current one."));
	}

#ifdef FRONTEND
	/* There is no principal key cache in the frontend. */
	pfree(principal_key);
#endif

	return rel_key;
}


#ifndef FRONTEND

/*****************************************
 * Functions for migrating old smgr keys into a new format file.
 *****************************************/

/*
 * A version-specific migration routine. It reads an entry of the specific
 * version from the given fd and offset, and transforms it into
 * TDEMapEntry (current version)
 */
typedef bool (*MapFromDiskEntry) (int fd, off_t *entry_offset, const TDEPrincipalKey *principal_key, TDEMapEntry *out);

typedef struct TDEMapEntryV3
{
	Oid			spcOid;			/* Part of AAD */
	RelFileNumber relNumber;	/* Part of AAD */
	uint32		type;			/* Part of AAD */
	uint32		_unused1;		/* Part of AAD */

	uint8		encrypted_key_data[16];
	uint8		key_base_iv[16];

	uint32		_unused2;		/* Will be 1 in existing files entries. */
	uint32		_unused3;
	uint64		_unused4;		/* Will be 0 in existing files entries. */

	/* IV and tag used when encrypting the key itself */
	unsigned char entry_iv[16];
	unsigned char aead_tag[16];
} TDEMapEntryV3;

static bool
read_one_map_entry_v3(int fd, TDEMapEntryV3 *entry, off_t *offset)
{
	off_t		bytes_read = 0;

	Assert(entry);
	Assert(offset);

	bytes_read = pg_pread(fd, entry, sizeof(TDEMapEntryV3), *offset);

	/* We've reached the end of the file. */
	if (bytes_read != sizeof(TDEMapEntryV3))
		return false;

	*offset += bytes_read;

	return true;
}

static void
ikey_from_map_entry_v3(TDEMapEntryV3 *entry, const TDEPrincipalKey *principal_key, InternalKey *out)
{
	out->key_len = sizeof(entry->encrypted_key_data);
	/* v3 entries always hold a 16-byte key, i.e. AES-128. */
	out->cipher = CIPHER_AES_128;

	memcpy(out->base_iv, entry->key_base_iv, sizeof(entry->key_base_iv));
	if (!AesGcmDecrypt(principal_key->keyData, principal_key->keyLength,
					   entry->entry_iv, sizeof(entry->entry_iv),
					   (unsigned char *) entry, offsetof(TDEMapEntryV3, encrypted_key_data),
					   entry->encrypted_key_data, out->key_len,
					   out->key,
					   entry->aead_tag, sizeof(entry->aead_tag)))
		ereport(ERROR,
				errmsg("failed to decrypt key, incorrect principal key or corrupted key file"));
}

static bool
map_from_disk_entry_v3(int fd, off_t *entry_offset, const TDEPrincipalKey *principal_key, TDEMapEntry *out)
{
	TDEMapEntryV3 disk_entry;
	InternalKey key;
	RelFileLocator rloc;

	if (!read_one_map_entry_v3(fd, &disk_entry, entry_offset))
		return false;

	/* Decrypting empty entries would fail, so just return an empty entry. */
	if (disk_entry.type == MAP_ENTRY_TYPE_EMPTY)
	{
		memset(out, 0, sizeof(TDEMapEntry));
		return true;
	}

	ikey_from_map_entry_v3(&disk_entry, principal_key, &key);

	rloc.spcOid = disk_entry.spcOid;
	rloc.dbOid = principal_key->keyInfo.databaseId;
	rloc.relNumber = disk_entry.relNumber;

	open_pg_tde_initialize_map_entry(out, principal_key, &rloc, &key);

	return true;
}

/*
 * Version 4 key map entry: identical layout to the current entry except the
 * key buffer, which was 32 bytes before AES-256-XTS added a 64-byte key.
 * Retained to migrate version-4 files to the current format.
 */
typedef struct TDEMapEntryV4
{
	uint32		cipher;
	Oid			spcOid;
	RelFileNumber relNumber;
	uint32		type;
	unsigned char entry_iv[MAP_ENTRY_IV_SIZE];
	unsigned char aead_tag[MAP_ENTRY_AEAD_TAG_SIZE];
	uint8		key_base_iv[INTERNAL_KEY_IV_LEN];
	uint8		encrypted_key_data[32];
}			TDEMapEntryV4;

static bool
read_one_map_entry_v4(int fd, TDEMapEntryV4 * entry, off_t *offset)
{
	int			bytes_read = pg_pread(fd, entry, sizeof(TDEMapEntryV4), *offset);

	if (bytes_read == 0)
		return false;
	if (bytes_read != sizeof(TDEMapEntryV4))
		ereport(ERROR,
				errmsg("keys migration: could not read map entry"));
	*offset += bytes_read;
	return true;
}

static void
ikey_from_map_entry_v4(TDEMapEntryV4 * entry, const TDEPrincipalKey *principal_key, InternalKey *out)
{
	out->key_len = open_pg_tde_cipher_key_length(entry->cipher);
	out->cipher = entry->cipher;

	memcpy(out->base_iv, entry->key_base_iv, sizeof(entry->key_base_iv));
	if (!AesGcmDecrypt(principal_key->keyData, principal_key->keyLength,
					   entry->entry_iv, sizeof(entry->entry_iv),
					   (unsigned char *) entry, offsetof(TDEMapEntryV4, entry_iv),
					   entry->encrypted_key_data, out->key_len,
					   out->key,
					   entry->aead_tag, sizeof(entry->aead_tag)))
		ereport(ERROR,
				errmsg("failed to decrypt key, incorrect principal key or corrupted key file"));
}

static bool
map_from_disk_entry_v4(int fd, off_t *entry_offset, const TDEPrincipalKey *principal_key, TDEMapEntry *out)
{
	TDEMapEntryV4 disk_entry;
	InternalKey key;
	RelFileLocator rloc;

	if (!read_one_map_entry_v4(fd, &disk_entry, entry_offset))
		return false;

	/* Decrypting empty entries would fail, so just return an empty entry. */
	if (disk_entry.type == MAP_ENTRY_TYPE_EMPTY)
	{
		memset(out, 0, sizeof(TDEMapEntry));
		return true;
	}

	ikey_from_map_entry_v4(&disk_entry, principal_key, &key);

	rloc.spcOid = disk_entry.spcOid;
	rloc.dbOid = principal_key->keyInfo.databaseId;
	rloc.relNumber = disk_entry.relNumber;

	open_pg_tde_initialize_map_entry(out, principal_key, &rloc, &key);

	return true;
}

/*
 * Version 5 key map entry: same fields as the current entry, but key_base_iv is
 * stored after the AEAD tag, so it was not covered by the AEAD additional
 * authenticated data. Retained to migrate version-5 files to the current
 * format, which authenticates key_base_iv.
 */
typedef struct TDEMapEntryV5
{
	uint32		cipher;
	Oid			spcOid;
	RelFileNumber relNumber;
	uint32		type;
	unsigned char entry_iv[MAP_ENTRY_IV_SIZE];
	unsigned char aead_tag[MAP_ENTRY_AEAD_TAG_SIZE];
	uint8		key_base_iv[INTERNAL_KEY_IV_LEN];
	uint8		encrypted_key_data[INTERNAL_KEY_MAX_LEN];
} TDEMapEntryV5;

static bool
read_one_map_entry_v5(int fd, TDEMapEntryV5 *entry, off_t *offset)
{
	int			bytes_read = pg_pread(fd, entry, sizeof(TDEMapEntryV5), *offset);

	if (bytes_read == 0)
		return false;
	if (bytes_read != sizeof(TDEMapEntryV5))
		ereport(ERROR,
				errmsg("keys migration: could not read map entry"));
	*offset += bytes_read;
	return true;
}

static void
ikey_from_map_entry_v5(TDEMapEntryV5 *entry, const TDEPrincipalKey *principal_key, InternalKey *out)
{
	out->key_len = open_pg_tde_cipher_key_length(entry->cipher);
	out->cipher = entry->cipher;

	memcpy(out->base_iv, entry->key_base_iv, sizeof(entry->key_base_iv));

	/* v5 excluded key_base_iv from the AAD (offsetof entry_iv). */
	if (!AesGcmDecrypt(principal_key->keyData, principal_key->keyLength,
					   entry->entry_iv, sizeof(entry->entry_iv),
					   (unsigned char *) entry, offsetof(TDEMapEntryV5, entry_iv),
					   entry->encrypted_key_data, out->key_len,
					   out->key,
					   entry->aead_tag, sizeof(entry->aead_tag)))
		ereport(ERROR,
				errmsg("failed to decrypt key, incorrect principal key or corrupted key file"));
}

static bool
map_from_disk_entry_v5(int fd, off_t *entry_offset, const TDEPrincipalKey *principal_key, TDEMapEntry *out)
{
	TDEMapEntryV5 disk_entry;
	InternalKey key;
	RelFileLocator rloc;

	if (!read_one_map_entry_v5(fd, &disk_entry, entry_offset))
		return false;

	/* Decrypting empty entries would fail, so just return an empty entry. */
	if (disk_entry.type == MAP_ENTRY_TYPE_EMPTY)
	{
		memset(out, 0, sizeof(TDEMapEntry));
		return true;
	}

	ikey_from_map_entry_v5(&disk_entry, principal_key, &key);

	rloc.spcOid = disk_entry.spcOid;
	rloc.dbOid = principal_key->keyInfo.databaseId;
	rloc.relNumber = disk_entry.relNumber;

	open_pg_tde_initialize_map_entry(out, principal_key, &rloc, &key);

	return true;
}

void
open_pg_tde_migrate_smgr_keys_file(void)
{
	DIR		   *dir;
	LWLock	   *lock_pk = tde_lwlock_enc_keys();
	struct dirent *file;

	/*
	 * No real need in lock here as the func should be called only on the
	 * server start, but GetPrincipalKey() expects lock.
	 */
	LWLockAcquire(lock_pk, LW_EXCLUSIVE);

	dir = opendir(open_pg_tde_get_data_dir());
	if (dir == NULL && errno != ENOENT)
		elog(ERROR, "could not open directory \"%s\": %m",
			 open_pg_tde_get_data_dir());

	while (errno = 0, (file = readdir(dir)) != NULL)
	{
		char		db_map_path[MAXPGPATH] = {0};
		char		tmp_db_map_path[MAXPGPATH] = {0};
		off_t		read_pos,
					write_pos;
		int			old_fd,
					new_fd;
		Oid			dbOid;
		char	   *suffix;
		TDEFileHeader fheader;
		MapFromDiskEntry read_map_entry;
		TDEMapEntry new_entry;
		TDEPrincipalKey *principal_key;
		TDESignedPrincipalKeyInfo signed_key_info;


		dbOid = strtoul(file->d_name, &suffix, 10);

		if (strcmp(suffix, "_keys") != 0)
			continue;

		open_pg_tde_set_db_file_path(dbOid, db_map_path);

		snprintf(tmp_db_map_path, MAXPGPATH, "%s.r", db_map_path);

		old_fd = open_pg_tde_open_file_basic(db_map_path, O_RDONLY | PG_BINARY, false);
		open_pg_tde_file_header_read(db_map_path, old_fd, &fheader, &read_pos);

		/*
		 * Check if we have anything to do. read_pos == 0 means the file had
		 * no header because it was empty.
		 */
		if (read_pos == 0 || fheader.file_version == OPEN_PG_TDE_SMGR_FILE_MAGIC)
		{
			CloseTransientFile(old_fd);
			continue;
		}

		/* The type check later, when extracting the principal key */
		if (FILEMAGIC_VERSION(fheader.file_version) == 3)
			read_map_entry = map_from_disk_entry_v3;
		else if (FILEMAGIC_VERSION(fheader.file_version) == 4)
			read_map_entry = map_from_disk_entry_v4;
		else if (FILEMAGIC_VERSION(fheader.file_version) == 5)
			read_map_entry = map_from_disk_entry_v5;
		else
			elog(ERROR, "keys migration: unsupported or corrupted version %d of file \"%s\"", FILEMAGIC_VERSION(fheader.file_version), db_map_path);

		/*
		 * The old file exists and it's not empty, hence a principal key
		 * should exist as well.
		 */
		principal_key = GetPrincipalKey(dbOid, LW_EXCLUSIVE);
		if (principal_key == NULL)
		{
			ereport(ERROR,
					errmsg("failed to retrieve principal key for database %u", dbOid),
					errdetail("Failed to migrate the keys file of database %u.", dbOid));
		}
		open_pg_tde_sign_principal_key_info(&signed_key_info, principal_key);

		new_fd = open_pg_tde_open_file_write(tmp_db_map_path, &signed_key_info, true, &write_pos);

		while (read_map_entry(old_fd, &read_pos, principal_key, &new_entry))
		{
			if (new_entry.type == MAP_ENTRY_TYPE_EMPTY)
				continue;
			open_pg_tde_write_one_map_entry(new_fd, &new_entry, &write_pos, db_map_path);
		}

		CloseTransientFile(old_fd);
		CloseTransientFile(new_fd);
		durable_rename(tmp_db_map_path, db_map_path, ERROR);
	}

	closedir(dir);
	LWLockRelease(lock_pk);
}

#endif