Wrap all controller wasm calls in mutex

cartridge-gg · Dec 10, 2024 · b2ec76c · b2ec76c
1 parent cdec774
commit b2ec76c
Show file tree

Hide file tree

Showing 4 changed files with 81 additions and 45 deletions.
diff --git a/packages/keychain/src/utils/connection/execute.ts b/packages/keychain/src/utils/connection/execute.ts
@@ -11,7 +11,6 @@ import {
 } from "starknet";
 import { ConnectionCtx, ControllerError, ExecuteCtx } from "./types";
 import { ErrorCode, JsCall } from "@cartridge/account-wasm/controller";
-import { mutex } from "./sync";
 
 export const ESTIMATE_FEE_PERCENTAGE = 10;
 
@@ -65,7 +64,6 @@ export function execute({
       });
     }
 
-    const release = await mutex.obtain();
     return await new Promise<InvokeFunctionResponse | ConnectError>(
       async (resolve, reject) => {
         // If a session call and there is no session available
@@ -151,9 +149,7 @@ export function execute({
           });
         }
       },
-    ).finally(() => {
-      release();
-    });
+    );
   };
 }
 

diff --git a/packages/keychain/src/utils/connection/sign.ts b/packages/keychain/src/utils/connection/sign.ts
@@ -5,7 +5,6 @@ import {
 } from "@cartridge/controller";
 import { Signature, TypedData } from "starknet";
 import { ConnectionCtx, SignMessageCtx } from "./types";
-import { mutex } from "./sync";
 import Controller from "utils/controller";
 import { parseControllerError } from "./execute";
 
@@ -30,7 +29,6 @@ export function signMessageFactory(setContext: (ctx: ConnectionCtx) => void) {
       });
     }
 
-    const release = await mutex.obtain();
     return await new Promise<Signature | ConnectError>(
       async (resolve, reject) => {
         // If a session call and there is no session available
@@ -64,8 +62,6 @@ export function signMessageFactory(setContext: (ctx: ConnectionCtx) => void) {
           });
         }
       },
-    ).finally(() => {
-      release();
-    });
+    );
   };
 }
diff --git a/packages/keychain/src/utils/connection/sync.ts b/packages/keychain/src/utils/connection/sync.ts
diff --git a/packages/keychain/src/utils/controller.ts b/packages/keychain/src/utils/controller.ts
@@ -13,9 +13,11 @@ import {
   Call,
   CallData,
 } from "starknet";
-
+import { Mutex } from "utils/mutex";
 import { toWasmPolicies } from "@cartridge/controller";
 
+const mutex = new Mutex();
+
 import {
   CartridgeAccount,
   JsCall,
@@ -92,7 +94,12 @@ export default class Controller extends Account {
       throw new Error("Account not found");
     }
 
-    await this.cartridge.createSession(toWasmPolicies(policies), expiresAt);
+    const release = await mutex.obtain();
+    try {
+      await this.cartridge.createSession(toWasmPolicies(policies), expiresAt);
+    } finally {
+      release();
+    }
   }
 
   registerSessionCalldata(
@@ -117,24 +124,39 @@ export default class Controller extends Account {
       throw new Error("Account not found");
     }
 
-    return await this.cartridge.registerSession(
-      toWasmPolicies(policies),
-      expiresAt,
-      publicKey,
-      num.toHex(maxFee),
-    );
+    const release = await mutex.obtain();
+    try {
+      return await this.cartridge.registerSession(
+        toWasmPolicies(policies),
+        expiresAt,
+        publicKey,
+        num.toHex(maxFee),
+      );
+    } finally {
+      release();
+    }
   }
 
   upgrade(new_class_hash: JsFelt): JsCall {
     return this.cartridge.upgrade(new_class_hash);
   }
 
   async executeFromOutsideV2(calls: Call[]): Promise<InvokeFunctionResponse> {
-    return await this.cartridge.executeFromOutsideV2(toJsCalls(calls));
+    const release = await mutex.obtain();
+    try {
+      return await this.cartridge.executeFromOutsideV2(toJsCalls(calls));
+    } finally {
+      release();
+    }
   }
 
   async executeFromOutsideV3(calls: Call[]): Promise<InvokeFunctionResponse> {
-    return await this.cartridge.executeFromOutsideV3(toJsCalls(calls));
+    const release = await mutex.obtain();
+    try {
+      return await this.cartridge.executeFromOutsideV3(toJsCalls(calls));
+    } finally {
+      release();
+    }
   }
 
   async execute(
@@ -149,10 +171,15 @@ export default class Controller extends Account {
       executionDetails.maxFee = num.toHex(executionDetails.maxFee);
     }
 
-    return await this.cartridge.execute(
-      toJsCalls(calls),
-      executionDetails as JsInvocationsDetails,
-    );
+    const release = await mutex.obtain();
+    try {
+      return await this.cartridge.execute(
+        toJsCalls(calls),
+        executionDetails as JsInvocationsDetails,
+      );
+    } finally {
+      release();
+    }
   }
 
   hasSession(calls: Call[]): boolean {
@@ -174,22 +201,27 @@ export default class Controller extends Account {
     calls: Call[],
     _: EstimateFeeDetails = {},
   ): Promise<EstimateFee> {
-    const res = await this.cartridge.estimateInvokeFee(toJsCalls(calls));
-
-    // The reason why we set the multiplier unseemingly high is to account
-    // for the fact that the estimation above is done without validation (ie SKIP_VALIDATE).
-    //
-    // Setting it lower might cause the actual transaction to fail due to
-    // insufficient max fee.
-    const MULTIPLIER_PERCENTAGE = 170; // x1.7
-
-    // This will essentially multiply the estimated fee by 1.7
-    const suggestedMaxFee = num.addPercent(
-      BigInt(res.overall_fee),
-      MULTIPLIER_PERCENTAGE,
-    );
+    const release = await mutex.obtain();
+    try {
+      const res = await this.cartridge.estimateInvokeFee(toJsCalls(calls));
+
+      // The reason why we set the multiplier unseemingly high is to account
+      // for the fact that the estimation above is done without validation (ie SKIP_VALIDATE).
+      //
+      // Setting it lower might cause the actual transaction to fail due to
+      // insufficient max fee.
+      const MULTIPLIER_PERCENTAGE = 170; // x1.7
+
+      // This will essentially multiply the estimated fee by 1.7
+      const suggestedMaxFee = num.addPercent(
+        BigInt(res.overall_fee),
+        MULTIPLIER_PERCENTAGE,
+      );
 
-    return { suggestedMaxFee, ...res };
+      return { suggestedMaxFee, ...res };
+    } finally {
+      release();
+    }
   }
 
   async verifyMessageHash(
@@ -211,15 +243,30 @@ export default class Controller extends Account {
   }
 
   async signMessage(typedData: TypedData): Promise<Signature> {
-    return this.cartridge.signMessage(JSON.stringify(typedData));
+    const release = await mutex.obtain();
+    try {
+      return await this.cartridge.signMessage(JSON.stringify(typedData));
+    } finally {
+      release();
+    }
   }
 
   async getNonce(_?: any): Promise<string> {
-    return await this.cartridge.getNonce();
+    const release = await mutex.obtain();
+    try {
+      return await this.cartridge.getNonce();
+    } finally {
+      release();
+    }
   }
 
   async delegateAccount(): Promise<string> {
-    return this.cartridge.delegateAccount();
+    const release = await mutex.obtain();
+    try {
+      return await this.cartridge.delegateAccount();
+    } finally {
+      release();
+    }
   }
 
   revoke(_origin: string) {