Merge pull request #9 from pluto/pushing_forward

Pushing forward

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);
@@ -57,23 +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] {
-    // assert(nk = 4 | nk = 6 | nk = 8);
-    // fix nk to 4 and nr to 10 for AES 128
-    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];
@@ -82,26 +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);
-        //  i ← i+1 
+        key_expanded[round] = word_xor(key_expanded[round - 4], temp);
     }
 
     key_expanded
@@ -111,18 +94,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],
@@ -156,8 +139,6 @@ fn tkey_expansion(){
       [0xe1, 0x3f, 0x0c, 0xc8],
       [0xb6, 0x63, 0x0c, 0xa6],
     ];
-    println("actual \n");
-    println(key_expansion(key));
     assert(key_expansion(key) == keyExpanded);
 }
 
@@ -259,10 +240,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
 }
 

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);
+}

src/gctr.nr

@@ -0,0 +1,88 @@
+
+
+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 mut counter_block = initial_counter_block;
+    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, block);
+        // update counter block
+        counter_block = increment_32(counter_block);
+    }
+    block_to_stream(ciphertext)
+}
+
+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);
+    stream[12] = incremented[0];
+    stream[13] = incremented[1];
+    stream[14] = incremented[2];
+    stream[15] = incremented[3];
+    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]];
+    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);
+}

src/lib.nr

@@ -1,4 +1,5 @@
 mod utils;
 mod ghash;
 mod aes;
-
+mod gctr;
+mod aes_gcm;

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