5 files changed, 212 insertions, 208 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/README.md b/README.md
index ebef638..704b2a7 100644
--- a/README.md
+++ b/README.md
@@ -7,16 +7,27 @@ You can override the default key-size of 128 bit with 192 or 256 bit by defining
 The API is very simple and looks like this (I am using C99 `<stdint.h>`-style annotated types):
 
 ```C
-void AES_ECB_encrypt(uint8_t* input, const uint8_t* key, uint8_t* output);
-void AES_ECB_decrypt(uint8_t* input, const uint8_t* key, uint8_t* output);
+//Init ctx with
+void AES_init_ctx(struct AES_ctx *ctx,const uint8_t* key);
+void AES_init_ctx_iv(struct AES_ctx *ctx,const uint8_t* key,const uint8_t* iv);
 
-void AES_CBC_encrypt_buffer(uint8_t* output, uint8_t* input, uint32_t length, const uint8_t* key, const uint8_t* iv);
-void AES_CBC_decrypt_buffer(uint8_t* output, uint8_t* input, uint32_t length, const uint8_t* key, const uint8_t* iv);
+//or reset iv at random point
+void AES_ctx_set_iv(struct AES_ctx *ctx,const uint8_t* iv);
 
-/* Same function for encrypting as for decrypting. Note no IV/nonce should ever be reused with the same key */
-void AES_CTR_xcrypt_buffer(uint8_t* output, uint8_t* input, uint32_t length, const uint8_t* key, const uint8_t* nonce);
+//then do 
+void AES_ECB_encrypt(struct AES_ctx *ctx, const uint8_t* buf);
+void AES_ECB_decrypt(struct AES_ctx *ctx, const uint8_t* buf);
+
+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);
+
+/* Same function for encrypting as for decrypting in CTR mode */
+void AES_CTR_xcrypt_buffer(struct AES_ctx* ctx, uint8_t* buf, uint32_t length);
 ```
 
+Note: 
+ * We don't provide any padding so all buffers should be mutiple of 16 bytes if you need padding we rocomend https://en.wikipedia.org/wiki/Padding_(cryptography)#PKCS7
+ * ECB mode is considered unsafe and is not implemented in streaming mode. If you realy need this just call the function for every block of 16 bytes you need encrypted. See https://en.wikipedia.org/wiki/Block_cipher_mode_of_operation#Electronic_Codebook_(ECB) for more details
 
 You can choose to use any or all of the modes-of-operations, by defining the symbols CBC, CTR or ECB. See the header file for clarification.
 
diff --git a/aes.c b/aes.c
index 0b4e764..9dc3aba 100644
--- a/aes.c
+++ b/aes.c
@@ -44,23 +44,16 @@ NOTE:   String length must be evenly divisible by 16byte (str_len % 16 == 0)
 /*****************************************************************************/
 // The number of columns comprising a state in AES. This is a constant in AES. Value=4
 #define Nb 4
-#define BLOCKLEN 16 // Block length in bytes AES is 128b block only
 
 #if defined(AES256) && (AES256 == 1)
     #define Nk 8
-    #define KEYLEN 32
     #define Nr 14
-    #define keyExpSize 240
 #elif defined(AES192) && (AES192 == 1)
     #define Nk 6
-    #define KEYLEN 24
     #define Nr 12
-    #define keyExpSize 208
 #else
     #define Nk 4        // The number of 32 bit words in a key.
-    #define KEYLEN 16   // Key length in bytes
     #define Nr 10       // The number of rounds in AES Cipher.
-    #define keyExpSize 176
 #endif
 
 // jcallan@github points out that declaring Multiply as a function 
@@ -71,23 +64,17 @@ NOTE:   String length must be evenly divisible by 16byte (str_len % 16 == 0)
 #endif
 
 
+
+
+
 /*****************************************************************************/
 /* Private variables:                                                        */
 /*****************************************************************************/
 // state - array holding the intermediate results during decryption.
 typedef uint8_t state_t[4][4];
