4 files changed, 69 insertions, 75 deletions

diff --git a/Makefile b/Makefile
index 55af9bc..3b875af 100644
--- a/Makefile
+++ b/Makefile
@@ -2,7 +2,11 @@
 #CFLAGS       = -Wall -mmcu=atmega16 -Os -Wl,-Map,test.map
 #OBJCOPY      = avr-objcopy
 CC           = gcc
-CFLAGS       = -Wall -Os -Wl,-Map,test.map
+LD           = gcc
+CFLAGS       = -Wall -Os -c
+LDFLAGS      = -Wall -Os -Wl,-Map,test.map
+
+OBJCOPYFLAFS = -j .text -O ihex
 OBJCOPY      = objcopy
 
 # include path to AVR library
@@ -10,28 +14,27 @@ INCLUDE_PATH = /usr/lib/avr/include
 # splint static check
 SPLINT       = splint test.c aes.c -I$(INCLUDE_PATH) +charindex -unrecog
 
+default: test.elf
+
 .SILENT:
 .PHONY:  lint clean
 
+test.hex : test.elf
+	echo copy object-code to new image and format in hex
+	$(OBJCOPY) ${OBJCOPYFLAFS} $< $@
 
-rom.hex : test.out
-	# copy object-code to new image and format in hex
-	$(OBJCOPY) -j .text -O ihex test.out rom.hex
-
-test.o : test.c
-	# compiling test.c
-	$(CC) $(CFLAGS) -c test.c -o test.o
+test.o : test.c aes.h aes.o
+	echo [CC] $@
+	$(CC) $(CFLAGS) -o  $@ $<
 
-aes.o : aes.h aes.c
-	# compiling aes.c
-	$(CC) $(CFLAGS) -c aes.c -o aes.o
+aes.o : aes.c aes.h
+	echo [CC] $@
+	$(CC) $(CFLAGS) -o $@ $<
 
-test.out : aes.o test.o
-	# linking object code to binary
-	$(CC) $(CFLAGS) aes.o test.o -o test.out
+test.elf : aes.o test.o
+	echo [LD] $@
+	$(LD) $(LDFLAGS) -o $@ $^
 
-small: test.out
-	$(OBJCOPY) -j .text -O ihex test.out rom.hex
 
 clean:
 	rm -f *.OBJ *.LST *.o *.gch *.out *.hex *.map
diff --git a/aes.c b/aes.c
index c42620b..030e8b1 100644
--- a/aes.c
+++ b/aes.c
@@ -482,21 +482,16 @@ static void InvCipher(state_t *state,uint8_t*RoundKey)
 #if defined(ECB) && (ECB == 1)
 
 
-void AES_ECB_encrypt(struct AES_ctx *ctx,const uint8_t* input,  uint8_t* output)
+void AES_ECB_encrypt(struct AES_ctx *ctx,const uint8_t* buf)
 {
-  // Copy input to output, and work in-memory on output
-  memcpy(output, input, AES_BLOCKLEN);
-
   // The next function call encrypts the PlainText with the Key using AES algorithm.
-  Cipher((state_t*)output,ctx->RoundKey);
+  Cipher((state_t*)buf,ctx->RoundKey);
 }
 
-void AES_ECB_decrypt(struct AES_ctx *ctx,const uint8_t* input,  uint8_t *output)
+void AES_ECB_decrypt(struct AES_ctx *ctx,const uint8_t* buf)
 {
-  // Copy input to output, and work in-memory on output
-  memcpy(output, input, AES_BLOCKLEN);
-  
-  InvCipher((state_t*)output,ctx->RoundKey);
+  // The next function call decrypts the PlainText with the Key using AES algorithm.
+  InvCipher((state_t*)buf,ctx->RoundKey);
 }
 
 
@@ -518,35 +513,33 @@ static void XorWithIv(uint8_t* buf,uint8_t*Iv)
   }
 }
 
