AES wrapper rework

src/cryptography/AES.c

@@ -1,52 +1,62 @@
 #include "cryptography.h"
 #include <assert.h>
 
-//TODO: use AES CTR encryption instead of my own padding algorithm
+// write data from src to array and increment array data pointer
+static inline void __Array_writeNext(Array(u8)* dst, void* src, size_t size){
+    memcpy(dst->data, src, size);
+    dst->data = (u8*)dst->data + size;
+    dst->size -= size;
+}
+
+// read data from array to dst and increment array data pointer
+static inline void __Array_readNext(void* dst, Array(u8)* src, size_t size){
+    memcpy(dst, src->data, size);
+    src->data = (u8*)src->data + size;
+    src->size -= size;
+}
+
 
 void EncryptorAES_construct(EncryptorAES* ptr, Array(u8) key){
     assert(key.size == 16 || key.size == 24 || key.size == 32);
     
+    //TODO: use AES CTR encryption?
     br_aes_ct64_cbcenc_init(&ptr->enc_ctx, key.data, key.size);
-    ptr->rng_ctx.vtable = &br_hmac_drbg_vtable;
-    rng_init_sha256_seedFromTime(&ptr->rng_ctx.vtable);
 
-    memset(ptr->buf, 0, __AES_BUFFER_SIZE);
+    ptr->rng_ctx.vtable = &br_hmac_drbg_vtable;
-    memset(ptr->iv, 0, sizeof(ptr->iv));
+    rng_init_sha256_seedFromSystem(&ptr->rng_ctx.vtable);
 }
 
 void EncryptorAES_encrypt(EncryptorAES* ptr, Array(u8) src, Array(u8) dst){
     assert(dst.size >= EncryptorAES_calcDstSize(src.size));
     
-    // write random bytes to the beginning of the buffer
+    // generate random initial vector
-    br_hmac_drbg_generate(&ptr->rng_ctx, ptr->buf, __AES_RANDOM_BYTES_N);
+    br_hmac_drbg_generate(&ptr->rng_ctx, ptr->iv, __AES_IV_SIZE);
-    const EncryptedBlockInfo block_info = { .padding_size = 16 - src.size % 16 };
+    // write IV to the beginning of dst
-    // write struct after random_bytes
+    __Array_writeNext(&dst, ptr->iv, __AES_IV_SIZE);
-    memcpy((u8*)ptr->buf + __AES_RANDOM_BYTES_N, &block_info, sizeof(EncryptedBlockInfo));
-    // encrypt buffer
-    const u32 header_size = __AES_RANDOM_BYTES_N + sizeof(EncryptedBlockInfo);
-    br_aes_ct64_cbcenc_run(&ptr->enc_ctx, ptr->iv, ptr->buf, header_size);
-    // write encrypted header to dst
-    memcpy(dst.data, ptr->buf, header_size);
-    dst.data = (u8*)dst.data + header_size;
-    dst.size -= header_size;
 
-    // write full blocks
+    const EncryptedBlockHeader header = { .padding_size = 16 - src.size % 16 };
+    // write header to buffer
+    memcpy(ptr->buf, &header, sizeof(header));
+    // encrypt header
+    br_aes_ct64_cbcenc_run(&ptr->enc_ctx, ptr->iv, ptr->buf, sizeof(header));
+    // write encrypted header to dst
+    __Array_writeNext(&dst, ptr->buf, sizeof(header));
+
+    // encrypt full EncryptorAES buffers
     while(src.size > __AES_BUFFER_SIZE){
-        memcpy(ptr->buf, src.data, __AES_BUFFER_SIZE);
+        __Array_readNext(ptr->buf, &src, __AES_BUFFER_SIZE);
-        src.data = (u8*)src.data + __AES_BUFFER_SIZE;
-        src.size -= __AES_BUFFER_SIZE;
         br_aes_ct64_cbcenc_run(&ptr->enc_ctx, ptr->iv, ptr->buf, __AES_BUFFER_SIZE);
-        memcpy(dst.data, ptr->buf, __AES_BUFFER_SIZE);
+        __Array_writeNext(&dst, ptr->buf, __AES_BUFFER_SIZE);
-        dst.data = (u8*)dst.data + __AES_BUFFER_SIZE;
-        dst.size -= __AES_BUFFER_SIZE;
     }
 
-    // write incomplete block
+    // encrypt buffer with remaining data
-    memset(ptr->buf, 0, __AES_BUFFER_SIZE);
+    if(src.size > 0){
         memcpy(ptr->buf, src.data, src.size);
-    u32 src_size_padded = src.size + block_info.padding_size;
+        u32 src_size_padded = src.size + header.padding_size;
+        memset(ptr->buf + src.size, 0, header.padding_size);
         br_aes_ct64_cbcenc_run(&ptr->enc_ctx, ptr->iv, ptr->buf, src_size_padded);
         memcpy(dst.data, ptr->buf, src_size_padded);
+    }
 }
 
 
