0009-Add-kmgr-plugin-test-module-kmgr_file.patch

From 4197b11bb91d9442777637dd2dfbadbb1deb9db1 Mon Sep 17 00:00:00 2001
From: Masahiko Sawada <sawada.mshk@gmail.com>
Date: Fri, 28 Jun 2019 18:12:18 +0900
Subject: [PATCH 9/9] Add kmgr plugin test module: kmgr_file.

---
 contrib/Makefile                 |   3 +-
 contrib/kmgr_file/Makefile       |  20 +++
 contrib/kmgr_file/kmgr_file.c    | 311 +++++++++++++++++++++++++++++++++++++++
 contrib/kmgr_file/t/001_basic.pl | 123 ++++++++++++++++
 4 files changed, 456 insertions(+), 1 deletion(-)
 create mode 100644 contrib/kmgr_file/Makefile
 create mode 100644 contrib/kmgr_file/kmgr_file.c
 create mode 100644 contrib/kmgr_file/t/001_basic.pl

diff --git a/contrib/Makefile b/contrib/Makefile
index 92184ed..1e70603 100644
--- a/contrib/Makefile
+++ b/contrib/Makefile
@@ -48,7 +48,8 @@ SUBDIRS = \
 		tsm_system_rows \
 		tsm_system_time \
 		unaccent	\
-		vacuumlo
+		vacuumlo	\
+		kmgr_file
 
 ifeq ($(with_openssl),yes)
 SUBDIRS += sslinfo