-void AES_CBC_encrypt_buffer(struct AES_ctx *ctx,uint8_t* output, uint8_t* input, uint32_t length)
+void AES_CBC_encrypt_buffer(struct AES_ctx *ctx,uint8_t* buf, uint32_t length)
 {
   uintptr_t i;
   uint8_t *Iv=ctx->Iv;
-  memcpy(output, input, length);
   for (i = 0; i < length; i += AES_BLOCKLEN)
   {
-    XorWithIv(output,Iv);
-    Cipher((state_t*)output,ctx->RoundKey);
-    Iv = output;
-    output += AES_BLOCKLEN;
+    XorWithIv(buf,Iv);
+    Cipher((state_t*)buf,ctx->RoundKey);
+    Iv = buf;
+    buf += AES_BLOCKLEN;
     //printf("Step %d - %d", i/16, i);
   }
   //store Iv in ctx for next call
   memcpy(ctx->Iv,Iv,AES_BLOCKLEN);
 }
 
-void AES_CBC_decrypt_buffer(struct AES_ctx *ctx, uint8_t* output, uint8_t* input, uint32_t length)
+void AES_CBC_decrypt_buffer(struct AES_ctx *ctx, uint8_t* buf,  uint32_t length)
 {
   uintptr_t i;
-  uint8_t *Iv=ctx->Iv;
-  memcpy(output, input, length);
+  uint8_t storeNextIv[AES_BLOCKLEN];
   for (i = 0; i < length; i += AES_BLOCKLEN)
   {
-    InvCipher((state_t*)output,ctx->RoundKey);
-    XorWithIv(output,Iv);
-    Iv = input; //we DO need original input stored here
-    input += AES_BLOCKLEN;
-    output += AES_BLOCKLEN;
+	memcpy(storeNextIv, buf, AES_BLOCKLEN);
+    InvCipher((state_t*)buf,ctx->RoundKey);
+    XorWithIv(buf,ctx->Iv);
+	memcpy(ctx->Iv, storeNextIv, AES_BLOCKLEN);
+    buf += AES_BLOCKLEN;
   }
 
 }
