Fix bug for non-multiple block length in AES-GCM

circuits/aes-gcm/gctr.circom

@@ -52,17 +52,17 @@ template GCTR(INPUT_LEN) {
     plainTextBlocks.stream <== plainText;
 
     // Step 1: Generate counter blocks
-    signal CounterBlocks[nBlocks][4][4];
+    signal CounterBlocks[nBlocks + 1][4][4];
     CounterBlocks[0] <== initialCounterBlock;
 
     // First counter block is passed in, as a combination of the IV right padded with zeros IV is 96 bits or 12 bytes
     // The next counter needs to be set by incrementing the right most 32 bits (4 bytes) of the previous counter block
     //
     // component to increment the last word of the counter block
-    component inc32[nBlocks - 1];
+    component inc32[nBlocks];
     // For i = 2 to nBlocks, let CBi = inc32(CBi-1).
 
-    for (var i = 1; i < nBlocks; i++) {
+    for (var i = 1; i < nBlocks + 1; i++) {
         inc32[i - 1] = IncrementWord();
         inc32[i - 1].in <== CounterBlocks[i - 1][3];
 
@@ -103,7 +103,7 @@ template GCTR(INPUT_LEN) {
     // encrypt the last counter block 
     aes[nBlocks] = Cipher();
     aes[nBlocks].key <== key;
-    aes[nBlocks].block <== CounterBlocks[nBlocks-1];
+    aes[nBlocks].block <== CounterBlocks[nBlocks];
     component aesCipherToStream = ToStream(1, 16);
     aesCipherToStream.blocks[0] <== aes[nBlocks].cipher;
 

circuits/aes-gcm/utils.circom

@@ -112,7 +112,6 @@ template BitwiseRightShift(n, r) {
     }
 }
 
-
 template BitwiseXor(n) {
     signal input a[n];
     signal input b[n];
@@ -135,26 +134,7 @@ template BitwiseOr(n) {
     }
 }
 
-// compute the OR of n inputs, each m bits wide
-template OrMultiple(n, m) {
-    signal input inputs[n][m];
-    signal output out[m];
-
-    signal mids[n][m];
-    mids[0] <== inputs[0];
-
-    component ors[n-1];
-    for(var i=0; i<n-1; i++) {
-        ors[i] = BitwiseOr(m);
-        ors[i].a <== mids[i];
-        ors[i].b <== inputs[i+1];
-        mids[i+1] <== ors[i].out;
-    }
-
-    out <== mids[n-1];
-}
-
-// compute the XOR of n inputs, each m bits wide
+// compute the XOR of n inputs, each of m bytes
 template XorMultiple(n, m) {
     signal input inputs[n][m];
     signal output out[m];
@@ -164,10 +144,10 @@ template XorMultiple(n, m) {
 
     component xors[n-1];
     for(var i=0; i<n-1; i++) {
-        xors[i] = BitwiseXor(m);
-        xors[i].a <== mids[i];
-        xors[i].b <== inputs[i+1];
-        mids[i+1] <== xors[i].out;
+         xors[i] = XORBLOCK(m);
+         xors[i].a <== mids[i];
+         xors[i].b <== inputs[i+1];
+         mids[i+1] <== xors[i].out;
     }
 
     out <== mids[n-1];

circuits/test/aes-gcm/aes-gcm.test.ts

@@ -67,6 +67,27 @@ describe("aes-gcm", () => {
     assert.deepEqual(witness.cipherText, hexBytesToBigInt(ct))
     // assert.deepEqual(witness.authTag, hexBytesToBigInt(auth_tag));
   });
+
+  it("should work in case of non-integer multiples of block size", async () => {
+    let circuit_one_block: WitnessTester<["key", "iv", "plainText", "aad"], ["cipherText", "tag"]>;
+    circuit_one_block = await circomkit.WitnessTester(`aes-gcm`, {
+      file: "aes-gcm/aes-gcm",
+      template: "AESGCM",
+      params: [30],
+    });
+
+    const key = hexToBytes('31313131313131313131313131313131');
+    const iv = hexToBytes('313131313131313131313131');
+    const msg = hexToBytes('7465737468656c6c6f776f726c64307465737468656c6c6f776f726c6431');
+    const aad = hexToBytes('00000000000000000000000000000000')
+    const ct = hexToBytes('2929d2bb1ae94804406cd135325933123763622c7c374c0542ae9198a16f');
+    const auth_tag = hexToBytes('31837b1974685501be791364a586fa6e');
+
+    const witness = await circuit_one_block.compute({ key: key, iv: iv, plainText: msg, aad: aad }, ["cipherText", "authTag"])
+
+    assert.deepEqual(witness.cipherText, hexBytesToBigInt(ct))
+    // assert.deepEqual(witness.authTag, hexBytesToBigInt(auth_tag));
+  });
 });
 
 // signal input key[16]; // 128-bit key

src/main.rs

@@ -98,6 +98,37 @@ mod tests {
         println!("ct={}", hex::encode(ct));
     }
 
+    /// Test AES-GCM with data length not a multiple of block size
+    #[tokio::test]
+    async fn test_aes_gcm_non_multiple_length() {
+        use aes_gcm::{
+            aead::{generic_array::GenericArray, Aead, NewAead, Payload},
+            Aes128Gcm,
+        };
+
+        let test_key = [
+            0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31,
+            0x31, 0x31,
+        ];
+        let test_iv = [0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31];
+
+        let mut payload: Vec<u8> = Vec::new();
+        for last_byte in 0..2 {
+            let message = format!("testhelloworld{}", last_byte);
+            payload.extend(message.as_bytes());
+        }
+        let aes_payload = Payload { msg: &payload, aad: &[] };
+
+        let cipher = Aes128Gcm::new_from_slice(&test_key).unwrap();
+        let nonce = GenericArray::from_slice(&test_iv);
+        let ct = cipher.encrypt(nonce, aes_payload).expect("error generating ct");
+
+        println!("key={}", hex::encode(test_key));
+        println!("iv={}", hex::encode(test_iv));
+        println!("msg={}", hex::encode(payload));
+        println!("ct={}", hex::encode(ct));
+    }
+
     #[tokio::test]
     async fn test_ghash() {
         use ghash::{