feat(Dataworker): ZkSync should count ETH and WETH towards balance needed to execute leaves

src/clients/BalanceAllocator.ts

@@ -1,4 +1,5 @@
-import { BigNumber, ERC20, ethers, ZERO_ADDRESS, min } from "../utils";
+import { CONTRACT_ADDRESSES } from "../common/ContractAddresses";
+import { BigNumber, ERC20, ethers, ZERO_ADDRESS, min, isDefined } from "../utils";
 
 // This type is used to map used and current balances of different users.
 export interface BalanceMap {
@@ -141,9 +142,19 @@ export class BalanceAllocator {
     this.balances = {};
   }
 
+  isEthAddress(chainId: number, tokenAddress: string): boolean {
+    // If there is an ETH address defined in CONTRACT_ADDRESSES, use it, otherwise assume ETH address is the zero
+    // address.
+    if (isDefined(CONTRACT_ADDRESSES[chainId]?.eth?.address)) {
+      return CONTRACT_ADDRESSES[chainId].eth.address === tokenAddress;
+    } else {
+      return ZERO_ADDRESS === tokenAddress;
+    }
+  }
+
   // This method is primarily here to be overriden for testing purposes.
   protected async _queryBalance(chainId: number, token: string, holder: string): Promise<BigNumber> {
-    return token === ZERO_ADDRESS
+    return this.isEthAddress(chainId, token)
       ? await this.providers[chainId].getBalance(holder)
       : await ERC20.connect(token, this.providers[chainId]).balanceOf(holder);
   }

src/clients/bridges/OptimismAdapter.ts

@@ -9,8 +9,6 @@ import { constants } from "@across-protocol/sdk-v2";
 import { CONTRACT_ADDRESSES } from "../../common";
 const { TOKEN_SYMBOLS_MAP } = constants;
 
-export const isOvmChain = (chainId: number): boolean => [10, 288].includes(chainId);
-
 export class OptimismAdapter extends BaseAdapter {
   public l2Gas: number;
   private txnClient: TransactionClient;

src/common/ContractAddresses.ts

@@ -3,7 +3,7 @@
   [chainId: number]: {
     [contractName: string]: {
       address?: string;
       abi?: any[];
     };
   };
 } = {
@@ -415,6 +415,14 @@
       ],
     },
   },
+  324: {
+    eth: {
+      address: "0x000000000000000000000000000000000000800A",
+    },
+    weth: {
+      address: "0x5AEa5775959fBC2557Cc8789bC1bf90A239D9a91",
+    },
+  },
   42161: {
     erc20Gateway: {
       abi: [

src/dataworker/Dataworker.ts

@@ -37,6 +37,7 @@ import {
   blockRangesAreInvalidForSpokeClients,
   getBlockRangeForChain,
   getImpliedBundleBlockRanges,
+  l2TokensToCountTowardsSpokePoolLeafExecutionCapital,
 } from "../dataworker/DataworkerUtils";
 import {
   getEndBlockBuffers,
@@ -46,8 +47,7 @@ import {
 } from "./DataworkerUtils";
 import { BalanceAllocator } from "../clients";
 import _ from "lodash";
-import { CONTRACT_ADDRESSES, spokePoolClientsToProviders } from "../common";
-import { isOvmChain } from "../clients/bridges";
+import { spokePoolClientsToProviders } from "../common";
 import * as sdk from "@across-protocol/sdk-v2";
 
 // Internal error reasons for labeling a pending root bundle as "invalid" that we don't want to submit a dispute
@@ -72,8 +72,6 @@ type ProposeRootBundleReturnType = {
   slowFillTree: MerkleTree<SlowFillLeaf>;
 };
 
-const ovmWethTokens = [CONTRACT_ADDRESSES[10].weth.address, CONTRACT_ADDRESSES[10].eth.address];
-
 export type PoolRebalanceRoot = {
   runningBalances: RunningBalances;
   realizedLpFees: RunningBalances;
@@ -1522,9 +1520,10 @@ export class Dataworker {
           const amountRequired = getRefund(relayData.amount.sub(amountFilled), relayData.realizedLpFeePct);
           const success = await balanceAllocator.requestBalanceAllocation(
             relayData.destinationChainId,
-            isOvmChain(relayData.destinationChainId) && ovmWethTokens.includes(relayData.destinationToken)
-              ? ovmWethTokens
-              : [relayData.destinationToken],
+            l2TokensToCountTowardsSpokePoolLeafExecutionCapital(
+              relayData.destinationToken,
+              relayData.destinationChainId
+            ),
             client.spokePool.address,
             amountRequired
           );
@@ -2096,9 +2095,7 @@ export class Dataworker {
           const totalSent = refundSum.add(leaf.amountToReturn.gte(0) ? leaf.amountToReturn : BigNumber.from(0));
           const success = await balanceAllocator.requestBalanceAllocation(
             leaf.chainId,
-            isOvmChain(leaf.chainId) && ovmWethTokens.includes(leaf.l2TokenAddress)
-              ? ovmWethTokens
-              : [leaf.l2TokenAddress],
+            l2TokensToCountTowardsSpokePoolLeafExecutionCapital(leaf.l2TokenAddress, leaf.chainId),
             client.spokePool.address,
             totalSent
           );

src/dataworker/DataworkerUtils.ts

@@ -15,6 +15,7 @@ import {
   buildPoolRebalanceLeafTree,
   buildRelayerRefundTree,
   buildSlowRelayTree,
+  isDefined,
   MerkleTree,
   winston,
 } from "../utils";
@@ -28,6 +29,7 @@ import { subtractExcessFromPreviousSlowFillsFromRunningBalances } from "./PoolRe
 import { getAmountToReturnForRelayerRefundLeaf } from "./RelayerRefundUtils";
 import { sortRefundAddresses, sortRelayerRefundLeaves } from "./RelayerRefundUtils";
 import { utils } from "@across-protocol/sdk-v2";
+import { CONTRACT_ADDRESSES } from "../common/ContractAddresses";
 export const { getImpliedBundleBlockRanges, getBlockRangeForChain, getBlockForChain } = utils;
 
 export function getEndBlockBuffers(
@@ -387,3 +389,26 @@ export async function _buildPoolRebalanceRoot(
     tree: buildPoolRebalanceLeafTree(leaves),
   };
 }
+
+// Some SpokePools will contain balances of ETH and WETH, while others will only contain balances of WETH,
+// so if the l2TokenAddress is WETH and the SpokePool is one such chain that holds both ETH and WETH,
+// then it should return both ETH and WETH. For other chains, it should only return the l2TokenAddress.
+function getWethAndEth(chainId): string[] {
+  const wethAndEth = [CONTRACT_ADDRESSES[chainId].weth.address, CONTRACT_ADDRESSES[chainId].eth.address];
+  if (wethAndEth.some((chainId) => !isDefined(chainId))) {
+    throw new Error(`WETH or ETH address not defined for chain ${chainId}`);
+  }
+  return wethAndEth;
+}
+export function l2TokensToCountTowardsSpokePoolLeafExecutionCapital(
+  l2TokenAddress: string,
+  l2ChainId: number
+): string[] {
+  const spokesThatHoldEthAndWeth = [10, 324];
+  if (!spokesThatHoldEthAndWeth.includes(l2ChainId)) {
+    return [l2TokenAddress];
+  }
+  // If we get to here, ETH and WETH addresses should be defined, or we'll throw an error.
+  const ethAndWeth = getWethAndEth(l2ChainId);
+  return ethAndWeth.includes(l2TokenAddress) ? ethAndWeth : [l2TokenAddress];
+}