From e2bda0653299cbf3409ca901480926ab7c1bf1ff Mon Sep 17 00:00:00 2001
From: Florin Dzeladini <florin.dzeladini@consensys.net>
Date: Mon, 25 Nov 2024 10:57:42 +0100
Subject: [PATCH] fix: rollback request state in case of error in
 input-event-controller

---
 .../user-input-event-controller.test.ts       | 78 +++++++++++++++++++
 .../user-input-event-controller.ts            | 22 +++++-
 2 files changed, 97 insertions(+), 3 deletions(-)

diff --git a/packages/starknet-snap/src/ui/controllers/user-input-event-controller.test.ts b/packages/starknet-snap/src/ui/controllers/user-input-event-controller.test.ts
index 97836233..5c7b0583 100644
--- a/packages/starknet-snap/src/ui/controllers/user-input-event-controller.test.ts
+++ b/packages/starknet-snap/src/ui/controllers/user-input-event-controller.test.ts
@@ -410,6 +410,84 @@ describe('UserInputEventController', () => {
       },
     );
 
+    it.each([STRK_SEPOLIA_TESTNET, ETHER_SEPOLIA_TESTNET])(
+      'rollsback the transaction request with the original request state if state update fails: feeToken - %symbol',
+      async (token: Erc20Token) => {
+        const feeToken = FeeToken[token.symbol];
+        const {
+          event,
+          account,
+          network,
+          getEstimatedFeesSpy,
+          hasSufficientFundsForFeeSpy,
+          updateExecuteTxnFlowSpy,
+          mockGetEstimatedFeesResponse,
+          upsertTransactionRequestSpy,
+          transactionRequest,
+        } = await prepareHandleFeeTokenChange(feeToken);
+        const feeTokenAddress = token.address;
+        const { signer, calls } = transactionRequest;
+        const { publicKey, privateKey, address } = account;
+        const { suggestedMaxFee } = mockGetEstimatedFeesResponse;
+        const rollbackSnapshot = {
+          maxFee: transactionRequest.maxFee,
+          selectedFeeToken: transactionRequest.selectedFeeToken,
+          includeDeploy: transactionRequest.includeDeploy,
+          resourceBounds: [...transactionRequest.resourceBounds],
+        };
+
+        upsertTransactionRequestSpy.mockRejectedValue(new Error('Failed!'));
+
+        const controller = createMockController(event);
+        await controller.handleFeeTokenChange();
+
+        expect(getEstimatedFeesSpy).toHaveBeenCalledWith(
+          network,
+          signer,
+          privateKey,
+          publicKey,
+          [
+            {
+              type: TransactionType.INVOKE,
+              payload: calls.map((call) => ({
+                calldata: call.calldata,
+                contractAddress: call.contractAddress,
+                entrypoint: call.entrypoint,
+              })),
+            },
+          ],
+          {
+            version: controller.feeTokenToTransactionVersion(feeToken),
+          },
+        );
+        expect(hasSufficientFundsForFeeSpy).toHaveBeenCalledWith({
+          address,
+          network,
+          calls,
+          feeTokenAddress,
+          suggestedMaxFee,
+        });
+        // transactionRequest will be pass by reference, so we can use this to check the updated value
+        expect(transactionRequest.maxFee).toStrictEqual(suggestedMaxFee);
+        expect(updateExecuteTxnFlowSpy).toHaveBeenCalledWith(
+          controller.eventId,
+          transactionRequest,
+        );
+        expect(upsertTransactionRequestSpy).toHaveBeenCalledWith(
+          transactionRequest,
+        );
+        expect(updateExecuteTxnFlowSpy).toHaveBeenCalledWith(
+          controller.eventId,
+          { ...transactionRequest, ...rollbackSnapshot },
+          {
+            errors: {
+              fees: `Failed to calculate the fees, switching back to ${transactionRequest.selectedFeeToken}`,
+            },
+          },
+        );
+      },
+    );
+
     it('updates the transaction request with an insufficient funds error message if the account balance is insufficient to cover the fee.', async () => {
       const {
         event,
diff --git a/packages/starknet-snap/src/ui/controllers/user-input-event-controller.ts b/packages/starknet-snap/src/ui/controllers/user-input-event-controller.ts
index 8a5f47d3..28520697 100644
--- a/packages/starknet-snap/src/ui/controllers/user-input-event-controller.ts
+++ b/packages/starknet-snap/src/ui/controllers/user-input-event-controller.ts
@@ -130,8 +130,20 @@ export class UserInputEventController {
     return token.address;
   }
 
+  protected createRollbackSnapshot(
+    request: TransactionRequest,
+  ): Partial<TransactionRequest> {
+    return {
+      maxFee: request.maxFee,
+      selectedFeeToken: request.selectedFeeToken,
+      includeDeploy: request.includeDeploy,
+      resourceBounds: [...request.resourceBounds],
+    };
+  }
+
   protected async handleFeeTokenChange() {
     const request = this.context?.request as TransactionRequest;
+    const rollbackSnapshot = this.createRollbackSnapshot(request);
     const { addressIndex, calls, signer, chainId } = request;
     const feeToken = (this.event as InputChangeEvent)
       .value as unknown as FeeToken;
@@ -196,9 +208,13 @@ export class UserInputEventController {
           : `Failed to calculate the fees, switching back to ${request.selectedFeeToken}`;
 
       // On failure, display ExecuteTxnUI with an error message
-      await updateExecuteTxnFlow(this.eventId, request, {
-        errors: { fees: errorMessage },
-      });
+      await updateExecuteTxnFlow(
+        this.eventId,
+        { ...request, ...rollbackSnapshot },
+        {
+          errors: { fees: errorMessage },
+        },
+      );
     }
   }
 }