Generate proof functions

src/implementations/AES128CTR.ts

@@ -1,66 +1,120 @@
-import { Struct, SelfProof, ZkProgram, Field } from "o1js";
+import { Struct, SelfProof, ZkProgram, Field, Poseidon } from "o1js";
 import { Byte16 } from "../primitives/Bytes.js";
 import { computeIterativeAes128Encryption } from "./IterativeAES128.js";
 
+/**
+ * Public input for the AES-128 CTR mode verification circuit.
+ *
+ * @property cipher - The ciphertext produced by encrypting a single block.
+ * @property iv - The initialization vector. This value must be randomly generated for each proof.
+ */
 class AES128CTRPublicInput extends Struct({
   cipher: Byte16,
-  // initialization vector: critical that this is provided randomly each time
-  // This can be publicly disclosed and gives the verifier the befit of checking they have the correct IV
-  // It can be made private in which case the verifier is not guaranteed they have the correct IV
-  iv: Field,
+  iv: Field, // Randomly generated IV that can be publicly disclosed.
 }) {}
 
-// Cipher under the CTR mode for a single block
+/**
+ * Public output for the AES-128 CTR mode verification circuit.
+ *
+ * @property counter - The current counter value (starting from 1).
+ * @property keyHash - The Poseidon hash of the key used in the encryption.
+ */
+class AES128CTRPublicOutput extends Struct({
+  counter: Field,
+  keyHash: Field,
+}) {}
+
+/**
+ * Computes the cipher for a single block in CTR mode.
+ *
+ * @param iv_plus_ctr - The sum of the IV and counter.
+ * @param key - The 128-bit key.
+ * @param message - The plaintext message block.
+ * @returns The ciphertext produced by XORing the plaintext with the key stream.
+ *
+ * The key stream is computed by applying an iterative AES-128 encryption on the IV+counter.
+ */
 export function computeCipher(
   iv_plus_ctr: Field,
   key: Byte16,
-  message: Byte16, // plaintext
+  message: Byte16,
 ): Byte16 {
-  // Use AES128 just to get the key
   const curr_key: Byte16 = computeIterativeAes128Encryption(
     Byte16.fromField(iv_plus_ctr),
     key,
   );
-  // compute curr_key by encyrpting counter + iv with key with AES128
-  // simply xor with the key to get ciphertext
   return Byte16.xor(message, curr_key);
 }
 