diff --git a/contrib/kmgr_file/Makefile b/contrib/kmgr_file/Makefile
new file mode 100644
index 0000000..7430445
--- /dev/null
+++ b/contrib/kmgr_file/Makefile
@@ -0,0 +1,20 @@
+# contrib/kmgr_file/Makefile
+
+MODULES = kmgr_file
+PGFILEDESC = "kmgr_file - example of a kmgr plugin"
+
+OBJS = kmgr_file.o
+
+PREGRESS = kmgr_test
+TAP_TESTS = 1
+
+ifdef USE_PGXS
+PG_CONFIG = pg_config
+PGXS := $(shell $(PG_CONFIG) --pgxs)
+include $(PGXS)
+else
+subdir = contrib/kmgr_file
+top_builddir = ../..
+include $(top_builddir)/src/Makefile.global
+include $(top_srcdir)/contrib/contrib-global.mk
+endif
diff --git a/contrib/kmgr_file/kmgr_file.c b/contrib/kmgr_file/kmgr_file.c
new file mode 100644
index 0000000..8465556
--- /dev/null
+++ b/contrib/kmgr_file/kmgr_file.c
@@ -0,0 +1,311 @@
+/*-------------------------------------------------------------------------
+ *
+ * keyring_file.c
+ *	  Test module for Key ring plugin
+ *
+ * Portions Copyright (c) 1996-2019, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ * IDENTIFICATION
+ *	  contrib/test_keyringfile/keyring_file.c
+ *-------------------------------------------------------------------------
+ */
+
+#include "postgres.h"
+
+#include <unistd.h>
+
+#include "fmgr.h"
+#include "utils/hsearch.h"
+#include "storage/fd.h"
+#include "storage/kmgr.h"
+#include "storage/kmgr_api.h"
+#include "storage/ipc.h"
+#include "storage/lwlock.h"
+#include "storage/shmem.h"
+#include "utils/memutils.h"
+#include "utils/hashutils.h"
+#include "utils/guc.h"
+#include "miscadmin.h"
+
+#define KEYRING_MAX_KEYS 128
+#define TEST_MASTERKEY_FILENAME "kmgr_test"
+
+
+PG_MODULE_MAGIC;
+
+typedef char KeyId[MASTER_KEY_ID_LEN];
+
+typedef struct MyKey
+{
+	char	id[MASTER_KEY_ID_LEN];
+	char	key[ENCRYPTION_KEY_SIZE];
+} MyKey;
+
+static HTAB *MyKeys;
+static LWLock *lock;
+static char *masterkey_filepath;
+
+/* function prototypes */
+extern void _PG_kmgr_init(KmgrPluginCallbacks *cb);
+static void test_startup(void);
+static char *test_getkey(const char *keyid);
+static void test_generatekey(const char *keyid);
+static void test_removekey(const char *keyid);
+static bool test_isexistkey(const char *keyid);
+static void load_all_keys(void);
+static void save_all_keys(void);
+static Size test_memsize(void);
+
+/*
+ * Specify output plugin callbacks
+ */
+void
+_PG_kmgr_init(KmgrPluginCallbacks *cb)
+{
+	AssertVariableIsOfType(&_PG_kmgr_init, KmgrPluginInit);
+
+	cb->startup_cb = test_startup;
+	cb->getkey_cb = test_getkey;
+	cb->generatekey_cb = test_generatekey;
+	cb->isexistkey_cb = test_isexistkey;
+	cb->removekey_cb = test_removekey;
+
+	DefineCustomStringVariable("kmgr_file.masterkey_filepath",
+							   "file location of key file",
+							   NULL,
+							   &masterkey_filepath,
+							   "global/",
+							   PGC_POSTMASTER,
+							   0,
+							   NULL, NULL, NULL);
+
+	RequestAddinShmemSpace(test_memsize());
+	RequestNamedLWLockTranche("kmgr_test", 1);
+}
+
+static Size
+test_memsize(void)
+{
+	Size size;
+
+	size = MAXALIGN(sizeof(lock));
+	size = add_size(size, hash_estimate_size(KEYRING_MAX_KEYS, sizeof(MyKey)));
+
+	return size;
+}
+
+/*
+ * Read all keys in my keyfile(kmgr_test.kr) and build hash map for keys.
+ */
+static void
+load_all_keys(void)
+{
+	char path[MAXPGPATH];
+	int fd;
+
+	LWLockAcquire(lock, LW_EXCLUSIVE);
+
+	sprintf(path, "%s/%s", masterkey_filepath, TEST_MASTERKEY_FILENAME);
+
+	fd = OpenTransientFile(path, O_RDONLY | PG_BINARY);
+
+	if (fd < 0)
+	{
+		if (errno == ENOENT)
+		{
+			LWLockRelease(lock);
+			return;
+		}
+
+		ereport(ERROR,
+				(errcode_for_file_access(),
+				 errmsg("could not open file \"%s\": %m", path)));
+	}
+
+	for (;;)
+	{
+		MyKey *key = palloc(sizeof(MyKey));
+		MyKey *keycache;
+		int read_len;
+
+		read_len = read(fd, key, sizeof(MyKey));
+
+		if (read_len < 0)
+			ereport(ERROR,
+					(errcode_for_file_access(),
+					 (errmsg("could not read from file \"%s\": %m", path))));
+		else if (read_len == 0) /* EOF */
+			break;
+		else if (read_len != sizeof(MyKey))
+			ereport(ERROR,
+					(errcode_for_file_access(),
+					 errmsg("could not read from file \"%s\": read %d instead of %d bytes",
+							path, read_len, (int32) sizeof(MyKey))));
+
+		keycache = hash_search(MyKeys, (void *) key->id, HASH_ENTER, NULL);
+
+		memcpy(keycache->key, key->key, ENCRYPTION_KEY_SIZE);
+	}
+
+	CloseTransientFile(fd);
+	LWLockRelease(lock);
+}
+
+/*
+ * Persistent all keys in my key file(kmgr_test).
+ */
+static void
+save_all_keys(void)
+{
+	HASH_SEQ_STATUS status;
+	char path[MAXPGPATH];
+	MyKey *key;
+	FILE *fpout;
+	int	rc;
+
+	sprintf(path, "%s/%s", masterkey_filepath, TEST_MASTERKEY_FILENAME);
+	fpout = AllocateFile(path, PG_BINARY_W);
+	if (fpout == NULL)
+	{
+		ereport(LOG,
+				(errcode_for_file_access(),
+				 errmsg("could not open temporary keyring file \"%s\": %m",
+						path)));
+		return;
+	}
+
+	/* Iterate all keys while writing them to the file */
+	hash_seq_init(&status, MyKeys);
+	while ((key = (MyKey *) hash_seq_search(&status)) != NULL)
+	{
+		rc = fwrite(key, sizeof(MyKey), 1, fpout);
+		(void) rc;
+	}
+
+	if (ferror(fpout))
+	{
+		ereport(LOG,
+				(errcode_for_file_access(),
+				 errmsg("could not write temporary keyring file \"%s\": %m",
+						path)));
+	}
+	else if (FreeFile(fpout) < 0)
+	{
+		ereport(LOG,
+				(errcode_for_file_access(),
+				 errmsg("could not close temporary keyring file \"%s\": %m",
+						path)));
+	}
+}
+
+static void
+test_startup(void)
+{
+	HASHCTL ctl;
+
+	ereport(LOG, (errmsg("keyring_file: starting up")));
+
+	//LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
+
+	MemSet(&ctl, 0, sizeof(ctl));
+	ctl.keysize = MASTER_KEY_ID_LEN;
+	ctl.entrysize = sizeof(MyKey);
+
+	MyKeys = ShmemInitHash("test_kmgr key shared hash map",
+						   KEYRING_MAX_KEYS, KEYRING_MAX_KEYS,
+						   &ctl,
+						   HASH_ELEM | HASH_BLOBS);
+
+	lock = &(GetNamedLWLockTranche("kmgr_test"))->lock;
+
+	//LWLockRelease(AddinShmemInitLock);
+
+	load_all_keys();
+}
+
+static char *
+test_getkey(const char *keyid)
+{
+	char *key;
+	MyKey *mykey;
+	bool found;
+
+	LWLockAcquire(lock, LW_SHARED);
+
+	mykey = hash_search(MyKeys, (void *) keyid, HASH_FIND,	&found);
+
+	if (!found)
+		ereport(ERROR,
+				(errmsg("kmgr_file: could not get master key \"%s\"",
+						keyid)));
+
+	/* Set master key */
+	key = palloc0(ENCRYPTION_KEY_SIZE);
+	memcpy(key, mykey->key, ENCRYPTION_KEY_SIZE);
+
+	LWLockRelease(lock);
+
+	return key;
+}
+
+static bool
+test_isexistkey(const char *keyid)
+{
+	bool found;
+
+	Assert(keyid != NULL);
+
+	LWLockAcquire(lock, LW_SHARED);
+
+	(void *) hash_search(MyKeys, (void *) keyid, HASH_FIND,	&found);
+
+	LWLockRelease(lock);
+
+	return found;
+}
+
+static void
+test_generatekey(const char *keyid)
+{
+	MyKey *mykey;
+	bool found;
+	char *newkey;
+	int ret;
+
+	Assert(keyid);
+
+	LWLockAcquire(lock, LW_EXCLUSIVE);
+
+	/* Set master key */
+	newkey = (char *) palloc0(ENCRYPTION_KEY_SIZE);
+	ret = pg_strong_random(newkey, ENCRYPTION_KEY_SIZE);
+	if (!ret)
+		ereport(ERROR,
+				(errmsg("failed to generate tablespace encryption key")));
+
+	/* Duplication check */
+	mykey = hash_search(MyKeys, (void *) keyid, HASH_FIND,	&found);
+
+	if (found)
+		ereport(ERROR,
+				(errmsg("key \"%s\" is already registered", keyid)));
+
+	/* Ok, register the key */
+	mykey = hash_search(MyKeys, (void *) keyid,	HASH_ENTER,	&found);
+	Assert(!found);
+
+	memcpy(mykey->key, newkey, ENCRYPTION_KEY_SIZE);
+
+	/* update key file */
+	save_all_keys();
+
+	LWLockRelease(lock);
+}
+
+static void
+test_removekey(const char *keyid)
+{
+	/* Not implemented yet */
+	return;
+}
diff --git a/contrib/kmgr_file/t/001_basic.pl b/contrib/kmgr_file/t/001_basic.pl
new file mode 100644
index 0000000..677b7cb
--- /dev/null
+++ b/contrib/kmgr_file/t/001_basic.pl
@@ -0,0 +1,123 @@
+use strict;
+use warnings;
+use PostgresNode;
+use TestLib;
+use Test::More tests => 9;
+
+my $keyword = "secret keyword";
+my $node = get_new_node('test');
+$node->init();
+$node->append_conf('postgresql.conf',
+				   qq(kmgr_plugin_library = 'kmgr_file'));
+$node->start;
+
+# Check is the given relation file is encrypted
+sub is_encrypted
+{
+	my $node = shift;
+	my $filepath = shift;
+	my $expected = shift;
+	my $testname = shift;
+	my $pgdata = $node->data_dir;
+
+	open my $file, '<' , "$pgdata/$filepath";
+	sysread $file, my $buffer, 8192;
+
+	my $ret = $buffer !~ /$keyword/ ? 1 : 0;
+
+	is($ret, $expected, $testname);
+
+	close $file;
+}
+
+# Prepare both encrypted and non-encrypted tablespaces
+my $enctblspc = $node->data_dir() . '_enctblspc';
+my $plntblspc = $node->data_dir() . '_plntblspc';
+mkdir $enctblspc or die;
+mkdir $plntblspc or die;
+$node->safe_psql('postgres',
+				 qq(
+				 CREATE TABLESPACE enctblspc LOCATION '$enctblspc' WITH (encryption = on);
+				 CREATE TABLESPACE plntblspc LOCATION '$plntblspc' WITH (encryption = off);
+				 ));
+
+# populate tables
+$node->safe_psql('postgres',
+				 qq(
+				 CREATE TABLE enc (a text) TABLESPACE enctblspc;
+				 CREATE TABLE pln (a text) TABLESPACE plntblspc;
+				 ));
+$node->safe_psql('postgres',
+				 qq(
+				 INSERT INTO enc VALUES ('$keyword');
+				 INSERT INTO pln SELECT ('$keyword');
+				 ));
+
+my $enc_filepath = $node->safe_psql('postgres', "SELECT pg_relation_filepath('enc')");
+my $pln_filepath = $node->safe_psql('postgres', "SELECT pg_relation_filepath('pln')");
+
+# Read encrypted table
+my $ret = $node->safe_psql('postgres', 'SELECT a FROM enc');
+is($ret, "$keyword", 'Read encrypted table');
+
+# Sync to disk
+$node->safe_psql('postgres', 'CHECKPOINT');
+
+# Encrypted table must be encrypted
+is_encrypted($node, $enc_filepath, 1, 'encrypted table is encrypted');
+
+# Plain table must not be encrypted
+is_encrypted($node, $pln_filepath, 0, 'non-encrypted table is not encrypted');
+
+# Read encrypted table from disk while decryption
+$node->restart;
+$ret = $node->safe_psql('postgres', 'SELECT a FROM enc');
+is($ret, "$keyword", 'Read encrypted table');
+
+# Rotate enryption keys, and read encrypted table
+$node->psql('postgres',
+			'SELECT pg_rotate_encryption_key()');
+$ret = $node->safe_psql('postgres', 'SELECT a FROM enc');
+is($ret, "$keyword", 'Read encrypted table');
+
+
+# Create an encrypted database, and create table
+$node->psql('postgres',
+			'CREATE DATABASE encdb TABLESPACE enctblspc');
+
+# System catalog on the encrypted database must be encrypted.
+my $pgclass_filepath = $node->safe_psql('encdb',
+									   "SELECT pg_relation_filepath('pg_class')");
+is_encrypted($node, $pgclass_filepath, 1, 'system catalog is encrypted');
+
+# Change encryption properly by moving tablespace
+$node->psql('postgres',
+			qq(
+			ALTER TABLE enc SET TABLESPACE plntblspc;
+			ALTER TABLE pln SET TABLESPACE enctblspc;
+			));
+
+# Now, the result of is_encryption is reverse
+$enc_filepath = $node->safe_psql('postgres', "SELECT pg_relation_filepath('enc')");
+$pln_filepath = $node->safe_psql('postgres', "SELECT pg_relation_filepath('pln')");
+$node->psql('postgres', 'CHECKPOINT');
+is_encrypted($node, $enc_filepath, 0, 'encrypted table gets plain');
+is_encrypted($node, $pln_filepath, 1, 'plain tble gets encrypted');
+
+# Create another encrypted tablespace and table
+my $enctblspc2 = $node->data_dir() . '_enctblspc2';
+mkdir $enctblspc2 or die;
+$node->safe_psql('postgres',
+				 qq(CREATE TABLESPACE enctblspc2 LOCATION '$enctblspc2' WITH (encryption = on);));
+$node->safe_psql('postgres',
+				 qq(
+				 CREATE TABLE enc2 (a text) TABLESPACE enctblspc2;
+				 INSERT INTO enc2 VALUES ('$keyword')));
+
+# We create the encrypted tablespace and table during recovery
+$node->teardown_node;
+$node->start;
+
+# Check the table is recovered successfully
+$ret = $node->safe_psql('postgres', 'SELECT a FROM enc2');
+is($ret, "$keyword", 'Read encrypted table');
-- 
1.8.3.1