-static state_t* state;
 
-// The array that stores the round keys.
-static uint8_t RoundKey[keyExpSize];
 
-// The Key input to the AES Program
-static const uint8_t* Key;
 
-#if defined(CBC) && CBC
-  // Initial Vector used only for CBC mode
-  static uint8_t* Iv;
-#endif
 
 // The lookup-tables are marked const so they can be placed in read-only storage instead of RAM
 // The numbers below can be computed dynamically trading ROM for RAM - 
@@ -148,23 +135,30 @@ static const uint8_t Rcon[11] = {
 /*****************************************************************************/
 /* Private functions:                                                        */
 /*****************************************************************************/
+/*
 static uint8_t getSBoxValue(uint8_t num)
 {
   return sbox[num];
 }
-
+*/
+#define getSBoxValue(num) (sbox[(num)])
+/*
 static uint8_t getSBoxInvert(uint8_t num)
 {
   return rsbox[num];
 }
+*/
+#define getSBoxInvert(num) (rsbox[(num)])
 
 // This function produces Nb(Nr+1) round keys. The round keys are used in each round to decrypt the states. 
-static void KeyExpansion(void)
+static void KeyExpansion(uint8_t* RoundKey,const uint8_t* Key)
 {
-  uint32_t i, k;
+  unsigned i, j, k;
   uint8_t tempa[4]; // Used for the column/row operations
   
   // The first round key is the key itself.
+  memcpy(RoundKey,Key,AES_keyExpSize);
+  /*
   for (i = 0; i < Nk; ++i)
   {
     RoundKey[(i * 4) + 0] = Key[(i * 4) + 0];
@@ -172,16 +166,18 @@ static void KeyExpansion(void)
     RoundKey[(i * 4) + 2] = Key[(i * 4) + 2];
     RoundKey[(i * 4) + 3] = Key[(i * 4) + 3];
   }
+  */
 
   // All other round keys are found from the previous round keys.
-  //i == Nk
-  for (; i < Nb * (Nr + 1); ++i)
+  for (i = Nk; i < Nb * (Nr + 1); ++i)
   {
     {
-      tempa[0] = RoundKey[(i-1) * 4 + 0];
-      tempa[1] = RoundKey[(i-1) * 4 + 1];
-      tempa[2] = RoundKey[(i-1) * 4 + 2];
-      tempa[3] = RoundKey[(i-1) * 4 + 3];
+      k=(i-1) * 4;
+      tempa[0]=RoundKey[k + 0];
+      tempa[1]=RoundKey[k + 1];
+      tempa[2]=RoundKey[k + 2];
+      tempa[3]=RoundKey[k + 3];
+
     }
 
     if (i % Nk == 0)
@@ -223,16 +219,30 @@ static void KeyExpansion(void)
       }
     }
 #endif
-    RoundKey[i * 4 + 0] = RoundKey[(i - Nk) * 4 + 0] ^ tempa[0];
-    RoundKey[i * 4 + 1] = RoundKey[(i - Nk) * 4 + 1] ^ tempa[1];
-    RoundKey[i * 4 + 2] = RoundKey[(i - Nk) * 4 + 2] ^ tempa[2];
-    RoundKey[i * 4 + 3] = RoundKey[(i - Nk) * 4 + 3] ^ tempa[3];
+    j=i * 4; k=(i - Nk) * 4;
+    RoundKey[j + 0] = RoundKey[k + 0] ^ tempa[0];
+    RoundKey[j + 1] = RoundKey[k + 1] ^ tempa[1];
+    RoundKey[j + 2] = RoundKey[k + 2] ^ tempa[2];
+    RoundKey[j + 3] = RoundKey[k + 3] ^ tempa[3];
   }
 }
 
