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);
-    memset(ptr->iv, 0, sizeof(ptr->iv));
+    ptr->rng_ctx.vtable = &br_hmac_drbg_vtable;
+    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
-    br_hmac_drbg_generate(&ptr->rng_ctx, ptr->buf, __AES_RANDOM_BYTES_N);
-    const EncryptedBlockInfo block_info = { .padding_size = 16 - src.size % 16 };
-    // write struct after random_bytes
-    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;
+    // generate random initial vector
+    br_hmac_drbg_generate(&ptr->rng_ctx, ptr->iv, __AES_IV_SIZE);
+    // write IV to the beginning of dst
+    __Array_writeNext(&dst, ptr->iv, __AES_IV_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);
-        src.data = (u8*)src.data + __AES_BUFFER_SIZE;
-        src.size -= __AES_BUFFER_SIZE;
+        __Array_readNext(ptr->buf, &src, __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);
-        dst.data = (u8*)dst.data + __AES_BUFFER_SIZE;
-        dst.size -= __AES_BUFFER_SIZE;
+        __Array_writeNext(&dst, ptr->buf, __AES_BUFFER_SIZE);
+    }
+
+    // encrypt buffer with remaining data
+    if(src.size > 0){
+        memcpy(ptr->buf, src.data, src.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);
     }
-    
-    // write incomplete block
-    memset(ptr->buf, 0, __AES_BUFFER_SIZE);
-    memcpy(ptr->buf, src.data, src.size);
-    u32 src_size_padded = src.size + block_info.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
-    const u32 header_size = __AES_RANDOM_BYTES_N + sizeof(EncryptedBlockInfo);
-    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;
+    // read IV from the beginning of src
+    __Array_readNext(ptr->iv, &src, __AES_IV_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);
-        src.data = (u8*)src.data + __AES_BUFFER_SIZE;
-        src.size -= __AES_BUFFER_SIZE;
+        __Array_readNext(ptr->buf, &src, __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);
-        dst.data = (u8*)dst.data + __AES_BUFFER_SIZE;
-        dst.size -= __AES_BUFFER_SIZE;
+        __Array_writeNext(&dst, ptr->buf, __AES_BUFFER_SIZE);
     }
     
-    // write incomplete block
-    memset(ptr->buf, 0, __AES_BUFFER_SIZE);
-    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);
+    // decrypt buffer with remaining data
+    if(src.size > 0){
+        memcpy(ptr->buf, src.data, src.size);
+        br_aes_ct64_cbcdec_run(&ptr->dec_ctx, ptr->iv, ptr->buf, src_size_padded);
+        memcpy(dst.data, ptr->buf, src.size);
+    }
 }