From dae977fc179b761d343957291e604ba244e39a99 Mon Sep 17 00:00:00 2001
From: Waylon Jepsen <waylonjepsen1@gmail.com>
Date: Thu, 10 Oct 2024 11:28:20 -0600
Subject: [PATCH 1/6] nior bug

---
 src/utils.nr | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/src/utils.nr b/src/utils.nr
index 790976a..198d89f 100644
--- a/src/utils.nr
+++ b/src/utils.nr
@@ -1,6 +1,25 @@
 type Stream = [u8; 16];
 type Block = [[u8; 4]; 4];
 
+pub(crate) fn byte_stream_blocks<N>(input: [u8; N]) -> [Block; N / 16] {
+    let nBlocks = input.len() / 16;
+    let mut blocks = [Block::default(); nBlocks];
+    for i in 0..nBlocks {
+        for j in 0..16 {
+            blocks[i][j] = input[i * 16 + j];
+        }
+    }
+    blocks
+}
+
+#[test]
+fn test_to_blocks() {
+    let input = [0u8; 16];
+    let blocks = byte_stream_to_blocks(input);
+    assert(blocks.len() == 1);
+    assert(blocks[0] == input);
+}
+
 pub(crate) fn stream_xor(stream1: Stream, stream2: Stream) -> Stream {
     let mut result: Stream = [0; 16];
     for i in 0..16 {

From 295a33eaec78a9eccffd7ae77b1a4b81118ddb98 Mon Sep 17 00:00:00 2001
From: Waylon Jepsen <waylonjepsen1@gmail.com>
Date: Thu, 10 Oct 2024 12:05:08 -0600
Subject: [PATCH 2/6] carry adder for 4 byte word

---
 src/gctr.nr  | 69 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 src/lib.nr   |  2 +-
 src/utils.nr | 19 ---------------
 3 files changed, 70 insertions(+), 20 deletions(-)
 create mode 100644 src/gctr.nr

diff --git a/src/gctr.nr b/src/gctr.nr
new file mode 100644
index 0000000..1f9f17b
--- /dev/null
+++ b/src/gctr.nr
@@ -0,0 +1,69 @@
+
+
+use crate::utils::{stream_xor, stream_to_block, Stream, Block};
+
+// starting with fixed size input
+fn gctr(key: Stream, initial_counter_block: Block, plaintext: Stream) -> Stream {
+    let mut ciphertext: [u8; 16] = [0; 16];
+    let mut counter_block = initial_counter_block;
+
+    let pt_blocks = [stream_to_block(plaintext)];
+    for chunk in pt_blocks {
+        let encrypted_counter_block = aes_encrypt(key, counter_block);
+        let mut block = [0u8; 16];
+        for (i, &byte) in chunk.iter().enumerate() {
+            block[i] = byte;
+        }
+        let encrypted_block = stream_xor(encrypted_counter_block, block);
+        ciphertext.extend_from_slice(&encrypted_block[..chunk.len()]);
+        counter_block = increment_counter_block(counter_block);
+    }
+
+    ciphertext
+}
+
+/// carry adder on 4 byte words
+fn increment_word(word: [u8; 4]) -> [u8; 4] {
+    let mut incremented = [word[3], word[2], word[1], word[0]];
+    let mut carry = 1;
+    for i in 0..4 {
+        if incremented[i] == 0xFF {
+            incremented[i] = 0x00;
+            carry = 1;
+        } else {
+            incremented[i] += carry;
+            carry = 0;
+        }
+    }
+    [incremented[3], incremented[2], incremented[1], incremented[0]]
+}
+#[test]
+fn test_increment_word() {
+    let word = [0x00, 0x00, 0x00, 0x00];
+    let expected_incremented_word = [0x00, 0x00, 0x00, 0x01];
+    assert(increment_word(word) == expected_incremented_word);
+
+    let word = [0x00, 0x00, 0x00, 0xFF];
+    let expected_incremented_word = [0x00, 0x00, 0x01, 0x00];
+    assert(increment_word(word) == expected_incremented_word);
+
+    let word = [0x00, 0x00, 0xFF, 0xFF];
+    let expected_incremented_word = [0x00, 0x01, 0x00, 0x00];
+    assert(increment_word(word) == expected_incremented_word);
+
+    let word = [0xFF, 0xFF, 0xFF, 0xFF];
+    let expected_incremented_word = [0x00, 0x00, 0x00, 0x00];
+    assert(increment_word(word) == expected_incremented_word);
+}
+
+
+#[test]
+fn test_gctr() {
+    let key = [0xca, 0xaa, 0x3f, 0x6f, 0xd3, 0x18, 0x22, 0xed, 0x2d, 0x21, 0x25, 0xf2, 0x25, 0xb0, 0x16, 0x9f];
+    let initial_counter_block: Block = [[0x7f,0x48,0x12,0x00],[0x6d,0x3e,0xfa,0x00],[0x90,0x8c,0x55,0x00],[0x41,0x14,0x2a,0x02]];
+    let plaintext = [0x84, 0xc9, 0x07, 0xb1, 0x1a, 0xe3, 0xb7, 0x9f,0xc4, 0x45, 0x1d, 0x1b, 0xf1, 0x7f, 0x4a, 0x99];
+    let expected_ciphertext = [0xfd, 0xb4, 0xaa, 0xfa, 0x35, 0x19, 0xd3, 0xc0,0x55, 0xbe, 0x8b, 0x34, 0x77, 0x64, 0xea, 0x33];
+
+    let ciphertext = gctr(key, initial_counter_block, plaintext);
+    assert(ciphertext == expected_ciphertext);
+}
diff --git a/src/lib.nr b/src/lib.nr
index 0231f2f..e247680 100644
--- a/src/lib.nr
+++ b/src/lib.nr
@@ -1,4 +1,4 @@
 mod utils;
 mod ghash;
 mod aes;
-
+mod gctr;
\ No newline at end of file
diff --git a/src/utils.nr b/src/utils.nr
index 198d89f..790976a 100644
--- a/src/utils.nr
+++ b/src/utils.nr
@@ -1,25 +1,6 @@
 type Stream = [u8; 16];
 type Block = [[u8; 4]; 4];
 
-pub(crate) fn byte_stream_blocks<N>(input: [u8; N]) -> [Block; N / 16] {
-    let nBlocks = input.len() / 16;
-    let mut blocks = [Block::default(); nBlocks];
-    for i in 0..nBlocks {
-        for j in 0..16 {
-            blocks[i][j] = input[i * 16 + j];
-        }
-    }
-    blocks
-}
-
-#[test]
-fn test_to_blocks() {
-    let input = [0u8; 16];
-    let blocks = byte_stream_to_blocks(input);
-    assert(blocks.len() == 1);
-    assert(blocks[0] == input);
-}
-
 pub(crate) fn stream_xor(stream1: Stream, stream2: Stream) -> Stream {
     let mut result: Stream = [0; 16];
     for i in 0..16 {

From 0f07c8a2bd8682ad82a87c88edefde1f81a65be4 Mon Sep 17 00:00:00 2001
From: Waylon Jepsen <waylonjepsen1@gmail.com>
Date: Thu, 10 Oct 2024 13:33:30 -0600
Subject: [PATCH 3/6] bug: gctr

---
 src/aes.nr  | 25 ++++++++++++-------------
 src/gctr.nr | 42 ++++++++++++++++++++++++++++--------------
 2 files changed, 40 insertions(+), 27 deletions(-)

diff --git a/src/aes.nr b/src/aes.nr
index e817e9c..4844904 100644
--- a/src/aes.nr
+++ b/src/aes.nr
@@ -101,7 +101,6 @@ pub(crate) fn key_expansion(key: [u8; 16]) -> [[u8; 4]; 44] {
 
         //  w[i] ← w[i−Nk]⊕temp
         key_expanded[round] = word_xor(key_expanded[round - nk], temp);
-        //  i ← i+1 
     }
 
     key_expanded
@@ -111,18 +110,18 @@ pub(crate) fn key_expansion(key: [u8; 16]) -> [[u8; 4]; 44] {
 fn tkey_expansion(){
     let key = [0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f, 0x3c];
     let keyExpanded = [
-      [0x2b, 0x7e, 0x15, 0x16], // got this
-      [0x28, 0xae, 0xd2, 0xa6], // got this
-      [0xab, 0xf7, 0x15, 0x88], // got this
-      [0x09, 0xcf, 0x4f, 0x3c], // got this
-      [0xa0, 0xfa, 0xfe, 0x17], // got this
-      [0x88, 0x54, 0x2c, 0xb1], // got this
-      [0x23, 0xa3, 0x39, 0x39], // got this
-      [0x2a, 0x6c, 0x76, 0x05], // got this
-      [0xf2, 0xc2, 0x95, 0xf2], // got this
-      [0x7a, 0x96, 0xb9, 0x43], // got this
-      [0x59, 0x35, 0x80, 0x7a], // got this
-      [0x73, 0x59, 0xf6, 0x7f], // got this
+      [0x2b, 0x7e, 0x15, 0x16],
+      [0x28, 0xae, 0xd2, 0xa6],
+      [0xab, 0xf7, 0x15, 0x88],
+      [0x09, 0xcf, 0x4f, 0x3c],
+      [0xa0, 0xfa, 0xfe, 0x17],
+      [0x88, 0x54, 0x2c, 0xb1],
+      [0x23, 0xa3, 0x39, 0x39],
+      [0x2a, 0x6c, 0x76, 0x05],
+      [0xf2, 0xc2, 0x95, 0xf2],
+      [0x7a, 0x96, 0xb9, 0x43],
+      [0x59, 0x35, 0x80, 0x7a],
+      [0x73, 0x59, 0xf6, 0x7f],
       [0x3d, 0x80, 0x47, 0x7d],
       [0x47, 0x16, 0xfe, 0x3e],
       [0x1e, 0x23, 0x7e, 0x44],
diff --git a/src/gctr.nr b/src/gctr.nr
index 1f9f17b..3b80ce3 100644
--- a/src/gctr.nr
+++ b/src/gctr.nr
@@ -1,25 +1,39 @@
 
 
-use crate::utils::{stream_xor, stream_to_block, Stream, Block};
+use crate::utils::{stream_xor, stream_to_block, block_to_stream, Stream, Block};
+use crate::aes::{aes, add_round_key};
 
 // starting with fixed size input
 fn gctr(key: Stream, initial_counter_block: Block, plaintext: Stream) -> Stream {
-    let mut ciphertext: [u8; 16] = [0; 16];
-    let mut counter_block = initial_counter_block;
+    let mut ciphertext: Block = [[0; 4]; 4];
+    let mut counter_block: Block = initial_counter_block;
 
-    let pt_blocks = [stream_to_block(plaintext)];
-    for chunk in pt_blocks {
-        let encrypted_counter_block = aes_encrypt(key, counter_block);
-        let mut block = [0u8; 16];
-        for (i, &byte) in chunk.iter().enumerate() {
-            block[i] = byte;
-        }
-        let encrypted_block = stream_xor(encrypted_counter_block, block);
-        ciphertext.extend_from_slice(&encrypted_block[..chunk.len()]);
+
+    let pt_blocks = stream_to_block(plaintext);
+    println("pt_blocks");
+    println(pt_blocks);
+    // for block in pt_blocks {
+        // encrypt counter block
+        let encrypted_counter_block = aes(counter_block, key);
+        // xor with plaintext
+        ciphertext = add_round_key(encrypted_counter_block, pt_blocks);
+        // update counter block
         counter_block = increment_counter_block(counter_block);
-    }
+    // }
+    println("ciphertext");
+    println(ciphertext);
+    block_to_stream(ciphertext)
+}
 
-    ciphertext
+fn increment_counter_block(counter_block: Block) -> Block {
+    let mut stream = block_to_stream(counter_block);
+    let mut word = [stream[12], stream[13], stream[14], stream[15]];
+    let incremented = increment_word(word);
+    stream[12] = incremented[0];
+    stream[13] = incremented[1];
+    stream[14] = incremented[2];
+    stream[15] = incremented[3];
+    stream_to_block(stream)
 }
 
 /// carry adder on 4 byte words

From 241734e875bca37b9469fd07044ef432b5ccd745 Mon Sep 17 00:00:00 2001
From: Waylon Jepsen <waylonjepsen1@gmail.com>
Date: Thu, 10 Oct 2024 13:51:49 -0600
Subject: [PATCH 4/6] passing test for gctr on one pt block

---
 src/aes.nr  |  2 --
 src/gctr.nr | 17 +++++++++++++----
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/src/aes.nr b/src/aes.nr
index 4844904..1610f30 100644
--- a/src/aes.nr
+++ b/src/aes.nr
@@ -155,8 +155,6 @@ fn tkey_expansion(){
       [0xe1, 0x3f, 0x0c, 0xc8],
       [0xb6, 0x63, 0x0c, 0xa6],
     ];
-    println("actual \n");
-    println(key_expansion(key));
     assert(key_expansion(key) == keyExpanded);
 }
 
diff --git a/src/gctr.nr b/src/gctr.nr
index 3b80ce3..e850f82 100644
--- a/src/gctr.nr
+++ b/src/gctr.nr
@@ -6,17 +6,15 @@ use crate::aes::{aes, add_round_key};
 // starting with fixed size input
 fn gctr(key: Stream, initial_counter_block: Block, plaintext: Stream) -> Stream {
     let mut ciphertext: Block = [[0; 4]; 4];
-    let mut counter_block: Block = initial_counter_block;
-
-
     let pt_blocks = stream_to_block(plaintext);
+    let mut counter_block = initial_counter_block;
     println("pt_blocks");
     println(pt_blocks);
     // for block in pt_blocks {
         // encrypt counter block
         let encrypted_counter_block = aes(counter_block, key);
         // xor with plaintext
-        ciphertext = add_round_key(encrypted_counter_block, pt_blocks);
+        ciphertext = add_cipher(encrypted_counter_block, pt_blocks);
         // update counter block
         counter_block = increment_counter_block(counter_block);
     // }
@@ -36,6 +34,17 @@ fn increment_counter_block(counter_block: Block) -> Block {
     stream_to_block(stream)
 }
 
+
+fn add_cipher(state: Block, key: Block) -> Block {
+    let mut new_state: Block = [[0; 4]; 4];
+    for i in 0..4 {
+        for j in 0..4 {
+            new_state[i][j] = state[i][j] ^ key[i][j];
+        }
+    }
+    new_state
+}
+
 /// carry adder on 4 byte words
 fn increment_word(word: [u8; 4]) -> [u8; 4] {
     let mut incremented = [word[3], word[2], word[1], word[0]];

From cea8e53b8f6cefde9afb3a8747a7be3c68167dcd Mon Sep 17 00:00:00 2001
From: Waylon Jepsen <waylonjepsen1@gmail.com>
Date: Thu, 10 Oct 2024 15:08:42 -0600
Subject: [PATCH 5/6] working gcrt

---
 src/aes.nr   | 11 ++++-------
 src/gctr.nr  | 14 +++++---------
 src/utils.nr |  2 +-
 3 files changed, 10 insertions(+), 17 deletions(-)

diff --git a/src/aes.nr b/src/aes.nr
index 1610f30..7b5e62c 100644
--- a/src/aes.nr
+++ b/src/aes.nr
@@ -20,7 +20,6 @@ use crate::utils::{byte_to_bits, bits_to_byte, Block, word_xor, Stream};
 //                                                                      ▼
 //                                                                 Ciphertext
 pub(crate) fn aes(block: Block, key: Stream) -> Block {
-    // TODO(WJ 2024-10-09): implement Key Expansion
     let key_expanded = key_expansion(key);
     let first_round_key: Block = [key_expanded[0], key_expanded[1], key_expanded[2], key_expanded[3]];
     let mut state = add_round_key(block, first_round_key);
@@ -62,8 +61,6 @@ fn test_aes() {
 // @inputs key: array of nk*4 bytes representing the key
 // @outputs keyExpanded: array of (nr+1)*4 words i.e for AES 128, 192, 256 it will be  44, 52, 60 words
 pub(crate) fn key_expansion(key: [u8; 16]) -> [[u8; 4]; 44] {
-    // assert(nk = 4 | nk = 6 | nk = 8);
-    // fix nk to 4 and nr to 10 for AES 128
     let nk = 4;
     let nr = 10;
 
@@ -256,10 +253,10 @@ fn mix_columns(state: Block) -> Block {
 
 fn mix_column(column: [u8; 4]) -> [u8; 4] {
     let mut new_column = [0; 4];
-    new_column[0] = s0(column); // s0
-    new_column[1] = s1(column); // s1
-    new_column[2] = s2(column); // s2
-    new_column[3] = s3(column); // s3
+    new_column[0] = s0(column);
+    new_column[1] = s1(column);
+    new_column[2] = s2(column);
+    new_column[3] = s3(column);
     new_column
 }
 
diff --git a/src/gctr.nr b/src/gctr.nr
index e850f82..4adeda7 100644
--- a/src/gctr.nr
+++ b/src/gctr.nr
@@ -4,22 +4,19 @@ use crate::utils::{stream_xor, stream_to_block, block_to_stream, Stream, Block};
 use crate::aes::{aes, add_round_key};
 
 // starting with fixed size input
+// TODO(WJ 2024-10-10): make this work for arbitrary length input
 fn gctr(key: Stream, initial_counter_block: Block, plaintext: Stream) -> Stream {
     let mut ciphertext: Block = [[0; 4]; 4];
-    let pt_blocks = stream_to_block(plaintext);
+    let pt_blocks = [stream_to_block(plaintext)];
     let mut counter_block = initial_counter_block;
-    println("pt_blocks");
-    println(pt_blocks);
-    // for block in pt_blocks {
+    for block in pt_blocks {
         // encrypt counter block
         let encrypted_counter_block = aes(counter_block, key);
         // xor with plaintext
-        ciphertext = add_cipher(encrypted_counter_block, pt_blocks);
+        ciphertext = add_cipher(encrypted_counter_block, block);
         // update counter block
         counter_block = increment_counter_block(counter_block);
-    // }
-    println("ciphertext");
-    println(ciphertext);
+    }
     block_to_stream(ciphertext)
 }
 
@@ -86,7 +83,6 @@ fn test_gctr() {
     let initial_counter_block: Block = [[0x7f,0x48,0x12,0x00],[0x6d,0x3e,0xfa,0x00],[0x90,0x8c,0x55,0x00],[0x41,0x14,0x2a,0x02]];
     let plaintext = [0x84, 0xc9, 0x07, 0xb1, 0x1a, 0xe3, 0xb7, 0x9f,0xc4, 0x45, 0x1d, 0x1b, 0xf1, 0x7f, 0x4a, 0x99];
     let expected_ciphertext = [0xfd, 0xb4, 0xaa, 0xfa, 0x35, 0x19, 0xd3, 0xc0,0x55, 0xbe, 0x8b, 0x34, 0x77, 0x64, 0xea, 0x33];
-
     let ciphertext = gctr(key, initial_counter_block, plaintext);
     assert(ciphertext == expected_ciphertext);
 }
diff --git a/src/utils.nr b/src/utils.nr
index 790976a..f2fa2d1 100644
--- a/src/utils.nr
+++ b/src/utils.nr
@@ -21,7 +21,7 @@ pub(crate) fn stream_to_block(stream: Stream) -> Block {
     let mut block: Block = [[0; 4]; 4];
     for i in 0..4 {
         for j in 0..4 {
-            block[j][i] = stream[i * 4 + j]; // Change the indexing to column-major order
+            block[j][i] = stream[i * 4 + j];
         }
     }
     block

From 3d6fba29266f761c43dc727e3ec09bbd6a50e534 Mon Sep 17 00:00:00 2001
From: Waylon Jepsen <waylonjepsen1@gmail.com>
Date: Thu, 10 Oct 2024 19:42:08 -0600
Subject: [PATCH 6/6] aes-gcm tests

---
 src/aes.nr     | 25 ++++++-----------------
 src/aes_gcm.nr | 55 ++++++++++++++++++++++++++++++++++++++++++++++++++
 src/gctr.nr    |  4 ++--
 src/lib.nr     |  3 ++-
 4 files changed, 65 insertions(+), 22 deletions(-)
 create mode 100644 src/aes_gcm.nr

diff --git a/src/aes.nr b/src/aes.nr
index 7b5e62c..ecf8dfc 100644
--- a/src/aes.nr
+++ b/src/aes.nr
@@ -56,21 +56,14 @@ fn test_aes() {
     assert(aes(block, key) == cipher);
 }
 
-// @param nk: number of keys can be 4, 6, or 8
-// @param nr: number of rounds which can be 10, 12, 14 for AES 128, 192, 256
 // @inputs key: array of nk*4 bytes representing the key
-// @outputs keyExpanded: array of (nr+1)*4 words i.e for AES 128, 192, 256 it will be  44, 52, 60 words
+// @outputs keyExpanded: array of (nr+1)*4 words i.e for AES 128 will be 44 words
 pub(crate) fn key_expansion(key: [u8; 16]) -> [[u8; 4]; 44] {
-    let nk = 4;
-    let nr = 10;
-
-    let total_words = 44;
-    let effective_rounds = 10;
 
     let mut key_expanded: [[u8; 4]; 44] = [[0; 4]; 44];
 
     // first nk words are the key divided into 4 byte words
-    for i in 0..nk {
+    for i in 0..4 {
         let mut temp = [0; 4];
         for j in 0..4 {
             temp[j] = key[i * 4 + j];
@@ -79,25 +72,19 @@ pub(crate) fn key_expansion(key: [u8; 16]) -> [[u8; 4]; 44] {
     }
 
     // while i ≤ 4 ∗Nr +3 do
-    for round in 4..(4 * nr + 3) + 1 {
+    for round in 4..44 {
         // temp ← w[i−1]
         let mut temp = key_expanded[round - 1];
 
         // if i mod Nk = 0 then
-        if round % nk == 0 {
+        if round % 4 == 0 {
             // temp ← SUBWORD(ROTWORD(temp))⊕Rcon[i/Nk]
             temp = sub_word(rot_word(temp));
-            temp = word_xor(temp, round_constant((round / nk) - 1));
+            temp = word_xor(temp, round_constant((round / 4) - 1));
 
         }
-        // else if Nk > 6 and i mod Nk = 4 then
-        // } else if nk > 6 & round % nk == 4 {
-        //     // temp ← SUBWORD(temp)
-        //     temp = sub_word(temp);
-        // }
-
         //  w[i] ← w[i−Nk]⊕temp
-        key_expanded[round] = word_xor(key_expanded[round - nk], temp);
+        key_expanded[round] = word_xor(key_expanded[round - 4], temp);
     }
 
     key_expanded
diff --git a/src/aes_gcm.nr b/src/aes_gcm.nr
new file mode 100644
index 0000000..183b281
--- /dev/null
+++ b/src/aes_gcm.nr
@@ -0,0 +1,55 @@
+use crate::utils::{Stream, Block, stream_to_block, block_to_stream};
+use crate::gctr::{gctr, increment_32};
+use crate::ghash::ghash;
+use crate::aes::aes;
+
+pub fn aes_gcm(key: Stream, plaintext: Stream, iv: [u8; 12], aad: Stream) -> (Stream, Stream) {
+
+    // step 1: generate hashkey as encryption of a zero block with AES
+    let zero_block: Block = [[0; 4]; 4];
+    let hashkey = block_to_stream(aes(zero_block, key));
+
+    // step 2: generate jO as iv || 0 ^{31} || 1 where || is concatenation
+    let mut j0 = [0; 16];
+    for i in 0..12 {
+        j0[i] = iv[i];
+    }
+    j0[15] = 1;
+    let j0_block = stream_to_block(j0);
+
+    // step 3: Let C=GCTRK(inc32(J0), P).
+    let C = gctr(key, increment_32(j0_block), plaintext);
+    // Define a block, S, as follows:
+    // A is 16 bytes, 128 bits, C is 16 bytes, 128 bits, then two 64 bit numbers is 16 bytes so 3 16 byte blocks
+    let mut S = [[0; 16]; 3];
+    S[0] = aad;
+    S[1] = C;
+    S[2] = [0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x80,];
+    let res = ghash(hashkey, S);
+    // S = GHASHH (A || 0v || C || 0u || [len(A)]64 || [len(C)]64). 
+    let T = gctr(key, j0_block, res);
+    // // fn gctr(key: Stream, initial_counter_block: Block, plaintext: Stream) -> Stream {
+    (C, T)
+}
+
+#[test]
+fn test_aes_gcm() {
+    let key = [0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00];
+    let plainText = [0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00];
+    let iv = [0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00];
+    let aad = [0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00];
+    let expected_output = [0x03, 0x88, 0xda, 0xce, 0x60, 0xb6, 0xa3, 0x92, 0xf3, 0x28, 0xc2, 0xb9, 0x71, 0xb2, 0xfe, 0x78];
+    let (ciphertext, _tag) = aes_gcm(key, plainText, iv, aad);
+    assert(ciphertext == expected_output);
+}
+
+#[test]
+fn test_aes_gcm_2() {
+    let key = [0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31];
+    let iv = [0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31];
+    let msg = [0x74, 0x65, 0x73, 0x74, 0x68, 0x65, 0x6c, 0x6c, 0x6f, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30];
+    let aad = [0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00];
+    let ct = [0x29, 0x29, 0xd2, 0xbb, 0x1a, 0xe9, 0x48, 0x04, 0x40, 0x2b, 0x8e, 0x77, 0x6e, 0x0d, 0x33, 0x56];
+    let (ciphertext, _tag) = aes_gcm(key, msg, iv, aad);
+    assert(ciphertext == ct);
+}
\ No newline at end of file
diff --git a/src/gctr.nr b/src/gctr.nr
index 4adeda7..5a9a1b1 100644
--- a/src/gctr.nr
+++ b/src/gctr.nr
@@ -15,12 +15,12 @@ fn gctr(key: Stream, initial_counter_block: Block, plaintext: Stream) -> Stream
         // xor with plaintext
         ciphertext = add_cipher(encrypted_counter_block, block);
         // update counter block
-        counter_block = increment_counter_block(counter_block);
+        counter_block = increment_32(counter_block);
     }
     block_to_stream(ciphertext)
 }
 
-fn increment_counter_block(counter_block: Block) -> Block {
+fn increment_32(counter_block: Block) -> Block {
     let mut stream = block_to_stream(counter_block);
     let mut word = [stream[12], stream[13], stream[14], stream[15]];
     let incremented = increment_word(word);
diff --git a/src/lib.nr b/src/lib.nr
index e247680..ccac088 100644
--- a/src/lib.nr
+++ b/src/lib.nr
@@ -1,4 +1,5 @@
 mod utils;
 mod ghash;
 mod aes;
-mod gctr;
\ No newline at end of file
+mod gctr;
+mod aes_gcm;