+void AES_init_ctx(struct AES_ctx *ctx,const uint8_t* key){
+  KeyExpansion(ctx->RoundKey,key);
+}
+#if defined(CBC) && (CBC == 1)
+void AES_init_ctx_iv(struct AES_ctx *ctx,const uint8_t* key,const uint8_t* iv){
+  KeyExpansion(ctx->RoundKey,key);
+  memcpy (ctx->Iv,iv,AES_BLOCKLEN);
+}
+void AES_ctx_set_iv(struct AES_ctx *ctx,const uint8_t* iv) {
+  memcpy (ctx->Iv,iv,AES_BLOCKLEN);
+}
+#endif
+
 // This function adds the round key to state.
 // The round key is added to the state by an XOR function.
-static void AddRoundKey(uint8_t round)
+static void AddRoundKey(uint8_t round,state_t *state,uint8_t*RoundKey)
 {
   uint8_t i,j;
   for (i=0;i<4;++i)
@@ -246,7 +256,7 @@ static void AddRoundKey(uint8_t round)
 
 // The SubBytes Function Substitutes the values in the
 // state matrix with values in an S-box.
-static void SubBytes(void)
+static void SubBytes(state_t *state)
 {
   uint8_t i, j;
   for (i = 0; i < 4; ++i)
@@ -261,7 +271,7 @@ static void SubBytes(void)
 // The ShiftRows() function shifts the rows in the state to the left.
 // Each row is shifted with different offset.
 // Offset = Row number. So the first row is not shifted.
-static void ShiftRows(void)
+static void ShiftRows(state_t *state)
 {
   uint8_t temp;
 
@@ -295,7 +305,7 @@ static uint8_t xtime(uint8_t x)
 }
 
 // MixColumns function mixes the columns of the state matrix
-static void MixColumns(void)
+static void MixColumns(state_t *state)
 {
   uint8_t i;
   uint8_t Tmp,Tm,t;
@@ -333,7 +343,7 @@ static uint8_t Multiply(uint8_t x, uint8_t y)
 // MixColumns function mixes the columns of the state matrix.
 // The method used to multiply may be difficult to understand for the inexperienced.
 // Please use the references to gain more information.
-static void InvMixColumns(void)
+static void InvMixColumns(state_t *state)
 {
   int i;
   uint8_t a, b, c, d;
@@ -354,7 +364,7 @@ static void InvMixColumns(void)
 
 // The SubBytes Function Substitutes the values in the
 // state matrix with values in an S-box.
-static void InvSubBytes(void)
+static void InvSubBytes(state_t *state)
 {
   uint8_t i,j;
   for (i = 0; i < 4; ++i)
@@ -366,7 +376,7 @@ static void InvSubBytes(void)
   }
 }
 
-static void InvShiftRows(void)
+static void InvShiftRows(state_t *state)
 {
   uint8_t temp;
 
@@ -396,54 +406,54 @@ static void InvShiftRows(void)
 
 
 // Cipher is the main function that encrypts the PlainText.
-static void Cipher(void)
+static void Cipher(state_t *state,uint8_t*RoundKey)
 {
   uint8_t round = 0;
 
   // Add the First round key to the state before starting the rounds.
-  AddRoundKey(0); 
+  AddRoundKey(0,state,RoundKey); 
   
   // There will be Nr rounds.
   // The first Nr-1 rounds are identical.
   // These Nr-1 rounds are executed in the loop below.
   for (round = 1; round < Nr; ++round)
   {
-    SubBytes();
-    ShiftRows();
-    MixColumns();
-    AddRoundKey(round);
+    SubBytes(state);
+    ShiftRows(state);
+    MixColumns(state);
+    AddRoundKey(round,state,RoundKey);
   }
   
   // The last round is given below.
   // The MixColumns function is not here in the last round.
-  SubBytes();
-  ShiftRows();
-  AddRoundKey(Nr);
+  SubBytes(state);
+  ShiftRows(state);
+  AddRoundKey(Nr,state,RoundKey);
 }
 
-static void InvCipher(void)
+static void InvCipher(state_t *state,uint8_t*RoundKey)
 {
   uint8_t round=0;
 
   // Add the First round key to the state before starting the rounds.
-  AddRoundKey(Nr); 
+  AddRoundKey(Nr,state,RoundKey); 
 
   // There will be Nr rounds.
   // The first Nr-1 rounds are identical.
   // These Nr-1 rounds are executed in the loop below.
   for (round = (Nr - 1); round > 0; --round)
   {
-    InvShiftRows();
-    InvSubBytes();
-    AddRoundKey(round);
-    InvMixColumns();
+    InvShiftRows(state);
+    InvSubBytes(state);
+    AddRoundKey(round,state,RoundKey);
+    InvMixColumns(state);
   }
   
   // The last round is given below.
   // The MixColumns function is not here in the last round.
-  InvShiftRows();
-  InvSubBytes();
-  AddRoundKey(0);
+  InvShiftRows(state);
+  InvSubBytes(state);
+  AddRoundKey(0,state,RoundKey);
 }
 
 
@@ -453,30 +463,16 @@ static void InvCipher(void)
 #if defined(ECB) && (ECB == 1)
 
 
-void AES_ECB_encrypt(const uint8_t* input, const uint8_t* key, uint8_t* output, const uint32_t length)
+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, length);
-  state = (state_t*)output;
-
-  Key = key;
-  KeyExpansion();
-
   // The next function call encrypts the PlainText with the Key using AES algorithm.
-  Cipher();
+  Cipher((state_t*)buf,ctx->RoundKey);
 }
 
-void AES_ECB_decrypt(const uint8_t* input, const uint8_t* key, uint8_t *output, const uint32_t length)
+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, length);
-  state = (state_t*)output;
-
-  // The KeyExpansion routine must be called before encryption.
-  Key = key;
-  KeyExpansion();
-
-  InvCipher();
+  // The next function call decrypts the PlainText with the Key using AES algorithm.
+  InvCipher((state_t*)buf,ctx->RoundKey);
 }
 
 
