Improving performance

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

@@ -43,16 +43,53 @@ SetIVPrefix(ItemPointerData* ip, 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;
@@ -107,6 +144,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);