Merge pull request #29 from jacksonwalters/add_pkcs7_padding

Add PKCS#7 padding

Author: jacksonwalters

include/ecb.h

@@ -12,4 +12,12 @@ void aes_ecb_encrypt(const uint8_t *plaintext, uint8_t *ciphertext,
 void aes_ecb_decrypt(const uint8_t *ciphertext, uint8_t *plaintext,
                      size_t length, const void *ctx);
 
+int aes_ecb_encrypt_padded(const uint8_t *in, size_t in_len,
+                           const void *ctx,
+                           uint8_t **out, size_t *out_len);
+
+int aes_ecb_decrypt_padded(const uint8_t *in, size_t in_len,
+                           const void *ctx,
+                           uint8_t **out, size_t *out_len);
+
 #endif

include/padding.h

@@ -0,0 +1,23 @@
+#ifndef PADDING_H
+#define PADDING_H
+
+#include <stddef.h>
+#include <stdint.h>
+
+#define PADDING_OK 0
+#define PADDING_ERR_INVALID -1
+#define PADDING_ERR_NOMEM -2
+
+/* Pad `in` (length in_len) to multiple of block_size.
+ * Allocates new buffer in *out (caller frees it).
+ */
+int pkcs7_pad(const uint8_t *in, size_t in_len, size_t block_size,
+              uint8_t **out, size_t *out_len);
+
+/* Remove PKCS#7 padding from `in` (length in_len),
+ * allocates unpadded buffer in *out (caller frees it).
+ */
+int pkcs7_unpad(const uint8_t *in, size_t in_len, size_t block_size,
+                uint8_t **out, size_t *out_len);
+
+#endif /* PADDING_H */

src/cbc.c

@@ -1,72 +1,45 @@
 #include "cbc.h"
+#include "padding.h"
 #include <string.h>
 #include <stdlib.h>
 
-/* PKCS#7 padding helper: returns number of padding bytes appended (1..16) */
-static size_t pkcs7_pad(const uint8_t *in, size_t in_len, uint8_t *out) {
-    size_t pad_len = AES_BLOCK - (in_len % AES_BLOCK);
-    if (pad_len == 0) pad_len = AES_BLOCK;
-    /* copy input */
-    memcpy(out, in, in_len);
-    /* append pad bytes */
-    for (size_t i = 0; i < pad_len; ++i) out[in_len + i] = (uint8_t)pad_len;
-    return pad_len;
-}
-
-/* remove PKCS#7 pad; returns 0 on success and sets *out_len to unpadded length.
- * Returns non-zero on invalid padding.
- */
-static int pkcs7_unpad(uint8_t *buf, size_t buf_len, size_t *out_len) {
-    if (buf_len == 0 || (buf_len % AES_BLOCK) != 0) return -1;
-    uint8_t pad = buf[buf_len - 1];
-    if (pad == 0 || pad > AES_BLOCK) return -2;
-    /* check that last pad bytes equal pad */
-    for (size_t i = 0; i < pad; ++i) {
-        if (buf[buf_len - 1 - i] != pad) return -3;
-    }
-    *out_len = buf_len - pad;
-    return 0;
-}
-
 /* XOR helper */
-static inline void xor_block(uint8_t out[AES_BLOCK], const uint8_t a[AES_BLOCK], const uint8_t b[AES_BLOCK]) {
-    for (int i = 0; i < AES_BLOCK; ++i) out[i] = a[i] ^ b[i];
+static inline void xor_block(uint8_t out[AES_BLOCK],
+                             const uint8_t a[AES_BLOCK],
+                             const uint8_t b[AES_BLOCK])
+{
+    for (int i = 0; i < AES_BLOCK; ++i)
+        out[i] = a[i] ^ b[i];
 }
 
 int aes_cbc_encrypt(const uint8_t *in, size_t in_len,
                     uint8_t *out, size_t *out_len,
                     const uint8_t iv[AES_BLOCK],
                     encrypt_block_fn encrypt, const void *ctx)
 {
-    if (!in || !out || !out_len || !iv || !encrypt) return -1;
-
-    /* padded buffer size = ceil(in_len/16)*16 + 16 (if in_len %16 ==0, add a full block) */
-    size_t padded_len = ((in_len + AES_BLOCK - 1) / AES_BLOCK) * AES_BLOCK;
-    if (padded_len == in_len) padded_len += AES_BLOCK;
-
-    uint8_t *buf = (uint8_t*)malloc(padded_len);
-    if (!buf) return -2;
+    if (!out || !out_len || !iv || !encrypt) return -1;
+    if (!in && in_len > 0) return -1; // NULL only allowed if in_len==0
+    if (!in) in = (const uint8_t *)""; // empty input is valid
 
-    /* create padded plaintext in buf */
-    size_t pad_len = pkcs7_pad(in, in_len, buf);
-    (void)pad_len; /* padded_len equals in_len + pad_len */
+    uint8_t *padded = NULL;
+    size_t padded_len = 0;
+    if (pkcs7_pad(in, in_len, AES_BLOCK, &padded, &padded_len) != PADDING_OK)
+        return -2;
 
     uint8_t prev[AES_BLOCK];
     memcpy(prev, iv, AES_BLOCK);
 
     for (size_t off = 0; off < padded_len; off += AES_BLOCK) {
         uint8_t block[AES_BLOCK];
-        xor_block(block, buf + off, prev);
+        xor_block(block, padded + off, prev);
         encrypt(block, out + off, ctx);
-        /* new prev = ciphertext block */
         memcpy(prev, out + off, AES_BLOCK);
     }
 
     *out_len = padded_len;
-    /* wipe sensitive buffer */
-    memset(buf, 0, padded_len);
-    free(buf);
     memset(prev, 0, AES_BLOCK);
+    memset(padded, 0, padded_len);
+    free(padded);
     return 0;
 }
 
@@ -83,22 +56,24 @@ int aes_cbc_decrypt(const uint8_t *in, size_t in_len,
 
     for (size_t off = 0; off < in_len; off += AES_BLOCK) {
         uint8_t tmp[AES_BLOCK];
-        decrypt(in + off, tmp, ctx); /* tmp = AES_DEC(Ci) */
-        xor_block(out + off, tmp, prev); /* plaintext block = tmp XOR prev */
+        decrypt(in + off, tmp, ctx);
+        xor_block(out + off, tmp, prev);
         memcpy(prev, in + off, AES_BLOCK);
     }
 
-    /* unpad in-place on out */
-    size_t unpadded_len = 0;
-    int r = pkcs7_unpad(out, in_len, &unpadded_len);
-    if (r != 0) {
-        /* wipe and return error */
+    /* Unpad using reusable PKCS#7 function */
+    uint8_t *unpadded = NULL;
+    size_t plain_len = 0;
+    if (pkcs7_unpad(out, in_len, AES_BLOCK, &unpadded, &plain_len) != PADDING_OK) {
         memset(out, 0, in_len);
         memset(prev, 0, AES_BLOCK);
         return -3;
     }
 
-    *out_len = unpadded_len;
+    memcpy(out, unpadded, plain_len);
+    *out_len = plain_len;
+    memset(unpadded, 0, plain_len);
+    free(unpadded);
     memset(prev, 0, AES_BLOCK);
     return 0;
 }

src/ecb.c

@@ -1,5 +1,7 @@
+#include <stdlib.h>
 #include <string.h>
 #include "ecb.h"
+#include "padding.h"
 
 void aes_ecb_encrypt(const uint8_t *plaintext, uint8_t *ciphertext,
                      size_t length, const void *ctx)
@@ -18,3 +20,53 @@ void aes_ecb_decrypt(const uint8_t *ciphertext, uint8_t *plaintext,
         aes_block_wrapper_dec(ciphertext + i*16, plaintext + i*16, ctx);
     }
 }
+
+int aes_ecb_encrypt_padded(const uint8_t *in, size_t in_len,
+                           const void *ctx,
+                           uint8_t **out, size_t *out_len)
+{
+    if (!out || !out_len || !ctx) return -1;      // mandatory pointers
+    if (!in && in_len > 0) return -1;            // NULL only allowed if in_len>0
+    if (!in) in = (const uint8_t *)"";          // empty input is valid
+
+    uint8_t *padded = NULL;
+    size_t padded_len = 0;
+    if (pkcs7_pad(in, in_len, 16, &padded, &padded_len) != 0) return -2;
+
+    uint8_t *cipher = (uint8_t *)malloc(padded_len);
+    if (!cipher) { free(padded); return -3; }
+
+    for (size_t i = 0; i < padded_len; i += 16)
+        aes_block_wrapper(padded + i, cipher + i, ctx);
+
+    free(padded);
+    *out = cipher;
+    *out_len = padded_len;
+    return 0;
+}
+
+int aes_ecb_decrypt_padded(const uint8_t *in, size_t in_len,
+                           const void *ctx,
+                           uint8_t **out, size_t *out_len)
+{
+    if (!in || !ctx || !out || !out_len || (in_len % 16) != 0) return -1;
+
+    uint8_t *decrypted = (uint8_t *)malloc(in_len);
+    if (!decrypted) return -2;
+
+    for (size_t i = 0; i < in_len; i += 16)
+        aes_block_wrapper_dec(in + i, decrypted + i, ctx);
+
+    uint8_t *unpadded = NULL;
+    size_t plain_len = 0;
+    if (pkcs7_unpad(decrypted, in_len, 16, &unpadded, &plain_len) != 0) {
+        free(decrypted);
+        return -3;
+    }
+
+    free(decrypted);
+    *out = unpadded;
+    *out_len = plain_len;
+    return 0;
+}
+

src/padding.c

@@ -0,0 +1,59 @@
+#include "padding.h"
+#include <stdlib.h>
+#include <string.h>
+
+/* Constant-time check: all bytes in buf[0..n-1] == val */
+static int ct_mem_is_value(const uint8_t *buf, size_t n, uint8_t val) {
+    uint8_t diff = 0;
+    for (size_t i = 0; i < n; ++i) {
+        diff |= buf[i] ^ val;
+    }
+    return diff == 0;
+}
+
+int pkcs7_pad(const uint8_t *in, size_t in_len, size_t block_size,
+              uint8_t **out, size_t *out_len)
+{
+    if (!in || !out || !out_len || block_size == 0 || block_size > 255)
+        return PADDING_ERR_INVALID;
+
+    size_t pad_len = block_size - (in_len % block_size);
+    if (pad_len == 0) pad_len = block_size; /* full block padding */
+
+    size_t total = in_len + pad_len;
+    uint8_t *buf = (uint8_t *)malloc(total);
+    if (!buf) return PADDING_ERR_NOMEM;
+
+    if (in_len > 0) memcpy(buf, in, in_len);
+    memset(buf + in_len, (uint8_t)pad_len, pad_len);
+
+    *out = buf;
+    *out_len = total;
+    return PADDING_OK;
+}
+
+int pkcs7_unpad(const uint8_t *in, size_t in_len, size_t block_size,
+                uint8_t **out, size_t *out_len)
+{
+    if (!in || !out || !out_len || block_size == 0 || block_size > 255)
+        return PADDING_ERR_INVALID;
+    if (in_len == 0 || (in_len % block_size) != 0)
+        return PADDING_ERR_INVALID;
+
+    uint8_t pad_val = in[in_len - 1];
+    if (pad_val == 0 || pad_val > block_size)
+        return PADDING_ERR_INVALID;
+
+    if (!ct_mem_is_value(in + (in_len - pad_val), pad_val, pad_val))
+        return PADDING_ERR_INVALID;
+
+    size_t plain_len = in_len - pad_val;
+    uint8_t *buf = (uint8_t *)malloc(plain_len ? plain_len : 1);
+    if (!buf) return PADDING_ERR_NOMEM;
+
+    if (plain_len > 0) memcpy(buf, in, plain_len);
+
+    *out = buf;
+    *out_len = plain_len;
+    return PADDING_OK;
+}

tests/test_ecb.c

@@ -1,15 +1,16 @@
 #include <stdio.h>
 #include <stdint.h>
 #include <string.h>
+#include <stdlib.h>
 
 #include "../include/ecb.h"
 #include "../include/aes_wrapper.h"
 #include "../include/aes_128.h"
 #include "../include/key_expansion_128.h"
 #include "../include/sbox.h"
+#include "../include/padding.h"
 
-int main(void) {
-    // NIST SP 800-38A AES-128 ECB test vector
+static int test_nist_vector(void) {
     const uint8_t key[16] = {
         0x2b,0x7e,0x15,0x16,0x28,0xae,0xd2,0xa6,
         0xab,0xf7,0x15,0x88,0x09,0xcf,0x4f,0x3c
@@ -29,38 +30,66 @@ int main(void) {
     uint8_t sbox[256];
     initialize_aes_sbox(sbox);
     aes_key_expansion_128(key, round_keys, sbox);
-
-    struct aes_ctx ctx = { .round_keys = round_keys, .sbox = sbox, .key_len = 16};
+    struct aes_ctx ctx = { .round_keys = round_keys, .sbox = sbox, .key_len = 16 };
 
     uint8_t ciphertext[16] = {0};
     uint8_t decrypted[16] = {0};
 
-    // Encrypt
     aes_ecb_encrypt(plaintext, ciphertext, 16, &ctx);
-
-    // Decrypt
     aes_ecb_decrypt(ciphertext, decrypted, 16, &ctx);
 
-    // Verify encryption matches NIST vector
-    if (memcmp(ciphertext, expected_ciphertext, 16) != 0) {
-        fprintf(stderr, "ECB encryption FAILED\n");
-        return 1;
-    }
+    if (memcmp(ciphertext, expected_ciphertext, 16) != 0) return 1;
+    if (memcmp(decrypted, plaintext, 16) != 0) return 2;
+    return 0;
+}
+
+static int test_partial_block(void) {
+    const uint8_t key[16] = {0};
+    uint8_t round_keys[176], sbox[256];
+    initialize_aes_sbox(sbox);
+    aes_key_expansion_128(key, round_keys, sbox);
+    struct aes_ctx ctx = { .round_keys = round_keys, .sbox = sbox, .key_len = 16 };
+
+    uint8_t plaintext[5] = {1,2,3,4,5};
+    uint8_t *cipher = NULL, *plain = NULL;
+    size_t cipher_len = 0, plain_len = 0;
+
+    if (aes_ecb_encrypt_padded(plaintext, 5, &ctx, &cipher, &cipher_len) != 0) return 1;
+    if (aes_ecb_decrypt_padded(cipher, cipher_len, &ctx, &plain, &plain_len) != 0) return 2;
+
+    if (plain_len != 5) return 3;
+    if (memcmp(plain, plaintext, 5) != 0) return 4;
 
-    // Verify decryption matches plaintext
-    if (memcmp(decrypted, plaintext, 16) != 0) {
-        fprintf(stderr, "ECB decryption FAILED\n");
-        return 2;
-    }
+    free(cipher);
+    free(plain);
+    return 0;
+}
+
+static int test_empty(void) {
+    const uint8_t key[16] = {0};
+    uint8_t round_keys[176], sbox[256];
+    initialize_aes_sbox(sbox);
+    aes_key_expansion_128(key, round_keys, sbox);
+    struct aes_ctx ctx = { .round_keys = round_keys, .sbox = sbox, .key_len = 16 };
+
+    uint8_t *cipher = NULL, *plain = NULL;
+    size_t cipher_len = 0, plain_len = 0;
 
-    printf("ECB round-trip OK — plaintext recovered correctly.\n");
-    printf("Plaintext: ");
-    for (size_t i = 0; i < 16; ++i) printf("%02x", plaintext[i]);
-    printf("\n");
+    if (aes_ecb_encrypt_padded(NULL, 0, &ctx, &cipher, &cipher_len) != 0) return 1;
+    if (aes_ecb_decrypt_padded(cipher, cipher_len, &ctx, &plain, &plain_len) != 0) return 2;
 
-    printf("Ciphertext (hex): ");
-    for (size_t i = 0; i < 16; ++i) printf("%02x", ciphertext[i]);
-    printf("\n");
+    if (plain_len != 0) return 3;
+
+    free(cipher);
+    free(plain);
+    return 0;
+}
+
+int main(void) {
+    if (test_nist_vector() != 0) { fprintf(stderr, "NIST ECB test FAILED\n"); return 1; }
+    if (test_partial_block() != 0) { fprintf(stderr, "Partial-block ECB test FAILED\n"); return 2; }
+    if (test_empty() != 0) { fprintf(stderr, "Empty ECB test FAILED\n"); return 3; }
 
+    printf("All ECB tests passed.\n");
     return 0;
 }