@@ -54,41 +64,37 @@ void DecryptorAES_construct(DecryptorAES* ptr, Array(u8) key){
     assert(key.size == 16 || key.size == 24 || key.size == 32);
     
     br_aes_ct64_cbcdec_init(&ptr->dec_ctx, key.data, key.size);
-
-    memset(ptr->buf, 0, __AES_BUFFER_SIZE);
-    memset(ptr->iv, 0, sizeof(ptr->iv));
 }
 
 void DecryptorAES_decrypt(DecryptorAES* ptr, Array(u8) src, Array(u8) dst, u32* decrypted_size){
+    assert(src.size >= EncryptorAES_calcDstSize(0));
+    assert(src.size % 16 == 0 && "src must be array of 16-byte blocks");
     assert(dst.size >= src.size);
     
-    // copy encrypted header from src to buffer
+    // read IV from the beginning of src
-    const u32 header_size = __AES_RANDOM_BYTES_N + sizeof(EncryptedBlockInfo);
+    __Array_readNext(ptr->iv, &src, __AES_IV_SIZE);
-    memcpy(ptr->buf, src.data, header_size);
-    src.data = (u8*)src.data + header_size;
-    src.size -= header_size;
-    // decrypt buffer
-    br_aes_ct64_cbcdec_run(&ptr->dec_ctx, ptr->iv, &ptr->buf, header_size);
-    // read EncryptedBlockInfo from buffer
-    EncryptedBlockInfo block_info;
-    memcpy(&block_info, (u8*)ptr->buf + __AES_RANDOM_BYTES_N, sizeof(EncryptedBlockInfo));
-    *decrypted_size = src.size - block_info.padding_size;
 
-    // write full blocks
+    EncryptedBlockHeader header;
+    // read encrypted header from src
+    __Array_readNext(&header, &src, sizeof(header));
+    // decrypt header
+    br_aes_ct64_cbcdec_run(&ptr->dec_ctx, ptr->iv, &header, sizeof(header));
+    // size of decrypted data without padding
+    *decrypted_size = src.size - header.padding_size;
+    const u32 src_size_padded = src.size;
+    src.size = *decrypted_size;
+
+    // decrypt full buffers
     while(src.size > __AES_BUFFER_SIZE){
-        memcpy(ptr->buf, src.data, __AES_BUFFER_SIZE);
+        __Array_readNext(ptr->buf, &src, __AES_BUFFER_SIZE);
-        src.data = (u8*)src.data + __AES_BUFFER_SIZE;
-        src.size -= __AES_BUFFER_SIZE;
         br_aes_ct64_cbcdec_run(&ptr->dec_ctx, ptr->iv, ptr->buf, __AES_BUFFER_SIZE);
-        memcpy(dst.data, ptr->buf, __AES_BUFFER_SIZE);
+        __Array_writeNext(&dst, ptr->buf, __AES_BUFFER_SIZE);
-        dst.data = (u8*)dst.data + __AES_BUFFER_SIZE;
-        dst.size -= __AES_BUFFER_SIZE;
     }
     
-    // write incomplete block
+    // decrypt buffer with remaining data
-    memset(ptr->buf, 0, __AES_BUFFER_SIZE);
+    if(src.size > 0){
         memcpy(ptr->buf, src.data, src.size);
-    u32 src_size_padded = src.size + block_info.padding_size;
         br_aes_ct64_cbcdec_run(&ptr->dec_ctx, ptr->iv, ptr->buf, src_size_padded);
-    memcpy(dst.data, ptr->buf, src_size_padded);
+        memcpy(dst.data, ptr->buf, src.size);
+    }
 }

src/cryptography/cryptography.h

@@ -51,22 +51,20 @@ void rng_init_sha256_seedFromTime(const br_prng_class** rng_vtable_ptr);
 #define __AES_SESSION_KEY_SIZE 32
 #define __AES_DB_KEY_SIZE 32
 
-typedef struct EncryptedBlockInfo {
+typedef struct EncryptedBlockHeader {
     u8 padding_size;
-    u32 _reserved;
+} __attribute__((aligned(16))) EncryptedBlockHeader;
-    u64 __reserved;
-} __attribute__((aligned(16))) EncryptedBlockInfo;
 
 // must be multiple of 16
-#define __AES_BUFFER_SIZE 512
+#define __AES_IV_SIZE 16
 // must be multiple of 16
-#define __AES_RANDOM_BYTES_N 16
+#define __AES_BUFFER_SIZE 512
 
 typedef struct EncryptorAES {
     br_aes_ct64_cbcenc_keys enc_ctx;
     br_hmac_drbg_context rng_ctx;
+    u8 iv[__AES_IV_SIZE];
     u8 buf[__AES_BUFFER_SIZE];
-    u8 iv[16];
 } EncryptorAES;
 
 /// @param key Array<u8, 16 | 24 | 32>
@@ -77,20 +75,20 @@ void EncryptorAES_construct(EncryptorAES* ptr, Array(u8) key);
 /// @param dst array of size >= EncryptorAES_calcDstSize(src.size)
 void EncryptorAES_encrypt(EncryptorAES* ptr, Array(u8) src, Array(u8) dst);
 
-#define EncryptorAES_calcDstSize(SRC_SIZE) (ALIGN_TO(SRC_SIZE, 16) + __AES_RANDOM_BYTES_N + sizeof(EncryptedBlockInfo))
+#define EncryptorAES_calcDstSize(SRC_SIZE) (__AES_IV_SIZE + sizeof(EncryptedBlockHeader) + ALIGN_TO(SRC_SIZE, 16))
 
 
 typedef struct DecryptorAES {
     br_aes_ct64_cbcdec_keys dec_ctx;
+    u8 iv[__AES_IV_SIZE];
     u8 buf[__AES_BUFFER_SIZE];
-    u8 iv[16];
 } DecryptorAES;
 
 /// @param key Array<u8, 16 | 24 | 32>
 void DecryptorAES_construct(DecryptorAES* ptr, Array(u8) key);
 
 /// @brief Decrypts `src` and writes output to `dst`.
-/// @param src array of any size
+/// @param src array of size at least EncryptorAES_calcDstSize(0). Size must be multiple of 16.
 /// @param dst array of size >= src.size
 /// @param decrypted_size size of original data without padding added by EncryptorAES_encrypt
 void DecryptorAES_decrypt(DecryptorAES* ptr, Array(u8) src, Array(u8) dst, u32* decrypted_size);