From 4098d957643641526f97416b04a59af43ce343bd Mon Sep 17 00:00:00 2001
From: Olga Kunyavskaya <olga.kunyavskaya@aurora.dev>
Date: Wed, 24 Jul 2024 11:36:23 +0300
Subject: [PATCH] add checkPreMigration flag

---
 eth-custodian/contracts/EthCustodianProxy.sol | 19 ++++++++++++++-----
 eth-custodian/test/EthCustodianProxy.js       | 16 ++++++++--------
 2 files changed, 22 insertions(+), 13 deletions(-)

diff --git a/eth-custodian/contracts/EthCustodianProxy.sol b/eth-custodian/contracts/EthCustodianProxy.sol
index 601902f..5e9d347 100644
--- a/eth-custodian/contracts/EthCustodianProxy.sol
+++ b/eth-custodian/contracts/EthCustodianProxy.sol
@@ -96,17 +96,18 @@ contract EthCustodianProxy is
 
     function withdraw(
         bytes calldata proofData,
-        uint64 proofBlockHeight
+        uint64 proofBlockHeight,
+        bool checkPreMigration
     ) external {
-        if (BlockHeightFromProofExtractor.getBlockHeightFromProof(proofData) > migrationBlockHeight) {
-            _requireNotPaused(PAUSED_WITHDRAW_POST_MIGRATION);
-            ethCustodianImpl.withdraw(proofData, proofBlockHeight);
-        } else {
+        if (isPreMigration(proofData, checkPreMigration)) {
             _requireNotPaused(PAUSED_WITHDRAW_PRE_MIGRATION);
             bytes memory postMigrationProducer = ethCustodianImpl.nearProofProducerAccount_();
             _writeProofProducerSlot(preMigrationProducerAccount);
             ethCustodianImpl.withdraw(proofData, proofBlockHeight);
             _writeProofProducerSlot(postMigrationProducer);
+        } else {
+            _requireNotPaused(PAUSED_WITHDRAW_POST_MIGRATION);
+            ethCustodianImpl.withdraw(proofData, proofBlockHeight);
         }
     }
 
@@ -150,6 +151,14 @@ contract EthCustodianProxy is
         _pause(flags);
     }
 
+    function isPreMigration(bytes calldata proofData, bool checkPreMigration) internal view returns(bool) {
+        if (checkPreMigration) {
+            return BlockHeightFromProofExtractor.getBlockHeightFromProof(proofData) <= migrationBlockHeight;
+        }
+
+        return false;
+    }
+
     /**
      * @dev Internal function called by the proxy contract to authorize an upgrade to a new implementation address
      * using the UUPS proxy upgrade pattern. Overrides the default `_authorizeUpgrade` function from the `UUPSUpgradeable` contract.
diff --git a/eth-custodian/test/EthCustodianProxy.js b/eth-custodian/test/EthCustodianProxy.js
index e1112ec..2d49721 100644
--- a/eth-custodian/test/EthCustodianProxy.js
+++ b/eth-custodian/test/EthCustodianProxy.js
@@ -144,7 +144,7 @@ describe('EthCustodianProxy contract', () => {
             await ethCustodianProxy.pauseProxy(PAUSED_WITHDRAW_POST_MIGRATION);
             const proof = require('./proof_template_from_testnet.json');
 
-            await expect(ethCustodianProxy.withdraw(borshifyOutcomeProof(proof), blockHeightFromProof + 1))
+            await expect(ethCustodianProxy.withdraw(borshifyOutcomeProof(proof), blockHeightFromProof + 1, true))
                 .to.be.revertedWith('Pausable: paused');
         });
 
@@ -153,7 +153,7 @@ describe('EthCustodianProxy contract', () => {
             await ethCustodianProxy.pauseProxy(PAUSED_WITHDRAW_PRE_MIGRATION);
             const proof = require('./proof_template_from_testnet.json');
 
-            await expect(ethCustodianProxy.withdraw(borshifyOutcomeProof(proof), blockHeightFromProof))
+            await expect(ethCustodianProxy.withdraw(borshifyOutcomeProof(proof), blockHeightFromProof, true))
                 .to.be.revertedWith('Pausable: paused');
         });
 
@@ -169,10 +169,10 @@ describe('EthCustodianProxy contract', () => {
 
             const proof = require('./proof_template_from_testnet.json');
 
-            await expect(ethCustodianProxy.withdraw(borshifyOutcomeProof(proof), 1099))
+            await expect(ethCustodianProxy.withdraw(borshifyOutcomeProof(proof), 1099, true))
                 .to.be.revertedWith('Pausable: paused');
 
-            await expect(ethCustodianProxy.withdraw(borshifyOutcomeProof(proof), blockHeightFromProof + 1))
+            await expect(ethCustodianProxy.withdraw(borshifyOutcomeProof(proof), blockHeightFromProof + 1, true))
                 .to.be.revertedWith('Pausable: paused');
         });
     });
@@ -234,7 +234,7 @@ describe('EthCustodianProxy contract', () => {
                 await ethers.provider.getBalance(user2.address));
 
             await expect(
-                ethCustodianProxy.withdraw(borshifyOutcomeProof(postMigrationProof), blockHeightFromProof + 1)
+                ethCustodianProxy.withdraw(borshifyOutcomeProof(postMigrationProof), blockHeightFromProof + 1, true)
             )
                 .to.emit(ethCustodian, 'Withdrawn')
                 .withArgs(user2.address, amount);
@@ -252,7 +252,7 @@ describe('EthCustodianProxy contract', () => {
             const balanceBefore = BigInt(await ethers.provider.getBalance(user2.address));
 
             await expect(
-                ethCustodianProxy.withdraw(borshifyOutcomeProof(proof), blockHeightFromProof + 2)
+                ethCustodianProxy.withdraw(borshifyOutcomeProof(proof), blockHeightFromProof + 2, true)
             )
                 .to.emit(ethCustodian, 'Withdrawn')
                 .withArgs(user2.address, amount);
@@ -268,7 +268,7 @@ describe('EthCustodianProxy contract', () => {
             const balanceBefore = BigInt(await ethers.provider.getBalance(user2.address));
 
             await expect(
-                ethCustodianProxy.withdraw(borshifyOutcomeProof(proof), blockHeightFromProof)
+                ethCustodianProxy.withdraw(borshifyOutcomeProof(proof), blockHeightFromProof, true)
             )
                 .to.emit(ethCustodian, 'Withdrawn')
                 .withArgs(user2.address, amount);
@@ -286,7 +286,7 @@ describe('EthCustodianProxy contract', () => {
             const migrationBlock = await ethCustodianProxy.migrationBlockHeight();
 
             await expect(
-                ethCustodianProxy.withdraw(borshifyOutcomeProof(proof), parseInt(migrationBlock))
+                ethCustodianProxy.withdraw(borshifyOutcomeProof(proof), parseInt(migrationBlock), true)
             )
                 .to.emit(ethCustodian, 'Withdrawn')
                 .withArgs(user2.address, amount);