1 files changed, 46 insertions, 39 deletions

diff --git a/aes_cbc.c b/aes_cbc.c
index 59df8aa..21c8430 100644
--- a/aes_cbc.c
+++ b/aes_cbc.c
@@ -1,6 +1,6 @@
 /*
 
-This is an implementation of the AES128 algorithm, specifically ECB mode.
+This is an implementation of the AES128 algorithm, specifically ECB and CBC mode.
 
 The implementation is verified against the test vectors in:
   National Institute of Standards and Technology Special Publication 800-38A 2001 ED
@@ -46,7 +46,7 @@ NOTE:   String length must be evenly divisible by 16byte (str_len % 16 == 0)
 // The number of 32 bit words in a key.
 #define Nk 4
 // Key length in bytes [128 bit]
-#define keyln 16
+#define KEYLEN 16
 // The number of rounds in AES Cipher.
 #define Nr 10
 
@@ -57,13 +57,6 @@ NOTE:   String length must be evenly divisible by 16byte (str_len % 16 == 0)
   #define MULTIPLY_AS_A_FUNCTION 0
 #endif
 
-#ifndef CBC
-  #define CBC 1
-#endif
-
-#ifndef ECB
-  #define ECB 0
-#endif
 
 /*****************************************************************************/
 /* Private variables:                                                        */
@@ -439,10 +432,10 @@ static void InvCipher(void)
   AddRoundKey(0);
 }
 
-static void BufferCopy(uint8_t* output, uint8_t* input)
+static void BlockCopy(uint8_t* output, uint8_t* input)
 {
   uint8_t i;
-  for (i=0;i<16;++i)
+  for (i=0;i<KEYLEN;++i)
   {
     output[i] = input[i];
   }
@@ -459,7 +452,7 @@ static void BufferCopy(uint8_t* output, uint8_t* input)
 void AES128_ECB_encrypt(uint8_t* input, const uint8_t* key, uint8_t* output)
 {
   // Copy input to output, and work in-memory on output
-  BufferCopy(output, input);
+  BlockCopy(output, input);
   state = (state_t*)output;
 
   Key = key;
@@ -472,7 +465,7 @@ void AES128_ECB_encrypt(uint8_t* input, const uint8_t* key, uint8_t* output)
 void AES128_ECB_decrypt(uint8_t* input, const uint8_t* key, uint8_t *output)
 {
   // Copy input to output, and work in-memory on output
-  BufferCopy(output, input);
+  BlockCopy(output, input);
   state = (state_t*)output;
 
   // The KeyExpansion routine must be called before encryption.
@@ -495,74 +488,88 @@ void AES128_ECB_decrypt(uint8_t* input, const uint8_t* key, uint8_t *output)
 static void XorWithIv(uint8_t* buf)
 {
   uint8_t i;
-  for(i = 0; i < 16; ++i)
+  for(i = 0; i < KEYLEN; ++i)
   {
     buf[i] ^= Iv[i];
   }
 }
 
-void AES128_CBC_encrypt_buffer(uint8_t* input, uint32_t length, const uint8_t* key, uint8_t* output, const uint8_t* iv)
+void AES128_CBC_encrypt_buffer(uint8_t* output, uint8_t* input, uint32_t length, const uint8_t* key, const uint8_t* iv)
 {
   intptr_t i;
-  uint8_t j;
-  uint8_t remainders = length % 16; /* Remaining bytes in the last non-full block */
+  uint8_t remainders = length % KEYLEN; /* Remaining bytes in the last non-full block */
 
-  BufferCopy(output, input);
+  BlockCopy(output, input);
   state = (state_t*)output;
 
-  Key = key;
-  KeyExpansion();
+  // Skip the key expansion if key is passed as 0
+  if(0 != key)
+  {
+    Key = key;
+    KeyExpansion();
+  }
 
-  Iv = (uint8_t*)iv;
+  if(iv != 0)
+  {
+    Iv = (uint8_t*)iv;
+  }
 
-  for(i = 0; i < length; i += 16)
+  for(i = 0; i < length; i += KEYLEN)
   {
     XorWithIv(input);
-    BufferCopy(output, input);
+    BlockCopy(output, input);
     state = (state_t*)output;
     Cipher();
     Iv = output;
-    input += 16;
-    output += 16;
+    input += KEYLEN;
+    output += KEYLEN;
   }
 
   if(remainders)
   {
-    BufferCopy(output, input);
-    memset(output + remainders, 0, 16 - remainders); /* add 0-padding */
+    BlockCopy(output, input);
+    memset(output + remainders, 0, KEYLEN - remainders); /* add 0-padding */
     state = (state_t*)output;
     Cipher();
   }
 }
 
-void AES128_CBC_decrypt_buffer(uint8_t* input, uint32_t length, const uint8_t* key, uint8_t* output, const uint8_t* iv)
+void AES128_CBC_decrypt_buffer(uint8_t* output, uint8_t* input, uint32_t length, const uint8_t* key, const uint8_t* iv)
 {
   intptr_t i;
-  uint8_t remainders = length % 16; /* Remaining bytes in the last non-full block */
+  uint8_t remainders = length % KEYLEN; /* Remaining bytes in the last non-full block */
   
-  BufferCopy(output, input);
+  BlockCopy(output, input);
   state = (state_t*)output;
 
-  Key = key;
-  KeyExpansion();
+  // Skip the key expansion if key is passed as 0
+  if(0 != key)
+  {
+    Key = key;
+    KeyExpansion();
+  }
 
-  Iv = (uint8_t*)iv;
+  // If iv is passed as 0, we continue to encrypt without re-setting the Iv
+  if(iv != 0)
+  {
+    Iv = (uint8_t*)iv;
+  }
 
-  for(i = 0; i < length; i += 16)
+  for(i = 0; i < length; i += KEYLEN)
   {
-    BufferCopy(output, input);
+    BlockCopy(output, input);
     state = (state_t*)output;
     InvCipher();
     XorWithIv(output);
     Iv = input;
-    input += 16;
-    output += 16;
+    input += KEYLEN;
+    output += KEYLEN;
   }
 
   if(remainders)
   {
-    BufferCopy(output, input);
-    memset(output+remainders, 0, 16-remainders); /* add 0-padding */
+    BlockCopy(output, input);
+    memset(output+remainders, 0, KEYLEN - remainders); /* add 0-padding */
     state = (state_t*)output;
     InvCipher();
   }