+/**
+ * ZK program for verifying AES-128 CTR mode encryption using recursive proofs.
+ *
+ * It supports both a base case (for a single block) and an inductive step (for multiple blocks).
+ * In addition to verifying the encryption, each proof computes and returns the Poseidon hash of the key.
+ * In the inductive case, the key hash is compared with the previous proof's hash to enforce consistency.
+ */
 const Aes128Ctr = ZkProgram({
   name: "aes-verify-iterative",
   publicInput: AES128CTRPublicInput,
-  publicOutput: Field, // The counter
+  publicOutput: AES128CTRPublicOutput,
 
   methods: {
-    // base case for a singleton block
+    // Base case: Verify a single block encryption.
     base: {
+      // Private inputs: plaintext message and key.
       privateInputs: [Byte16, Byte16],
 
       async method(input: AES128CTRPublicInput, message: Byte16, key: Byte16) {
-        // ctr = 0, so iv passed as is
         const cipher = computeCipher(input.iv, key, message);
         cipher.assertEquals(input.cipher);
-        return { publicOutput: Field(1) };
+        const keyHash = Poseidon.hash([key.toField()]);
+
+        return {
+          publicOutput: new AES128CTRPublicOutput({
+            counter: Field(1),
+            keyHash,
+          }),
+        };
       },
     },
+
+    // Inductive step: Verify subsequent block encryptions.
     inductive: {
-      // the output type of the SelfProof is the Ctr
-      privateInputs: [SelfProof<AES128CTRPublicInput, Field>, Byte16, Byte16],
+      // Private inputs:
+      // - A recursive proof of the previous block's encryption.
+      // - The plaintext message for the current block.
+      // - The key for the current block.
+      privateInputs: [
+        SelfProof<AES128CTRPublicInput, AES128CTRPublicOutput>,
+        Byte16,
+        Byte16,
+      ],
+
       async method(
         input: AES128CTRPublicInput,
-        previousProof: SelfProof<AES128CTRPublicInput, Field>,
+        previousProof: SelfProof<AES128CTRPublicInput, AES128CTRPublicOutput>,
         message: Byte16,
-        key: Byte16, // TODO: How do I assert the prover always provides the same key to this function
+        key: Byte16,
       ) {
+        const currentKeyHash = Poseidon.hash([key.toField()]);
+        currentKeyHash.assertEquals(previousProof.publicOutput.keyHash);
         previousProof.verify();
         input.iv.assertEquals(previousProof.publicInput.iv);
+
         const cipher = computeCipher(
-          input.iv.add(previousProof.publicOutput),
+          input.iv.add(previousProof.publicOutput.counter),
           key,
           message,
         );
         cipher.assertEquals(input.cipher);
-        return { publicOutput: previousProof.publicOutput.add(Field(1)) };
+
+        const newCounter = previousProof.publicOutput.counter.add(Field(1));
+        return {
+          publicOutput: new AES128CTRPublicOutput({
+            counter: newCounter,
+            keyHash: currentKeyHash,
+          }),
+        };
       },
     },
   },

src/implementations/IterativeAES128.ts

@@ -1,16 +1,24 @@
-import { Struct, ZkProgram } from "o1js";
+import { Proof, Struct, ZkProgram } from "o1js";
 import { Byte16 } from "../primitives/Bytes.js";
 import { shiftRows } from "../lib/ShiftRows.js";
 import { sbox } from "../lib/SBox.js";
 import { mixColumn } from "../lib/MixColumns.js";
 import { addRoundKey } from "../lib/AddRoundKey.js";
 import { NUM_ROUNDS_128 as NUM_ROUNDS } from "../utils/constants.js";
 import { expandKey128 } from "../lib/KeyExpansion.js";
+import { encryptAES128, stringToHex } from "../utils/crypto.js";
 
 class IterativeAES128PublicInput extends Struct({
   cipher: Byte16,
 }) {}
 
+/**
+ * Computes the AES-128 encryption of a message using the given key.
+ *
+ * @param message The message to encrypt
+ * @param key The key to use for encryption
+ * @returns The encrypted message
+ */
 export function computeIterativeAes128Encryption(
   message: Byte16,
   key: Byte16,
@@ -36,6 +44,9 @@ export function computeIterativeAes128Encryption(
   return state;
 }
 
+/**
+ * A zkProgram that verifies a proof that a message was encrypted with AES-128 using the given key.
+ */
 const IterativeAes128 = ZkProgram({
   name: "aes-verify-iterative",
   publicInput: IterativeAES128PublicInput,
@@ -56,4 +67,44 @@ const IterativeAes128 = ZkProgram({
   },
 });
 
-export { IterativeAes128, IterativeAES128PublicInput };
+/**
+ * Generates a proof that the given message was encrypted with AES-128 using the given key.
+ * The key must be in hex form.
+ *
+ * @param message The message to generate a proof for
+ * @param keyHex The key to use for encryption in hex form
+ * @returns A proof that the message was encrypted with AES-128 using the given key and encrypted message
+ * @throws If the message is not 16 characters long or the key is not 32 characters long
+ * @throws If the proof generation fails
+ */
+// NO TEST NOW AS IT WILL CHANGE SOON
+async function generateIterativeAes128Proof(
+  message: string,
+  keyHex: string, // Should we allow non hex strings?
+): Promise<[Proof<IterativeAES128PublicInput, void>, string]> {
+  if (message.length !== 16) {
+    throw new Error("Message must be 16 characters long");
+  }
+
+  if (keyHex.length !== 32) {
+    throw new Error("Key must be 32 characters long");
+  }
+
+  const messageHex = stringToHex(message);
+  const cipher = encryptAES128(messageHex, keyHex);
+  const { proof } = await IterativeAes128.verifyAES128(
+    new IterativeAES128PublicInput({
+      cipher: Byte16.fromHex(cipher),
+    }),
+    Byte16.fromHex(messageHex),
+    Byte16.fromHex(keyHex),
+  );
+
+  return [proof, cipher];
+}
+
+export {
+  generateIterativeAes128Proof,
+  IterativeAes128,
+  IterativeAES128PublicInput,
+};

src/index.ts

@@ -5,3 +5,4 @@ import {
 import { Byte16 } from "./primitives/Bytes.js";
 
 export { IterativeAes128, IterativeAES128PublicInput, Byte16 };
+export { generateIterativeAes128Proof as generateAes128Proof } from "./implementations/IterativeAES128";

src/lib/KeyExpansion.ts

@@ -1,6 +1,6 @@
 import { Field, Gadgets } from "o1js";
 import { Byte16 } from "../primitives/Bytes.js";
-import { sbox_byte } from "./SBox.js";
+import { sboxByte } from "./SBox.js";
 
 // Each word consists of four 8-bit fields.
 export type Word = [Field, Field, Field, Field];
@@ -63,7 +63,7 @@ export function rotWord(word: Word): Word {
  * @returns The substituted 32-bit word.
  */
 export function subWord(word: Word): Word {
-  return word.map((field) => sbox_byte(field)) as Word;
+  return word.map((field) => sboxByte(field)) as Word;
 }
 
 /**

src/lib/SBox.ts

@@ -11,7 +11,7 @@ import {
  * @param {Field} input a byte represented within a field
  * @returns {Field} the substituted output
  */
-function sbox_byte(input: Field): Field {
+function sboxByte(input: Field): Field {
   const byte = RijndaelFiniteField.fromField(input);
   const byte_sbox = affineTransform(byte);
   return byte_sbox;
@@ -29,12 +29,12 @@ function sbox(input: Byte16): Byte16 {
   for (let i = 0; i < 4; i++) {
     const arr: Field[] = [];
     for (let j = 0; j < 4; j++) {
-      arr.push(sbox_byte(cols[i][j]));
+      arr.push(sboxByte(cols[i][j]));
     }
     newCols.push(arr);
   }
 
   return Byte16.fromColumns(newCols);
 }
 
-export { sbox, sbox_byte };
+export { sbox, sboxByte };

src/utils/crypto.ts

@@ -0,0 +1,30 @@
+import crypto from "crypto";
+
+/**
+ * Encrypts a message using AES-128 encryption.
+ *
+ * @param messageHex Plaintext in hex form to encrypt
+ * @param keyHex Key in hex form to use for encryption
+ * @returns The encrypted message in hex form
+ */
+export function encryptAES128(messageHex: string, keyHex: string): string {
+  const keyBuffer = Buffer.from(keyHex, "hex");
+  const plaintextBuffer = Buffer.from(messageHex, "hex");
+  const cipher = crypto.createCipheriv("aes-128-ecb", keyBuffer, null);
+  cipher.setAutoPadding(false);
+  const nodeEncrypted = Buffer.concat([
+    cipher.update(plaintextBuffer),
+    cipher.final(),
+  ]);
+  return nodeEncrypted.toString("hex");
+}
+
+/**
+ * Converts a string to its hex representation.
+ *
+ * @param str The string to convert to hex
+ * @returns The hex representation of the string
+ */
+export function stringToHex(str: string): string {
+  return Buffer.from(str).toString("hex");
+}

test/circuitSummary.ts

@@ -3,9 +3,10 @@ import {
   IterativeAes128,
   IterativeAES128PublicInput as AESPublicInput,
 } from "../src/implementations/IterativeAES128.js";
+import { Aes128Ctr } from "../src/implementations/AES128CTR.js";
 import { addRoundKey } from "../src/lib/AddRoundKey.js";
 import { mixColumn } from "../src/lib/MixColumns.js";
-import { sbox, sbox_byte } from "../src/lib/SBox.js";
+import { sbox, sboxByte } from "../src/lib/SBox.js";
 import { shiftRows } from "../src/lib/ShiftRows.js";
 import { Byte16 } from "../src/primitives/Bytes.js";
 import { expandKey128 } from "../src/lib/KeyExpansion.js";
@@ -42,7 +43,7 @@ const libZkProgram = ZkProgram({
     sboxByte: {
       privateInputs: [Field],
       async method(input_ignore: AESPublicInput, input: Field) {
-        sbox_byte(input);
+        sboxByte(input);
       },
     },
     expandKey128: {
@@ -58,10 +59,19 @@ const main = async () => {
   const { sboxByte, sbox, mixColumns, shiftRows, addRoundKey, expandKey128 } =
     await libZkProgram.analyzeMethods();
   const { verifyAES128 } = await IterativeAes128.analyzeMethods();
+  const { base, inductive } = await Aes128Ctr.analyzeMethods();
 
+  console.log("------------ Implementations Summary ------------");
   console.log("AES128 Iterative Summary:");
   console.log(verifyAES128.summary());
 
+  console.log("AES128CTR Base Summary:");
+  console.log(base.summary());
+
+  console.log("AES128CTR Inductive Summary:");
+  console.log(inductive.summary());
+
+  console.log("------------ Libraries Summary ------------");
   console.log("SBox Summary:");
   console.log(sbox.summary());
 

test/verifyAES128.test.ts → test/implementations/verifyAES128.test.ts

@@ -1,10 +1,10 @@
-import crypto, { randomBytes } from "crypto";
 import {
   computeIterativeAes128Encryption,
   IterativeAes128,
   IterativeAES128PublicInput as AESPublicInput,
-} from "../src/implementations/IterativeAES128.js";
-import { Byte16 } from "../src/primitives/Bytes.js";
+} from "../../src/implementations/IterativeAES128.js";
+import { encryptAES128 } from "../../src/utils/crypto.js";
+import { Byte16 } from "../../src/primitives/Bytes.js";
 import { verify } from "o1js";
 
 const RUN_ZK_TESTS = process.env.RUN_ZK_TESTS === "true";
@@ -20,16 +20,8 @@ const testVector1: TestVector = {
   keyHex: "2b7e151628aed2a6abf7158809cf4f3c",
 };
 
-const getCipherText = (tv: TestVector) => {
-  const keyBuffer = Buffer.from(tv.keyHex, "hex");
-  const plaintextBuffer = Buffer.from(tv.plaintextHex, "hex");
-  const cipher = crypto.createCipheriv("aes-128-ecb", keyBuffer, null);
-  cipher.setAutoPadding(false);
-  const nodeEncrypted = Buffer.concat([
-    cipher.update(plaintextBuffer),
-    cipher.final(),
-  ]);
-  return nodeEncrypted.toString("hex");
+const getCipherText = (tv: TestVector): string => {
+  return encryptAES128(tv.plaintextHex, tv.keyHex);
 };
 
 const testVectorToByte16 = (tv: TestVector) => ({
@@ -44,21 +36,6 @@ describe("Iterative AES128 Encryption", () => {
     expect(getCipherText(testVector1)).toBe(customCipher.toHex());
   });
 
-  it("should match random test vectors", () => {
-    // This produces random test vectors each time, so it's not deterministic
-    for (let i = 0; i < 200; i++) {
-      const tv: TestVector = {
-        plaintextHex: randomBytes(16).toString("hex"),
-        keyHex: randomBytes(16).toString("hex"),
-      };
-
-      const { plaintext, key } = testVectorToByte16(tv);
-
-      const provableCipher = computeIterativeAes128Encryption(plaintext, key);
-      expect(getCipherText(tv)).toBe(provableCipher.toHex());
-    }
-  });
-
   (RUN_ZK_TESTS ? it : it.skip)(
     "should verify the proof using the zkProgram",
     async () => {