ZK-590: remove id from state of shielder-sdk (#20)

Cardinal-Cryptography · Dec 3, 2024 · c80b6b4 · c80b6b4
1 parent a59543d
commit c80b6b4
Show file tree

Hide file tree

Showing 5 changed files with 53 additions and 35 deletions.
diff --git a/ts/shielder-sdk-tests/tests/chain/contract/getMerklePath.test.ts b/ts/shielder-sdk-tests/tests/chain/contract/getMerklePath.test.ts
@@ -3,14 +3,14 @@ import { getChainConfig } from "@tests/chain/config";
 import { sdkTest } from "@tests/playwrightTestUtils";
 import { generatePrivateKey } from "viem/accounts";
 
-import type { Contract, Scalar } from "shielder-sdk/__internal__";
+import { type Contract, type Scalar } from "shielder-sdk/__internal__";
 
 // Custom test that creates:
 //  - `playwrightFixture`: an object initialized outside the browser environment,
 //  - `webFixture`: a `JSHandle` to an object accessible only in the browser environment.
 export const getMerklePathTest = sdkTest.extend<Fixtures>({
   // eslint-disable-next-line no-empty-pattern
-  playwrightFixture: async ({ }, use) => {
+  playwrightFixture: async ({}, use) => {
     const playwrightFixture = await createPlaywrightFixture();
     await use(playwrightFixture);
   },
@@ -65,7 +65,11 @@ async function createWebFixture({
     window.shielder.actions.createNewAccountAction(contract);
 
   const amount = 5n;
-  const state = await window.state.emptyAccountState(privateKeyAlice);
+  const state = await window.state.emptyAccountState(
+    await window.wasmClientWorker
+      .getWorker()
+      .privateKeyToScalar(privateKeyAlice),
+  );
 
   const newAccountCalldata = await newAccountAction.generateCalldata(
     state,

diff --git a/ts/shielder-sdk-tests/tests/chain/contract/newAccount.test.ts b/ts/shielder-sdk-tests/tests/chain/contract/newAccount.test.ts
@@ -3,9 +3,9 @@ import { getChainConfig } from "@tests/chain/config";
 import { sdkTest } from "@tests/playwrightTestUtils";
 import { generatePrivateKey } from "viem/accounts";
 
-import type {
-  AccountState,
-  NewAccountCalldata,
+import {
+  type AccountState,
+  type NewAccountCalldata,
 } from "shielder-sdk/__internal__";
 import type { ContractTestFixture } from "@/chain/testUtils";
 
@@ -14,7 +14,7 @@ import type { ContractTestFixture } from "@/chain/testUtils";
 //  - `webFixture`: a `JSHandle` to an object accessible only in the browser environment.
 export const newAccountTest = sdkTest.extend<Fixtures>({
   // eslint-disable-next-line no-empty-pattern
-  playwrightFixture: async ({ }, use) => {
+  playwrightFixture: async ({}, use) => {
     const playwrightFixture = await createPlaywrightFixture();
     await use(playwrightFixture);
   },
@@ -70,7 +70,11 @@ async function createWebFixture({
     window.shielder.actions.createNewAccountAction(contract);
 
   const amount = 5n;
-  const state = await window.state.emptyAccountState(privateKeyAlice);
+  const state = await window.state.emptyAccountState(
+    await window.wasmClientWorker
+      .getWorker()
+      .privateKeyToScalar(privateKeyAlice),
+  );
 
   const newAccountCalldata = await newAccountAction.generateCalldata(
     state,

diff --git a/ts/shielder-sdk-tests/web/EntryPoint.tsx b/ts/shielder-sdk-tests/web/EntryPoint.tsx
@@ -166,7 +166,7 @@ declare global {
     };
 
     state: {
-      emptyAccountState: (privateKey: `0x${string}`) => Promise<AccountState>;
+      emptyAccountState: (id: Scalar) => AccountState;
       stateChangingEvents: (
         state: AccountState,
         noteEvents: NoteEvent[],

diff --git a/ts/shielder-sdk/src/shielder/state/index.ts b/ts/shielder-sdk/src/shielder/state/index.ts
@@ -51,6 +51,8 @@ export const eventToTransaction = (event: NoteEvent): ShielderTransaction => {
 export class StateManager {
   private storage: StorageInterface;
   privateKey: Hex;
+  id: Scalar | undefined;
+  idHash: Scalar | undefined;
 
   constructor(privateKey: Hex, storage: StorageInterface) {
     this.privateKey = privateKey;
@@ -59,18 +61,16 @@ export class StateManager {
 
   async accountState(): Promise<AccountState> {
     const res = await this.storage.getItem("accountState");
+    const id = await this.getId();
     if (res) {
-      const obj = res;
-      if (
-        !scalarsEqual(
-          Scalar.fromBigint(BigInt(obj.id)),
-          await wasmClientWorker.privateKeyToScalar(this.privateKey)
-        )
-      ) {
-        throw new Error("Account id does not match private key.");
+      const expectedIdHash = await this.getIdHash();
+      const storageIdHash = Scalar.fromBigint(res.idHash);
+      if (!scalarsEqual(expectedIdHash, storageIdHash)) {
+        throw new Error("Id hash in storage does not matched the configured.");
       }
+      const obj = res;
       return {
-        id: Scalar.fromBigint(BigInt(obj.id)),
+        id,
         nonce: BigInt(obj.nonce),
         balance: BigInt(obj.balance),
         currentNote: Scalar.fromBigint(BigInt(obj.currentNote)),
@@ -80,39 +80,49 @@ export class StateManager {
             : undefined
       };
     }
-    return await emptyAccountState(this.privateKey);
+    return await this.emptyAccountState();
   }
 
   async updateAccountState(accountState: AccountState) {
     if (accountState.currentNoteIndex == undefined) {
       throw new Error("currentNoteIndex must be set.");
     }
-    if (
-      !scalarsEqual(
-        accountState.id,
-        await wasmClientWorker.privateKeyToScalar(this.privateKey)
-      )
-    ) {
-      throw new Error("Account id does not match private key.");
+    if (!scalarsEqual(accountState.id, await this.getId())) {
+      throw new Error("New account id does not match the configured.");
     }
     await this.storage.setItem("accountState", {
-      ...accountState,
-      id: scalarToBigint(accountState.id),
-      currentNote: scalarToBigint(accountState.currentNote)
+      idHash: scalarToBigint(
+        await wasmClientWorker.poseidonHash([accountState.id])
+      ),
+      nonce: accountState.nonce,
+      balance: accountState.balance,
+      currentNote: scalarToBigint(accountState.currentNote),
+      currentNoteIndex: accountState.currentNoteIndex
     });
   }
 
   async emptyAccountState() {
-    return await emptyAccountState(this.privateKey);
+    return emptyAccountState(await this.getId());
+  }
+
+  private async getId(): Promise<Scalar> {
+    if (!this.id) {
+      this.id = await wasmClientWorker.privateKeyToScalar(this.privateKey);
+    }
+    return this.id;
+  }
+  private async getIdHash(): Promise<Scalar> {
+    if (!this.idHash) {
+      this.idHash = await wasmClientWorker.poseidonHash([await this.getId()]);
+    }
+    return this.idHash;
   }
 }
 
-const emptyAccountState = async (
-  privateKey: `0x${string}`
-): Promise<AccountState> => {
+const emptyAccountState = (id: Scalar): AccountState => {
   return {
     /// Since the private key is an arbitrary 32byte number, this is a non-reversible mapping
-    id: await wasmClientWorker.privateKeyToScalar(privateKey),
+    id,
     nonce: 0n,
     balance: 0n,
     currentNote: Scalar.fromBigint(0n)

diff --git a/ts/shielder-sdk/src/shielder/state/storageSchema.ts b/ts/shielder-sdk/src/shielder/state/storageSchema.ts
@@ -19,7 +19,7 @@ const storageSchema = {
   accountState: z.object({
     nonce: validateBigInt,
     balance: validateBigInt,
-    id: validateBigInt,
+    idHash: validateBigInt,
     currentNote: validateBigInt,
     currentNoteIndex: z.union([validateBigInt, z.undefined()]).optional()
   })