@@ -558,34 +551,35 @@ void AES_CBC_decrypt_buffer(struct AES_ctx *ctx, uint8_t* output, uint8_t* input
 #if defined(CTR) && (CTR == 1)
 
 /* Symmetrical operation: same function for encrypting as for decrypting. Note any IV/nonce should never be reused with the same key */
-void AES_CTR_xcrypt_buffer(struct AES_ctx *ctx,uint8_t* output, uint8_t* input, uint32_t length)
+void AES_CTR_xcrypt_buffer(struct AES_ctx* ctx, uint8_t* buf, uint32_t length)
 {
   uint8_t buffer[AES_BLOCKLEN];
   
-  int j;
   unsigned i;
-  for (i = 0; i < length; ++i)
+  int bi;
+  for (i = 0,bi=AES_BLOCKLEN; i < length; ++i,bi++)
   {
-    if ((i & (AES_BLOCKLEN - 1)) == 0) //we need to regen xor compliment in buff
+    if (bi == AES_BLOCKLEN) //we need to regen xor compliment in buffer
     {
       
       memcpy(buffer, ctx->Iv, AES_BLOCKLEN);
       Cipher((state_t*)buffer,ctx->RoundKey);
 
-      /* Increment counter and handle overflow */
-      for (j = (AES_BLOCKLEN - 1); j >= 0; --j)
+      /* Increment Iv and handle overflow */
+      for (bi = (AES_BLOCKLEN - 1); bi >= 0; --bi)
       {
-        ctx->Iv[j] += 1;
+        if (ctx->Iv[bi] == 255) { //inc will owerflow
+          ctx->Iv[bi]=0;
+          continue;
+        } 
+        ctx->Iv[bi] += 1;
+        break;
    
-        /* Break if no overflow, keep going otherwise */
-        if (ctx->Iv[j] != 0)
-        {
-          break;
-        }
       }
+      bi=0;
     }
 
-    output[i] = (input[i] ^ buffer[(i & (AES_BLOCKLEN - 1))]);
+    buf[i] = (buf[i] ^ buffer[bi]);
   }
 }
 
diff --git a/aes.h b/aes.h
index 021f0c2..54d994a 100644
--- a/aes.h
+++ b/aes.h
@@ -58,8 +58,8 @@ void AES_ctx_set_iv(struct AES_ctx *ctx,const uint8_t* iv);
 // buffer size is exactly AES_BLOCKLEN bytes; 
 // you need only AES_init_ctx as Iv is not used in ECB 
 // NB: ECB s considered insecure
-void AES_ECB_encrypt(struct AES_ctx *ctx, const uint8_t* input, uint8_t *output);
-void AES_ECB_decrypt(struct AES_ctx *ctx, const uint8_t* input, uint8_t *output);
+void AES_ECB_encrypt(struct AES_ctx *ctx, const uint8_t* buf);
+void AES_ECB_decrypt(struct AES_ctx *ctx, const uint8_t* buf);
 
 #endif // #if defined(ECB) && (ECB == !)
 
@@ -69,8 +69,8 @@ void AES_ECB_decrypt(struct AES_ctx *ctx, const uint8_t* input, uint8_t *output)
 // We suggest https://en.wikipedia.org/wiki/Padding_(cryptography)#PKCS7 if you need one
 // you need to set iv in ctx via AES_init_ctx_iv or AES_ctx_set_iv
 // NB: no IV should ever be reused with the same key 
-void AES_CBC_encrypt_buffer(struct AES_ctx *ctx, uint8_t* output, uint8_t* input, uint32_t length);
-void AES_CBC_decrypt_buffer(struct AES_ctx *ctx, uint8_t* output, uint8_t* input, uint32_t length);
+void AES_CBC_encrypt_buffer(struct AES_ctx *ctx, uint8_t* buf, uint32_t length);
+void AES_CBC_decrypt_buffer(struct AES_ctx *ctx, uint8_t* buf, uint32_t length);
 
 #endif // #if defined(CBC) && (CBC == 1)
 
@@ -83,7 +83,7 @@ void AES_CBC_decrypt_buffer(struct AES_ctx *ctx, uint8_t* output, uint8_t* input
 // We suggest https://en.wikipedia.org/wiki/Padding_(cryptography)#PKCS7 if you need one
 // you need to set iv in ctx via AES_init_ctx_iv or AES_ctx_set_iv
 // NB: no IV should ever be reused with the same key 
-void AES_CTR_xcrypt_buffer(struct AES_ctx *ctx, uint8_t* output, uint8_t* input, uint32_t length);
+void AES_CTR_xcrypt_buffer(struct AES_ctx* ctx, uint8_t* buf, uint32_t length);
 
 #endif // #if defined(CTR) && (CTR == 1)
 
diff --git a/test.c b/test.c
index 2bd2159..1460003 100644
--- a/test.c
+++ b/test.c
@@ -101,8 +101,8 @@ static void test_encrypt_ecb_verbose(void)
     AES_init_ctx(&ctx,key);
     for(i = 0; i < 4; ++i)
     {
-      AES_ECB_encrypt(&ctx,plain_text + (i*16),  buf+(i*16));
-        phex(buf + (i*16));
+      AES_ECB_encrypt(&ctx,plain_text + (i*16));
+        phex(plain_text + (i*16));
     }
     printf("\n");
 }
@@ -124,14 +124,13 @@ static void test_encrypt_ecb(void)
 #endif
 
   uint8_t in[]  = {0x6b, 0xc1, 0xbe, 0xe2, 0x2e, 0x40, 0x9f, 0x96, 0xe9, 0x3d, 0x7e, 0x11, 0x73, 0x93, 0x17, 0x2a};
-  uint8_t buffer[16];
 
   struct AES_ctx ctx; AES_init_ctx(&ctx,key);
-  AES_ECB_encrypt(&ctx,in, buffer);
+  AES_ECB_encrypt(&ctx,in);
 
   printf("ECB encrypt: ");
 
-  if(0 == memcmp((char*) out, (char*) buffer, 16))
+  if(0 == memcmp((char*) out, (char*) in, 16))
   {
     printf("SUCCESS!\n");
   }
@@ -169,14 +168,14 @@ static void test_decrypt_cbc(void)
                      0xae, 0x2d, 0x8a, 0x57, 0x1e, 0x03, 0xac, 0x9c, 0x9e, 0xb7, 0x6f, 0xac, 0x45, 0xaf, 0x8e, 0x51,
                      0x30, 0xc8, 0x1c, 0x46, 0xa3, 0x5c, 0xe4, 0x11, 0xe5, 0xfb, 0xc1, 0x19, 0x1a, 0x0a, 0x52, 0xef,
                      0xf6, 0x9f, 0x24, 0x45, 0xdf, 0x4f, 0x9b, 0x17, 0xad, 0x2b, 0x41, 0x7b, 0xe6, 0x6c, 0x37, 0x10 };
