From d8c5e5661ed7a0aeac4ba0b060a4e8dec0f879f6 Mon Sep 17 00:00:00 2001
From: yash25198 <ydumpala@gmail.com>
Date: Sat, 23 Nov 2024 22:18:25 +0530
Subject: [PATCH] added tests

---
 circuits/hkdf.circom | 46 +++++++++++++++++++++++---------------------
 tests/add.test.ts    | 20 -------------------
 tests/hkdf.test.ts   | 34 ++++++++++++++++++++++++++++++++
 3 files changed, 58 insertions(+), 42 deletions(-)
 delete mode 100644 tests/add.test.ts
 create mode 100644 tests/hkdf.test.ts

diff --git a/circuits/hkdf.circom b/circuits/hkdf.circom
index bcdd959..b4c77ec 100644
--- a/circuits/hkdf.circom
+++ b/circuits/hkdf.circom
@@ -1,8 +1,12 @@
 pragma circom 2.1.5;
 
 include "./hmac/circuits/hmac.circom";
-import "circomlib/circuits/comparators.circom";
 
+// ss : secret length
+// is : info length
+// k : key length
+// m : number of keys to extract
+// s : key length
 template HKDFSha256(ss,is,k,m,s){
   signal input info[is];
   signal input secret[ss];
@@ -42,14 +46,14 @@ template Expand(n,k){
 // m : number of keys to extract
 // s : key length
 template Extract(n,k,m,s){
-  signal input info[n]
-  signal input key[k]
-  signal counter = 1; // counter is byte(1)
-  signal size = 32 + n + 1; // 32 bytes for hmac, n bytes for info, 1 byte for counter
+  signal input info[n];
+  signal input key[k];
+  
+  var size = 32 + n + 1; // 32 bytes for hmac, n bytes for info, 1 byte for counter
 
   // hash size is 32 bytes 
-  signal rounds = (32 * m)\s;
-  rounds = (rounds * s) < (32 * m) ? rounds + 1 : rounds;
+  var rounds = (m*s)\(32);
+  rounds = (rounds * 32) < (m*s) ? rounds + 1 : rounds;
 
 
   component hmac[rounds];
@@ -57,34 +61,32 @@ template Extract(n,k,m,s){
   signal expandedKeys [rounds][32];
   signal output out[m][s];
 
-  hmac[0] = HmacSha256(0, k);
+  hmac[0] = HmacSha256(1, k);
+  hmac[0].message[0] <== 1; // here counter is byte(1)
   hmac[0].key <== key;
   expandedKeys[0] <== hmac[0].hmac;
+  
+  var counter = 2; // counter is byte(2)
 
   for(var i = 1; i < rounds; i++){
     hmac[i] = HmacSha256(size, k);
-    for (var j = 0; j < n; j++){
+    for (var j = 0; j < 32; j++){
       hmac[i].message[j] <== expandedKeys[i-1][j];
     }
-    for (var j = 0; j < 32; j++){
-      hmac[i].message[n+j] <== info[j];
+    for (var j = 0; j < n; j++){
+      hmac[i].message[32+j] <== info[j];
     }
     hmac[i].message[32+n] <== counter;
     hmac[i].key <== key;
     expandedKeys[i] <== hmac[i].hmac;
-    counter <== counter + 1;
+    counter = counter + 1;
   }
 
-  signal xindex = 0;
-  signal yindex = 0;
-  for(var i = 0; i < m; i++){
-    for(var j = 0; j < s; j++){
-      out[i][j] <== expandedKeys[xindex][yindex];
-      yindex = yindex + 1;
-      if(yindex == 32){
-        xindex = xindex + 1;
-        yindex = 0;
-      }
+  var byteIndex = 0;
+  for (var i = 0; i < m; i++) {
+    for (var j = 0; j < s; j++) {
+      out[i][j] <== expandedKeys[byteIndex \ 32][byteIndex % 32];
+      byteIndex++;
     }
   }
 }
diff --git a/tests/add.test.ts b/tests/add.test.ts
deleted file mode 100644
index 0e7cc01..0000000
--- a/tests/add.test.ts
+++ /dev/null
@@ -1,20 +0,0 @@
-import { WitnessTester } from "circomkit";
-import { circomkit } from "./common";
-
-describe("Add", () => {
-  let circuit: WitnessTester<["a", "b"], ["out"]>;
-
-  describe("Add", () => {
-    before(async () => {
-      circuit = await circomkit.WitnessTester(`Add`, {
-        file: "add",
-        template: "Add",
-      });
-      console.log("#constraints:", await circuit.getConstraintCount());
-    });
-
-    it("should add two numbers", async () => {
-      await circuit.expectPass({ a: 1, b: 2 }, { out: 3 });
-    });
-  });
-});
diff --git a/tests/hkdf.test.ts b/tests/hkdf.test.ts
new file mode 100644
index 0000000..14ea821
--- /dev/null
+++ b/tests/hkdf.test.ts
@@ -0,0 +1,34 @@
+import { WitnessTester } from "circomkit";
+import { circomkit } from "./common";
+
+describe("HKDF", () => {
+  describe("Extract", () => {
+    let circuit: WitnessTester<["info", "key"], ["out"]>;
+    before(async () => {
+      circuit = await circomkit.WitnessTester(`Extract`, {
+        file: "hkdf",
+        template: "Extract",
+        params: [0, 32, 2, 16],
+      });
+      console.log("#constraints:", await circuit.getConstraintCount());
+    });
+
+    it("should extract two 16 bytes keys from key", async () => {
+      await circuit.expectPass(
+        {
+          info: [],
+          key: [
+            0x8b, 0xeb, 0x33, 0x8d, 0x43, 0x1d, 0x24, 0x3c, 0xee, 0xaa, 0xa6, 0xf0, 0xcb, 0x57, 0x26, 0xfb, 0xc5, 0xa3,
+            0x5c, 0x5e, 0x45, 0xbf, 0x99, 0x2c, 0xc3, 0xe2, 0x3b, 0x5b, 0xc2, 0xe4, 0xcc, 0xea,
+          ],
+        },
+        {
+          out: [
+            [0x5b, 0x02, 0xd2, 0x11, 0x3a, 0xbb, 0x74, 0x49, 0xc3, 0x7d, 0x57, 0xe0, 0xc7, 0x7a, 0x99, 0xc4],
+            [0x43, 0x7a, 0xb4, 0xc1, 0x85, 0x2f, 0xa9, 0xcc, 0x8e, 0xc5, 0xbd, 0x64, 0x97, 0xf0, 0x31, 0x91],
+          ],
+        }
+      );
+    });
+  });
+});