From 6b6ef5324f93a753ec25d128337500b3f2638e57 Mon Sep 17 00:00:00 2001
From: jordan <41rjordan@proton.me>
Date: Fri, 15 Nov 2024 23:46:31 +0700
Subject: [PATCH] push state to the calldata for easy reconstruction

---
 .github/workflows/ci.yaml                |  9 ++++-----
 contracts/src/UntronCore.sol             | 10 ++++++----
 contracts/src/interfaces/IUntronCore.sol |  3 ++-
 3 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
index b9f5e68..b1c4caf 100644
--- a/.github/workflows/ci.yaml
+++ b/.github/workflows/ci.yaml
@@ -17,11 +17,10 @@ jobs:
         with:
           submodules: recursive
 
-      - name: Download foundryup-zksync
-        run: curl -L https://raw.githubusercontent.com/matter-labs/foundry-zksync/main/install-foundry-zksync | bash
-
       - name: Install Foundry
-        run: foundryup-zksync
+        uses: foundry-rs/foundry-toolchain@v1
+        with:
+          version: nightly
 
       - name: Run tests
-        run: cd contracts && forge test -vvv --zksync
\ No newline at end of file
+        run: forge test -vvv
\ No newline at end of file
diff --git a/contracts/src/UntronCore.sol b/contracts/src/UntronCore.sol
index a7b782b..4146ce7 100644
--- a/contracts/src/UntronCore.sol
+++ b/contracts/src/UntronCore.sol
@@ -316,10 +316,8 @@ contract UntronCore is Initializable, OwnableUpgradeable, UntronTransfers, Untro
         }
     }
 
-    /// @notice Closes the orders and sends the funds to the providers or order creators, if not fulfilled.
-    /// @param proof The ZK proof.
-    /// @param publicValues The public values for the proof and order closure.
-    function closeOrders(bytes calldata proof, bytes calldata publicValues) external {
+    /// @inheritdoc IUntronCore
+    function closeOrders(bytes calldata proof, bytes calldata publicValues, bytes calldata newState) external {
         // verify the ZK proof with the public values
         // verifying logic is defined in the UntronZK contract.
         // currently it wraps SP1 zkVM verifier.
@@ -339,6 +337,10 @@ contract UntronCore is Initializable, OwnableUpgradeable, UntronTransfers, Untro
         // check that the old state hash is equal to the current state hash
         // this is needed to prevent the relayer from modifying the state in the ZK program.
         require(oldStateHash == stateHash, "Old state hash is invalid");
+        // check that the new state is valid
+        // we need to pass the new state in the calldata for easy state reconstruction by the relayers.
+        // on ZKsync Era (our deployment chain), the calldata is not published to the L1, so it's cheap to pass.
+        require(newStateHash == sha256(newState), "New state is invalid");
 
         // update the state hash
         stateHash = newStateHash;
diff --git a/contracts/src/interfaces/IUntronCore.sol b/contracts/src/interfaces/IUntronCore.sol
index 64d29ae..3fb563c 100644
--- a/contracts/src/interfaces/IUntronCore.sol
+++ b/contracts/src/interfaces/IUntronCore.sol
@@ -173,7 +173,8 @@ interface IUntronCore is IUntronTransfers, IUntronFees, IUntronZK {
     /// @notice Closes the orders and sends the funds to the providers or order creators, if not fulfilled.
     /// @param proof The ZK proof.
     /// @param publicValues The public values for the proof and order closure.
-    function closeOrders(bytes calldata proof, bytes calldata publicValues) external;
+    /// @param newState The new state of the ZK engine.
+    function closeOrders(bytes calldata proof, bytes calldata publicValues, bytes calldata newState) external;
 
     /// @notice Sets the liquidity provider details.
     /// @param liquidity The liquidity of the provider in USDT L2.