Improving performance (#117)

Improving performance

This commit tries to improve the performance of our encryption, based on the results
reported in #89.

It contains two main changes which results in improvents:

* Our pg_tde_crypt loop has several conditions to support all possible cornercases
it might encounter. These conditions are not neccessary most of the time, but slow
down our throughput considerably. This commit introduces a "simple" version of the
function for smaller data sizes, which only has a simple one instruction for loop.
* OpenSSL supports in place encryption when certain conditions are met - which is
true for our zeroblocks function. The zero block function no longer uses an
additional local array, and the filling of the input data is also simplified.

These changes together results in the following encryption overheads:

* json testcase: 140% -> 120%
* jsonb testcase: 230% -> 175%

Closes #89

src/encryption/enc_aes.c

@@ -164,33 +164,29 @@ void AesDecrypt(const unsigned char* key, const unsigned char* iv, const unsigne
 	AesRunCbc(0, key, iv, in, in_len, out, out_len);
 }
 
-/*
- * We want to avoid dynamic memory allocation, so the function only allows
- * to process NUM_AES_BLOCKS_IN_BATCH number of blocks at a time.
- * If the caller wants to process more than NUM_AES_BLOCKS_IN_BATCH * AES_BLOCK_SIZE
- * data it should divide the data into batches and call this function for each batch.
+/* This function assumes that the out buffer is big enough: at least (blockNumber2 - blockNumber1) * 16 bytes
  */
 void Aes128EncryptedZeroBlocks(void* ctxPtr, const unsigned char* key, const char* iv_prefix, uint64_t blockNumber1, uint64_t blockNumber2, unsigned char* out)
 {
-	unsigned char iv[16] = { 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00 };
+	const unsigned char iv[16] = { 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00 };
 
-	unsigned dataLen = (blockNumber2 - blockNumber1) * 16;
-	unsigned char data[DATA_BYTES_PER_AES_BATCH];
+	const unsigned dataLen = (blockNumber2 - blockNumber1) * 16;
 	int outLen;
 
 	Assert(blockNumber2 >= blockNumber1);
-	Assert(dataLen <= DATA_BYTES_PER_AES_BATCH);
 
-	memset(data, 0, dataLen);
 	for(int j=blockNumber1;j<blockNumber2;++j)
 	{
-		memcpy(data + (16*(j-blockNumber1)), iv_prefix, 16);
-
-		for(int i =0; i<8;++i) {
-			data[16*(j-blockNumber1)+15-i] = (j >> (8*i)) & 0xFF;
-		}
+		/*
+  		 * We have 16 bytes, and a 4 byte counter. The counter is the last 4 bytes.
+		 * Technically, this isn't correct: the byte order of the counter depends 
+		 * on the endianness of the CPU running it.
+		 * As this is a generic limitation of Postgres, it's fine.
+		 */
+		memcpy(out + (16*(j-blockNumber1)), iv_prefix, 12);
+		memcpy(out + (16*(j-blockNumber1)) + 12, (char*)&j, 4);
 	}
 
-	AesRunCtr(ctxPtr, 1, key, iv, data, dataLen, out, &outLen);
+	AesRunCtr(ctxPtr, 1, key, iv, out, dataLen, out, &outLen);
 	Assert(outLen == dataLen);
 }

src/encryption/enc_tde.c

@@ -47,16 +47,53 @@ SetIVPrefix(ItemPointerData* ip, CommandId cid, char* iv_prefix)
  */
 
 /* 
- * pg_tde_crypt:
+ * pg_tde_crypt_simple:
  * Encrypts/decrypts `data` with a given `key`. The result is written to `out`.
  * start_offset: is the absolute location of start of data in the file.
+ * This function assumes that everything is in a single block, and has an assertion ensuring this
  */
-void
-pg_tde_crypt(const char* iv_prefix, uint32 start_offset, const char* data, uint32 data_len, char* out, RelKeyData* key, const char* context)
+static void
+pg_tde_crypt_simple(const char* iv_prefix, uint32 start_offset, const char* data, uint32 data_len, char* out, RelKeyData* key, const char* context)
+{
+	const uint64 aes_start_block = start_offset / AES_BLOCK_SIZE;
+	const uint64 aes_end_block = (start_offset + data_len + (AES_BLOCK_SIZE - 1)) / AES_BLOCK_SIZE;
+	const uint64 aes_block_no = start_offset % AES_BLOCK_SIZE;
+
+	unsigned char enc_key[DATA_BYTES_PER_AES_BATCH + AES_BLOCK_SIZE];
+
+	Assert(aes_end_block - aes_start_block <= NUM_AES_BLOCKS_IN_BATCH + 1);
+
+	Aes128EncryptedZeroBlocks(&(key->internal_key.ctx), key->internal_key.key, iv_prefix, aes_start_block, aes_end_block, enc_key);
+
+#ifdef ENCRYPTION_DEBUG
 {
-	uint64 aes_start_block = start_offset / AES_BLOCK_SIZE;
-	uint64 aes_end_block = (start_offset + data_len + (AES_BLOCK_SIZE -1)) / AES_BLOCK_SIZE;
-	uint64 aes_block_no = start_offset % AES_BLOCK_SIZE;
+	char ivp_debug[33];
+	iv_prefix_debug(iv_prefix, ivp_debug);
+	ereport(LOG,
+		(errmsg("%s: Start offset: %lu Data_Len: %u, aes_start_block: %lu, aes_end_block: %lu, IV prefix: %s",
+			context?context:"", start_offset, data_len, aes_start_block, aes_end_block, ivp_debug)));
+}
+#endif
+
+	for(uint32 i = 0; i < data_len; ++i)
+	{
+		out[i] = data[i] ^ enc_key[i + aes_block_no];
+	}
+}
+
+
+/* 
+ * pg_tde_crypt_complex:
+ * Encrypts/decrypts `data` with a given `key`. The result is written to `out`.
+ * start_offset: is the absolute location of start of data in the file.
+ * This is a generic function indented for large data, that od not fit into a single block
+ */
+static void
+pg_tde_crypt_complex(const char* iv_prefix, uint32 start_offset, const char* data, uint32 data_len, char* out, RelKeyData* key, const char* context)
+{
+	const uint64 aes_start_block = start_offset / AES_BLOCK_SIZE;
+	const uint64 aes_end_block = (start_offset + data_len + (AES_BLOCK_SIZE -1)) / AES_BLOCK_SIZE;
+	const uint64 aes_block_no = start_offset % AES_BLOCK_SIZE;
 	uint32 batch_no = 0;
 	uint32 data_index = 0;
 	uint64 batch_end_block;
@@ -111,6 +148,24 @@ pg_tde_crypt(const char* iv_prefix, uint32 start_offset, const char* data, uint3
     }
 }
 
+/* 
+ * pg_tde_crypt:
+ * Encrypts/decrypts `data` with a given `key`. The result is written to `out`.
+ * start_offset: is the absolute location of start of data in the file.
+ * This function simply selects between the two above variations based on the data length
+ */
+void
+pg_tde_crypt(const char* iv_prefix, uint32 start_offset, const char* data, uint32 data_len, char* out, RelKeyData* key, const char* context)
+{
+	if(data_len >= DATA_BYTES_PER_AES_BATCH)
+	{
+		pg_tde_crypt_complex(iv_prefix, start_offset, data, data_len, out, key, context);
+	} else
+	{
+		pg_tde_crypt_simple(iv_prefix, start_offset, data, data_len, out, key, context);
+	}
+}
+
 /*
  * pg_tde_crypt_tuple:
  * Does the encryption/decryption of tuple data in place

src/include/encryption/enc_aes.h

@@ -13,7 +13,7 @@
 #include <stdint.h>
 
 #define AES_BLOCK_SIZE 		        16
-#define NUM_AES_BLOCKS_IN_BATCH     100
+#define NUM_AES_BLOCKS_IN_BATCH     200
 #define DATA_BYTES_PER_AES_BATCH    (NUM_AES_BLOCKS_IN_BATCH * AES_BLOCK_SIZE)
 
 void AesInit(void);