From 4da83e60ce43bf72c89acd553dbc6bf71039d232 Mon Sep 17 00:00:00 2001 From: Zsolt Parragi Date: Mon, 15 Apr 2024 19:48:22 +0100 Subject: [PATCH] Simple SMGR based AES-CBC encryption --- meson.build | 2 + src/include/smgr/pg_tde_smgr.h | 4 + src/pg_tde.c | 3 + src/smgr/pg_tde_smgr.c | 155 +++++++++++++++++++++++++++++++++ 4 files changed, 164 insertions(+) create mode 100644 src/include/smgr/pg_tde_smgr.h create mode 100644 src/smgr/pg_tde_smgr.c diff --git a/meson.build b/meson.build index 366770e0..fd15921e 100644 --- a/meson.build +++ b/meson.build @@ -38,6 +38,8 @@ pg_tde_sources = files( 'src/keyring/keyring_vault.c', 'src/keyring/keyring_api.c', + 'src/smgr/pg_tde_smgr.c', + 'src/catalog/tde_keyring.c', 'src/catalog/tde_master_key.c', 'src/common/pg_tde_shmem.c', diff --git a/src/include/smgr/pg_tde_smgr.h b/src/include/smgr/pg_tde_smgr.h new file mode 100644 index 00000000..90df06c4 --- /dev/null +++ b/src/include/smgr/pg_tde_smgr.h @@ -0,0 +1,4 @@ + +#pragma once + +extern void RegisterStorageMgr(); \ No newline at end of file diff --git a/src/pg_tde.c b/src/pg_tde.c index 7b076592..ed996cff 100644 --- a/src/pg_tde.c +++ b/src/pg_tde.c @@ -31,6 +31,7 @@ #include "keyring/keyring_vault.h" #include "utils/builtins.h" #include "pg_tde_defs.h" +#include "smgr/pg_tde_smgr.h" #define MAX_ON_INSTALLS 5 @@ -98,6 +99,8 @@ _PG_init(void) InstallFileKeyring(); InstallVaultV2Keyring(); RegisterCustomRmgr(RM_TDERMGR_ID, &pg_tde_rmgr); + + RegisterStorageMgr(); } Datum pg_tde_extension_initialize(PG_FUNCTION_ARGS) diff --git a/src/smgr/pg_tde_smgr.c b/src/smgr/pg_tde_smgr.c new file mode 100644 index 00000000..9e679ec5 --- /dev/null +++ b/src/smgr/pg_tde_smgr.c @@ -0,0 +1,155 @@ + +#include "smgr/pg_tde_smgr.h" +#include "postgres.h" +#include "storage/smgr.h" +#include "storage/md.h" +#include "catalog/catalog.h" +#include "encryption/enc_aes.h" + +#if PG_VERSION_NUM >= 170000 + +// TODO: implement proper key/IV +static char key[16] = {0,}; +// iv should be based on blocknum, available in the API +static char iv[16] = {0,}; + +void +tde_mdwritev(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum, + const void **buffers, BlockNumber nblocks, bool skipFsync) +{ + AesInit(); + + char* local_blocks = malloc( BLCKSZ * (nblocks+1) ); + char* local_blocks_aligned = (char*)TYPEALIGN(PG_IO_ALIGN_SIZE, local_blocks); + const void** local_buffers = malloc ( sizeof(void*) * nblocks ); + + // TODO: add check to only encrypt/decrypt tables with specific AM/DB? + + if(IsCatalogRelationOid(reln->smgr_rlocator.locator.spcOid)) + { + // Don't try to encrypt catalog tables: + // Issues with bootstrap and encryption metadata + + mdwritev(reln, forknum, blocknum, buffers, nblocks, skipFsync); + + return; + } + + for(int i = 0; i < nblocks; ++i ) + { + local_buffers[i] = &local_blocks_aligned[i*BLCKSZ]; + int out_len = BLCKSZ; + AesEncrypt(key, iv, ((char**)buffers)[i], BLCKSZ, local_buffers[i], &out_len); + } + + mdwritev(reln, forknum, blocknum, + local_buffers, nblocks, skipFsync); + + free(local_blocks); + free(local_buffers); +} + +void +tde_mdextend(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum, + const void *buffer, bool skipFsync) +{ + AesInit(); + + char* local_blocks = malloc( BLCKSZ * (1+1) ); + char* local_blocks_aligned = (char*)TYPEALIGN(PG_IO_ALIGN_SIZE, local_blocks); + + // TODO: add check to only encrypt/decrypt tables with specific AM/DB? + + if(IsCatalogRelationOid(reln->smgr_rlocator.locator.spcOid)) + { + // Don't try to encrypt catalog tables: + // Issues with bootstrap and encryption metadata + mdextend(reln, forknum, blocknum, buffer, skipFsync); + + return; + } + + int out_len = BLCKSZ; + AesEncrypt(key, iv, ((char*)buffer), BLCKSZ, local_blocks_aligned, &out_len); + + mdextend(reln, forknum, blocknum, local_blocks_aligned, skipFsync); + + + free(local_blocks); +} + +void +tde_mdreadv(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum, + void **buffers, BlockNumber nblocks) +{ + AesInit(); + + mdreadv(reln, forknum, blocknum, buffers, nblocks); + + // TODO: add check to only encrypt/decrypt tables with specific AM/DB? + + // Don't try to decrypt catalog tables, those are not encrypted + if(IsCatalogRelationOid(reln->smgr_rlocator.locator.spcOid)) + { + return; + } + + for(int i = 0; i < nblocks; ++i) + { + bool allZero = true; + for(int j = 0; j < 32; ++j) + { + if(((char**)buffers)[i][j] != 0) + { + // Postgres creates all zero blocks in an optimized route, which we do not try + // to encrypt. + // Instead we detect if a block is all zero at decryption time, and + // leave it as is. + // This could be a security issue later, but it is a good first prototype + allZero = false; + break; + } + } + if(allZero) continue; + + int out_len = BLCKSZ; + AesDecrypt(key, iv, ((char**)buffers)[i], BLCKSZ, ((char**)buffers)[i], &out_len); + } + + // And now decrypt buffers in place + // We check the first few bytes of the page: if all zero, we assume it is zero and keep it as is +} +static SMgrId tde_smgr_id; +static const struct f_smgr tde_smgr = { + .name = "tde", + .smgr_init = mdinit, + .smgr_shutdown = NULL, + .smgr_open = mdopen, + .smgr_close = mdclose, + .smgr_create = mdcreate, + .smgr_exists = mdexists, + .smgr_unlink = mdunlink, + .smgr_extend = tde_mdextend, + .smgr_zeroextend = mdzeroextend, + .smgr_prefetch = mdprefetch, + .smgr_readv = tde_mdreadv, + .smgr_writev = tde_mdwritev, + .smgr_writeback = mdwriteback, + .smgr_nblocks = mdnblocks, + .smgr_truncate = mdtruncate, + .smgr_immedsync = mdimmedsync, +}; + +void RegisterStorageMgr() +{ + tde_smgr_id = smgr_register(&tde_smgr, 0); + + // TODO: figure out how this part should work in a real extension + storage_manager_id = tde_smgr_id; +} + +#else +void RegisterStorageMgr() +{ +} +#endif \ No newline at end of file