-  uint8_t buffer[64];
+//  uint8_t buffer[64];
   struct AES_ctx ctx;
   AES_init_ctx_iv(&ctx,key,iv);
-  AES_CBC_decrypt_buffer(&ctx,buffer, in, 64);
+  AES_CBC_decrypt_buffer(&ctx,in, 64);
 
   printf("CBC decrypt: ");
 
-  if(0 == memcmp((char*) out, (char*) buffer, 64))
+  if(0 == memcmp((char*) out, (char*) in, 64))
   {
     printf("SUCCESS!\n");
   }
@@ -214,14 +213,13 @@ static void test_encrypt_cbc(void)
                     0x30, 0xc8, 0x1c, 0x46, 0xa3, 0x5c, 0xe4, 0x11, 0xe5, 0xfb, 0xc1, 0x19, 0x1a, 0x0a, 0x52, 0xef,
                     0xf6, 0x9f, 0x24, 0x45, 0xdf, 0x4f, 0x9b, 0x17, 0xad, 0x2b, 0x41, 0x7b, 0xe6, 0x6c, 0x37, 0x10 };
 
-  uint8_t buffer[64];
   struct AES_ctx ctx;
   AES_init_ctx_iv(&ctx,key,iv);
-  AES_CBC_encrypt_buffer(&ctx,buffer, in, 64);
+  AES_CBC_encrypt_buffer(&ctx, in, 64);
 
   printf("CBC encrypt: ");
 
-  if(0 == memcmp((char*) out, (char*) buffer, 64))
+  if(0 == memcmp((char*) out, (char*) in, 64))
   {
     printf("SUCCESS!\n");
   }
@@ -270,16 +268,15 @@ static void test_xcrypt_ctr(const char* xcrypt)
                       0xae, 0x2d, 0x8a, 0x57, 0x1e, 0x03, 0xac, 0x9c, 0x9e, 0xb7, 0x6f, 0xac, 0x45, 0xaf, 0x8e, 0x51,
                       0x30, 0xc8, 0x1c, 0x46, 0xa3, 0x5c, 0xe4, 0x11, 0xe5, 0xfb, 0xc1, 0x19, 0x1a, 0x0a, 0x52, 0xef,
                       0xf6, 0x9f, 0x24, 0x45, 0xdf, 0x4f, 0x9b, 0x17, 0xad, 0x2b, 0x41, 0x7b, 0xe6, 0x6c, 0x37, 0x10 };
-  uint8_t buffer[64];
   
   struct AES_ctx ctx;
   AES_init_ctx_iv(&ctx,key,iv);
   
-  AES_CTR_xcrypt_buffer(&ctx,buffer, in, 64);
+  AES_CTR_xcrypt_buffer(&ctx, in, 64);
   
   printf("CTR %s: ", xcrypt);
   
-  if (0 == memcmp((char *) out, (char *) buffer, 64))
+  if (0 == memcmp((char *) out, (char *) in, 64))
   {
     printf("SUCCESS!\n");
   }
@@ -306,14 +303,14 @@ static void test_decrypt_ecb(void)
 #endif
 
   uint8_t out[]  = {0x6b, 0xc1, 0xbe, 0xe2, 0x2e, 0x40, 0x9f, 0x96, 0xe9, 0x3d, 0x7e, 0x11, 0x73, 0x93, 0x17, 0x2a};
-  uint8_t buffer[16];
+
   struct AES_ctx ctx;
   AES_init_ctx(&ctx,key);
-  AES_ECB_decrypt(&ctx,in, buffer);
+  AES_ECB_decrypt(&ctx,in);
 
   printf("ECB decrypt: ");
 
-  if(0 == memcmp((char*) out, (char*) buffer, 16))
+  if(0 == memcmp((char*) out, (char*) in, 16))
   {
     printf("SUCCESS!\n");
   }