Padded struct disposal function

c-write/lib/aes.c

@@ -222,6 +222,20 @@ aes_padded_data_t aes_create_padded_data_container() {
     return out;
 }
 
+// NOTE: Only drops padding, not the original data. Resets struct to defaults and returns the original data pointer.
+uint8_t* aes_drop_padded_data_container(aes_padded_data_t* data) {
+    uint8_t* ptr = data->data;
+    free(data->padded_block);
+
+    data->padded_block = NULL;
+    data->data = NULL;
+    data->data_length = 0;
+    data->data_length_before_pad = 0;
+    data->length = 0;
+
+    return ptr;
+}
+
 // NOTE: malloc
 int aes_pad_data(aes_padded_data_t* ptr, uint8_t* data, size_t length) {
     if (length == 0) {
@@ -267,7 +281,7 @@ int aes_pad_data(aes_padded_data_t* ptr, uint8_t* data, size_t length) {
     return 0;
 }
 
-int aes_ecb(aes_encryption_mode_t mode, aes_padded_data_t* data, uint8_t* key) {
+void aes_ecb(aes_encryption_mode_t mode, aes_padded_data_t* data, uint8_t* key) {
     uint8_t* round_keys = aes__expand_key(key);
 
     for (size_t idx = 0; idx < data->data_length_before_pad; idx += BLOCK_LENGTH) {
@@ -285,10 +299,9 @@ int aes_ecb(aes_encryption_mode_t mode, aes_padded_data_t* data, uint8_t* key) {
     }
 
     free(round_keys);
-    return 0;
 }
 
-int aes_cbc(aes_encryption_mode_t mode, aes_padded_data_t* data, uint8_t* key) {
+void aes_cbc(aes_encryption_mode_t mode, aes_padded_data_t* data, uint8_t* key) {
     uint8_t* round_keys = aes__expand_key(key);
 
     // TODO: Derive this properly
@@ -320,5 +333,4 @@ int aes_cbc(aes_encryption_mode_t mode, aes_padded_data_t* data, uint8_t* key) {
 
     free(previous_block);
     free(round_keys);
-    return 0;
 }

c-write/lib/aes.h

@@ -19,15 +19,16 @@ typedef struct aes_padded_data_t {
     size_t length;
 } aes_padded_data_t;
 
-aes_padded_data_t aes_create_padded_data_container();
-
 typedef enum aes_encryption_mode_t {
     encryption,
     decryption,
 } aes_encryption_mode_t;
 
+aes_padded_data_t aes_create_padded_data_container();
+uint8_t* aes_drop_padded_data_container(aes_padded_data_t* data);
 int aes_pad_data(aes_padded_data_t* ptr, uint8_t* data, size_t length);
-int aes_ecb(aes_encryption_mode_t mode, aes_padded_data_t* data, uint8_t* key);
-int aes_cbc(aes_encryption_mode_t mode, aes_padded_data_t* data, uint8_t* key);
+
+void aes_ecb(aes_encryption_mode_t mode, aes_padded_data_t* data, uint8_t* key);
+void aes_cbc(aes_encryption_mode_t mode, aes_padded_data_t* data, uint8_t* key);
 
 #endif // AES_H_

c-write/lib/aes_tests.c

@@ -46,10 +46,10 @@ void aes_test_expand_key() {
     };
 
     uint8_t* round_keys = aes__expand_key(test_key);
-
-    assert(arrays_eq(round_keys, expected_round_keys, BLOCK_LENGTH));
-
+    int result = arrays_eq(round_keys, expected_round_keys, BLOCK_LENGTH);
     free(round_keys);
+
+    assert(result);
 }
 
 void aes_test_block_cipher() {
@@ -62,9 +62,9 @@ void aes_test_block_cipher() {
     aes__encrypt_block(test_data, round_keys);
     aes__decrypt_block(test_data, round_keys);
 
-    assert(arrays_eq(test_data, original_data, BLOCK_LENGTH));
-
     free(round_keys);
+
+    assert(arrays_eq(test_data, original_data, BLOCK_LENGTH));
 }
 
 void aes_test_padding_validation() {
@@ -108,6 +108,9 @@ void aes_test_create_padding() {
     assert(padded.data_length_before_pad == 16);
     assert(padded.length == 32);
     assert(arrays_eq(padded.padded_block, expected_padding, BLOCK_LENGTH));
+
+    uint8_t* data_ret = aes_drop_padded_data_container(&padded);
+    assert(data_ret == data);
 }
 
 // TODO: Doesn't test padded block
@@ -135,4 +138,6 @@ void aes_test_cipher() {
 
     aes_cbc(decryption, &padded, key);
     assert(arrays_eq(padded.data, original, padded.data_length_before_pad));
+
+    aes_drop_padded_data_container(&padded);
 }