@@ -489,89 +485,44 @@ void AES_ECB_decrypt(const uint8_t* input, const uint8_t* key, uint8_t *output,
 #if defined(CBC) && (CBC == 1)
 
 
-static void XorWithIv(uint8_t* buf)
+static void XorWithIv(uint8_t* buf,uint8_t*Iv)
 {
   uint8_t i;
-  for (i = 0; i < BLOCKLEN; ++i) //WAS for(i = 0; i < KEYLEN; ++i) but the block in AES is always 128bit so 16 bytes!
+  for (i = 0; i < AES_BLOCKLEN; ++i) //WAS for(i = 0; i < KEYLEN; ++i) but the block in AES is always 128bit so 16 bytes!
   {
     buf[i] ^= Iv[i];
   }
 }
 
-void AES_CBC_encrypt_buffer(uint8_t* output, uint8_t* input, uint32_t length, const uint8_t* key, const uint8_t* iv)
+void AES_CBC_encrypt_buffer(struct AES_ctx *ctx,uint8_t* buf, uint32_t length)
 {
   uintptr_t i;
-  uint8_t extra = length % BLOCKLEN; /* Remaining bytes in the last non-full block */
-
-  // Skip the key expansion if key is passed as 0
-  if (0 != key)
-  {
-    Key = key;
-    KeyExpansion();
-  }
-
-  if (iv != 0)
-  {
-    Iv = (uint8_t*)iv;
-  }
-
-  for (i = 0; i < length; i += BLOCKLEN)
+  uint8_t *Iv=ctx->Iv;
+  for (i = 0; i < length; i += AES_BLOCKLEN)
   {
-    memcpy(output, input, BLOCKLEN);
-    XorWithIv(output);
-    state = (state_t*)output;
-    Cipher();
-    Iv = output;
-    input += BLOCKLEN;
-    output += BLOCKLEN;
+    XorWithIv(buf,Iv);
+    Cipher((state_t*)buf,ctx->RoundKey);
+    Iv = buf;
+    buf += AES_BLOCKLEN;
     //printf("Step %d - %d", i/16, i);
   }
-
-  if (extra)
-  {
-    memcpy(output, input, extra);
-    memset((output + extra), 0, (BLOCKLEN - extra));
-    XorWithIv(output);
-    state = (state_t*)output;
-    Cipher();
-  }
+  //store Iv in ctx for next call
+  memcpy(ctx->Iv,Iv,AES_BLOCKLEN);
 }
 
-void AES_CBC_decrypt_buffer(uint8_t* output, uint8_t* input, uint32_t length, const uint8_t* key, const uint8_t* iv)
+void AES_CBC_decrypt_buffer(struct AES_ctx *ctx, uint8_t* buf,  uint32_t length)
 {
   uintptr_t i;
-  uint8_t extra = length % BLOCKLEN; /* Remaining bytes in the last non-full block */
-
-  // Skip the key expansion if key is passed as 0
-  if (0 != key)
-  {
-    Key = key;
-    KeyExpansion();
-  }
-
-  // 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 += BLOCKLEN)
+  uint8_t storeNextIv[AES_BLOCKLEN];
+  for (i = 0; i < length; i += AES_BLOCKLEN)
   {
-    memcpy(output, input, BLOCKLEN);
-    state = (state_t*)output;
-    InvCipher();
-    XorWithIv(output);
-    Iv = input;
-    input += BLOCKLEN;
-    output += 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;
   }
 
-  if (extra)
-  {
-    memcpy(output, input, extra);
-    state = (state_t*)output;
-    InvCipher();
-  }
 }
 
 #endif // #if defined(CBC) && (CBC == 1)
@@ -581,38 +532,35 @@ void AES_CBC_decrypt_buffer(uint8_t* output, uint8_t* input, uint32_t length, co
 #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(uint8_t* output, uint8_t* input, uint32_t length, const uint8_t* key, const uint8_t* nonce)
+void AES_CTR_xcrypt_buffer(struct AES_ctx* ctx, uint8_t* buf, uint32_t length)
 {
-  uint8_t buffer[BLOCKLEN], counter[BLOCKLEN];
-  
-  memcpy(counter, nonce, BLOCKLEN);
-  Key = key;
-  KeyExpansion();
+  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 & (BLOCKLEN - 1)) == 0)
+    if (bi == AES_BLOCKLEN) //we need to regen xor compliment in buffer
     {
-      memcpy(buffer, counter, BLOCKLEN);
-      state = (state_t*)buffer;
-      Cipher();
+      
+      memcpy(buffer, ctx->Iv, AES_BLOCKLEN);
+      Cipher((state_t*)buffer,ctx->RoundKey);
 
-      /* Increment counter and handle overflow */
-      for (j = (BLOCKLEN - 1); j >= 0; --j)
+      /* Increment Iv and handle overflow */
+      for (bi = (AES_BLOCKLEN - 1); bi >= 0; --bi)
       {
-        counter[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 (counter[j] != 0)
-        {
-          break;
-        }
       }
+      bi=0;
     }
 
-    output[i] = (input[i] ^ buffer[(i & (BLOCKLEN - 1))]);
+    buf[i] = (buf[i] ^ buffer[bi]);
   }
 }
 
