From 7c3c7d1cc63de75efe5cbf3ee5091793630e4d6f Mon Sep 17 00:00:00 2001
From: Evgeny Gusarov <evgeny@stakewise.io>
Date: Thu, 11 Jan 2024 14:27:21 +0300
Subject: [PATCH] Fix harvest params

---
 src/common/contracts.py |  3 +++
 src/common/ipfs.py      | 29 +++++++++++++++++++++--------
 src/config/networks.py  |  8 ++++++++
 3 files changed, 32 insertions(+), 8 deletions(-)

diff --git a/src/common/contracts.py b/src/common/contracts.py
index 37771e7f..bfee7955 100644
--- a/src/common/contracts.py
+++ b/src/common/contracts.py
@@ -105,6 +105,9 @@ async def get_registered_validators_public_keys(
         )
         return [Web3.to_hex(event['args']['publicKey']) for event in events]
 
+    async def mev_escrow(self) -> ChecksumAddress:
+        return await self.contract.functions.mevEscrow().call()
+
 
 class V2PoolContract(ContractWrapper):
     abi_path = 'abi/IV2Pool.json'
diff --git a/src/common/ipfs.py b/src/common/ipfs.py
index 432a896b..e8803c5a 100644
--- a/src/common/ipfs.py
+++ b/src/common/ipfs.py
@@ -5,7 +5,9 @@
 from web3.types import Wei
 
 from src.common.clients import ipfs_fetch_client
+from src.common.contracts import vault_contract
 from src.common.typings import HarvestParams
+from src.config.settings import settings
 
 logger = logging.getLogger(__name__)
 
@@ -14,19 +16,30 @@ async def fetch_harvest_params(
     vault_address: ChecksumAddress, ipfs_hash: str, rewards_root: bytes
 ) -> HarvestParams | None:
     ipfs_data = await ipfs_fetch_client.fetch_json(ipfs_hash)
+    mev_escrow = await vault_contract.mev_escrow()
+
     for vault_data in ipfs_data['vaults']:  # type: ignore
-        if vault_address == Web3.to_checksum_address(vault_data['vault']):
+        if vault_address != Web3.to_checksum_address(vault_data['vault']):
+            continue
+
+        if mev_escrow == settings.network_config.SHARED_MEV_ESCROW_CONTRACT_ADDRESS:
+            # shared mev vault
             unlocked_mev_reward = Wei(vault_data['unlocked_mev_reward'])
             reward = Wei(
                 vault_data['consensus_reward']
                 + unlocked_mev_reward
-                + vault_data.get('locked_mev_reward', 0)
-            )
-            return HarvestParams(
-                rewards_root=rewards_root,
-                reward=reward,
-                unlocked_mev_reward=unlocked_mev_reward,
-                proof=[Web3.to_bytes(hexstr=x) for x in vault_data['proof']],
+                + vault_data['locked_mev_reward']
             )
+        else:
+            # own mev vault
+            unlocked_mev_reward = Wei(0)
+            reward = Wei(vault_data['consensus_reward'])
+
+        return HarvestParams(
+            rewards_root=rewards_root,
+            reward=reward,
+            unlocked_mev_reward=unlocked_mev_reward,
+            proof=[Web3.to_bytes(hexstr=x) for x in vault_data['proof']],
+        )
 
     return None
diff --git a/src/config/networks.py b/src/config/networks.py
index b5822278..fb0b6d4a 100644
--- a/src/config/networks.py
+++ b/src/config/networks.py
@@ -37,6 +37,7 @@ class NetworkConfig:
     SHAPELLA_FORK_VERSION: bytes
     SHAPELLA_EPOCH: int
     MULTICALL_CONTRACT_ADDRESS: ChecksumAddress
+    SHARED_MEV_ESCROW_CONTRACT_ADDRESS: ChecksumAddress
 
     @property
     def SHAPELLA_FORK(self) -> ConsensusFork:
@@ -92,6 +93,9 @@ def IS_SUPPORT_V2_MIGRATION(self) -> bool:
         MULTICALL_CONTRACT_ADDRESS=Web3.to_checksum_address(
             '0xcA11bde05977b3631167028862bE2a173976CA11'
         ),
+        SHARED_MEV_ESCROW_CONTRACT_ADDRESS=Web3.to_checksum_address(
+            '0x48319f97E5Da1233c21c48b80097c0FB7a20Ff86'
+        ),
     ),
     HOLESKY: NetworkConfig(
         SYMBOL='HolETH',
@@ -123,6 +127,9 @@ def IS_SUPPORT_V2_MIGRATION(self) -> bool:
         MULTICALL_CONTRACT_ADDRESS=Web3.to_checksum_address(
             '0xcA11bde05977b3631167028862bE2a173976CA11'
         ),
+        SHARED_MEV_ESCROW_CONTRACT_ADDRESS=Web3.to_checksum_address(
+            '0xc98F25BcAA6B812a07460f18da77AF8385be7b56'
+        ),
     ),
     GNOSIS: NetworkConfig(
         SYMBOL='xDAI',
@@ -155,5 +162,6 @@ def IS_SUPPORT_V2_MIGRATION(self) -> bool:
         MULTICALL_CONTRACT_ADDRESS=Web3.to_checksum_address(
             '0xcA11bde05977b3631167028862bE2a173976CA11'
         ),
+        SHARED_MEV_ESCROW_CONTRACT_ADDRESS=Web3.to_checksum_address(EMPTY_ADDR_HEX),
     ),
 }