diff --git a/aes.h b/aes.h
index e5eb4c9..54d994a 100644
--- a/aes.h
+++ b/aes.h
@@ -28,26 +28,62 @@
 //#define AES192 1
 //#define AES256 1
 
-#if defined(ECB) && (ECB == 1)
+#define AES_BLOCKLEN 16 //Block length in bytes AES is 128b block only
+
+#if defined(AES256) && (AES256 == 1)
+    #define AES_KEYLEN 32
+    #define AES_keyExpSize 240
+#elif defined(AES192) && (AES192 == 1)
+    #define AES_KEYLEN 24
+    #define AES_keyExpSize 208
+#else
+    #define AES_KEYLEN 16   // Key length in bytes
+    #define AES_keyExpSize 176
+#endif
+
+struct AES_ctx {
+  uint8_t RoundKey[AES_keyExpSize];
+  #if (defined(CBC) && (CBC == 1)) || (defined(CTR) && (CTR == 1))
+  uint8_t Iv[AES_BLOCKLEN];
+  #endif
+};
+
+void AES_init_ctx(struct AES_ctx *ctx,const uint8_t* key);
+#if defined(CBC) && (CBC == 1)
+void AES_init_ctx_iv(struct AES_ctx *ctx,const uint8_t* key,const uint8_t* iv);
+void AES_ctx_set_iv(struct AES_ctx *ctx,const uint8_t* iv);
+#endif
 
-void AES_ECB_encrypt(const uint8_t* input, const uint8_t* key, uint8_t *output, const uint32_t length);
-void AES_ECB_decrypt(const uint8_t* input, const uint8_t* key, uint8_t *output, const uint32_t length);
+#if defined(ECB) && (ECB == 1)
+// 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* buf);
+void AES_ECB_decrypt(struct AES_ctx *ctx, const uint8_t* buf);
 
 #endif // #if defined(ECB) && (ECB == !)
 
 
 #if defined(CBC) && (CBC == 1)
-
-void AES_CBC_encrypt_buffer(uint8_t* output, uint8_t* input, uint32_t length, const uint8_t* key, const uint8_t* iv);
-void AES_CBC_decrypt_buffer(uint8_t* output, uint8_t* input, uint32_t length, const uint8_t* key, const uint8_t* iv);
+// buffer size MUST be mutile of AES_BLOCKLEN;
+// 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* 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)
 
 
 #if defined(CTR) && (CTR == 1)
 
-/* Same function for encrypting as for decrypting. Note no IV/nonce should ever be reused with the same key */
-void AES_CTR_xcrypt_buffer(uint8_t* output, uint8_t* input, uint32_t length, const uint8_t* key, const uint8_t* nonce);
+// Same function for encrypting as for decrypting. 
+// iv is incremented for every block, and usesd after  encryption as xor compliment for output
+// buffer size MUST be mutile of AES_BLOCKLEN;
+// 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* buf, uint32_t length);
 
 #endif // #if defined(CTR) && (CTR == 1)
 
diff --git a/test.c b/test.c
index 5b55eb9..1460003 100644
--- a/test.c
+++ b/test.c
@@ -97,10 +97,12 @@ static void test_encrypt_ecb_verbose(void)
 
     // print the resulting cipher as 4 x 16 byte strings
     printf("ciphertext:\n");
+    struct AES_ctx ctx;
+    AES_init_ctx(&ctx,key);
     for(i = 0; i < 4; ++i)
     {
-        AES_ECB_encrypt(plain_text + (i*16), key, buf+(i*16), 16);
-        phex(buf + (i*16));
+      AES_ECB_encrypt(&ctx,plain_text + (i*16));
+        phex(plain_text + (i*16));
     }
     printf("\n");
 }
@@ -122,13 +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];
 
-  AES_ECB_encrypt(in, key, buffer, 16);
+  struct AES_ctx ctx; AES_init_ctx(&ctx,key);
+  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");
   }
@@ -166,13 +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];
-
-  AES_CBC_decrypt_buffer(buffer, in, 64, key, iv);
+//  uint8_t buffer[64];
+  struct AES_ctx ctx;
+  AES_init_ctx_iv(&ctx,key,iv);
+  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");
   }
@@ -210,13 +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];
-
-  AES_CBC_encrypt_buffer(buffer, in, 64, key, iv);
+  struct AES_ctx ctx;
+  AES_init_ctx_iv(&ctx,key,iv);
+  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");
   }
@@ -265,13 +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];
-
-  AES_CTR_xcrypt_buffer(buffer, in, 64, key, iv);
+  
+  struct AES_ctx ctx;
+  AES_init_ctx_iv(&ctx,key,iv);
+  
+  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");
   }
@@ -298,13 +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];
 
-  AES_ECB_decrypt(in, key, buffer, 16);
+  struct AES_ctx ctx;
+  AES_init_ctx(&ctx,key);
+  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");
   }