feat(balance_overrides): use debug_traceCall to find slots

crates/contracts/artifacts/NonStandardERC20Balances.json

@@ -0,0 +1,132 @@
+{
+  "abi": [
+    {
+      "inputs": [
+        {
+          "internalType": "address",
+          "name": "user",
+          "type": "address"
+        },
+        {
+          "internalType": "address",
+          "name": "spender",
+          "type": "address"
+        }
+      ],
+      "name": "allowance",
+      "outputs": [
+        {
+          "internalType": "uint256",
+          "name": "",
+          "type": "uint256"
+        }
+      ],
+      "stateMutability": "view",
+      "type": "function"
+    },
+    {
+      "inputs": [
+        {
+          "internalType": "address",
+          "name": "spender",
+          "type": "address"
+        },
+        {
+          "internalType": "uint256",
+          "name": "amount",
+          "type": "uint256"
+        }
+      ],
+      "name": "approve",
+      "outputs": [],
+      "stateMutability": "nonpayable",
+      "type": "function"
+    },
+    {
+      "inputs": [
+        {
+          "internalType": "address",
+          "name": "user",
+          "type": "address"
+        }
+      ],
+      "name": "balanceOf",
+      "outputs": [
+        {
+          "internalType": "uint256",
+          "name": "",
+          "type": "uint256"
+        }
+      ],
+      "stateMutability": "view",
+      "type": "function"
+    },
+    {
+      "inputs": [
+        {
+          "internalType": "address",
+          "name": "user",
+          "type": "address"
+        },
+        {
+          "internalType": "uint256",
+          "name": "amount",
+          "type": "uint256"
+        }
+      ],
+      "name": "mint",
+      "outputs": [],
+      "stateMutability": "nonpayable",
+      "type": "function"
+    },
+    {
+      "inputs": [
+        {
+          "internalType": "address",
+          "name": "to",
+          "type": "address"
+        },
+        {
+          "internalType": "uint256",
+          "name": "amount",
+          "type": "uint256"
+        }
+      ],
+      "name": "transfer",
+      "outputs": [],
+      "stateMutability": "nonpayable",
+      "type": "function"
+    },
+    {
+      "inputs": [
+        {
+          "internalType": "address",
+          "name": "from",
+          "type": "address"
+        },
+        {
+          "internalType": "address",
+          "name": "to",
+          "type": "address"
+        },
+        {
+          "internalType": "uint256",
+          "name": "amount",
+          "type": "uint256"
+        }
+      ],
+      "name": "transferFrom",
+      "outputs": [],
+      "stateMutability": "nonpayable",
+      "type": "function"
+    }
+  ],
+  "bytecode": "0x6080604052348015600e575f5ffd5b506105398061001c5f395ff3fe608060405234801561000f575f5ffd5b506004361061006f575f3560e01c806370a082311161004d57806370a082311461010b578063a9059cbb14610155578063dd62ed3e14610168575f5ffd5b8063095ea7b31461007357806323b872dd146100bb57806340c10f19146100ce575b5f5ffd5b6100b96100813660046103c4565b335f9081526020818152604080832073ffffffffffffffffffffffffffffffffffffffff959095168352600190940190529190912055565b005b6100b96100c93660046103ec565b61017b565b6100b96100dc3660046103c4565b73ffffffffffffffffffffffffffffffffffffffff9091165f9081526020819052604090206001815560020155565b610143610119366004610426565b73ffffffffffffffffffffffffffffffffffffffff165f9081526020819052604090206002015490565b60405190815260200160405180910390f35b6100b96101633660046103c4565b6102a7565b610143610176366004610446565b610361565b73ffffffffffffffffffffffffffffffffffffffff83165f90815260208181526040808320338452600101909152812080548392906101bb9084906104a4565b909155505073ffffffffffffffffffffffffffffffffffffffff83165f90815260208190526040812060020180548392906101f79084906104a4565b909155505073ffffffffffffffffffffffffffffffffffffffff83165f90815260208190526040812080549161022c836104b7565b909155505073ffffffffffffffffffffffffffffffffffffffff82165f90815260208190526040812060020180548392906102689084906104ee565b909155505073ffffffffffffffffffffffffffffffffffffffff82165f90815260208190526040812080549161029d836104b7565b9190505550505050565b335f90815260208190526040812060020180548392906102c89084906104a4565b9091555050335f9081526020819052604081208054916102e7836104b7565b909155505073ffffffffffffffffffffffffffffffffffffffff82165f90815260208190526040812060020180548392906103239084906104ee565b909155505073ffffffffffffffffffffffffffffffffffffffff82165f908152602081905260408120805491610358836104b7565b91905055505050565b73ffffffffffffffffffffffffffffffffffffffff8083165f9081526020818152604080832093851683526001909301905220545b92915050565b803573ffffffffffffffffffffffffffffffffffffffff811681146103bf575f5ffd5b919050565b5f5f604083850312156103d5575f5ffd5b6103de8361039c565b946020939093013593505050565b5f5f5f606084860312156103fe575f5ffd5b6104078461039c565b92506104156020850161039c565b929592945050506040919091013590565b5f60208284031215610436575f5ffd5b61043f8261039c565b9392505050565b5f5f60408385031215610457575f5ffd5b6104608361039c565b915061046e6020840161039c565b90509250929050565b7f4e487b71000000000000000000000000000000000000000000000000000000005f52601160045260245ffd5b8181038181111561039657610396610477565b5f7fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff82036104e7576104e7610477565b5060010190565b808201808211156103965761039661047756fea164736f6c63782d302e382e33302d646576656c6f702e323032352e31302e33302b636f6d6d69742e37333731326130312e6d6f640035",
+  "deployedBytecode": "0x608060405234801561000f575f5ffd5b506004361061006f575f3560e01c806370a082311161004d57806370a082311461010b578063a9059cbb14610155578063dd62ed3e14610168575f5ffd5b8063095ea7b31461007357806323b872dd146100bb57806340c10f19146100ce575b5f5ffd5b6100b96100813660046103c4565b335f9081526020818152604080832073ffffffffffffffffffffffffffffffffffffffff959095168352600190940190529190912055565b005b6100b96100c93660046103ec565b61017b565b6100b96100dc3660046103c4565b73ffffffffffffffffffffffffffffffffffffffff9091165f9081526020819052604090206001815560020155565b610143610119366004610426565b73ffffffffffffffffffffffffffffffffffffffff165f9081526020819052604090206002015490565b60405190815260200160405180910390f35b6100b96101633660046103c4565b6102a7565b610143610176366004610446565b610361565b73ffffffffffffffffffffffffffffffffffffffff83165f90815260208181526040808320338452600101909152812080548392906101bb9084906104a4565b909155505073ffffffffffffffffffffffffffffffffffffffff83165f90815260208190526040812060020180548392906101f79084906104a4565b909155505073ffffffffffffffffffffffffffffffffffffffff83165f90815260208190526040812080549161022c836104b7565b909155505073ffffffffffffffffffffffffffffffffffffffff82165f90815260208190526040812060020180548392906102689084906104ee565b909155505073ffffffffffffffffffffffffffffffffffffffff82165f90815260208190526040812080549161029d836104b7565b9190505550505050565b335f90815260208190526040812060020180548392906102c89084906104a4565b9091555050335f9081526020819052604081208054916102e7836104b7565b909155505073ffffffffffffffffffffffffffffffffffffffff82165f90815260208190526040812060020180548392906103239084906104ee565b909155505073ffffffffffffffffffffffffffffffffffffffff82165f908152602081905260408120805491610358836104b7565b91905055505050565b73ffffffffffffffffffffffffffffffffffffffff8083165f9081526020818152604080832093851683526001909301905220545b92915050565b803573ffffffffffffffffffffffffffffffffffffffff811681146103bf575f5ffd5b919050565b5f5f604083850312156103d5575f5ffd5b6103de8361039c565b946020939093013593505050565b5f5f5f606084860312156103fe575f5ffd5b6104078461039c565b92506104156020850161039c565b929592945050506040919091013590565b5f60208284031215610436575f5ffd5b61043f8261039c565b9392505050565b5f5f60408385031215610457575f5ffd5b6104608361039c565b915061046e6020840161039c565b90509250929050565b7f4e487b71000000000000000000000000000000000000000000000000000000005f52601160045260245ffd5b8181038181111561039657610396610477565b5f7fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff82036104e7576104e7610477565b5060010190565b808201808211156103965761039661047756fea164736f6c63782d302e382e33302d646576656c6f702e323032352e31302e33302b636f6d6d69742e37333731326130312e6d6f640035",
+  "devdoc": {
+    "methods": {}
+  },
+  "userdoc": {
+    "methods": {}
+  }
+}

crates/contracts/build.rs

@@ -764,7 +764,9 @@ fn main() {
                     SEPOLIA => "0x0625aFB445C3B6B7B929342a04A22599fd5dBB59",
                     ARBITRUM_ONE => "0xcb8b5CD20BdCaea9a010aC1F8d835824F5C87A04",
                     BASE => "0xc694a91e6b071bF030A18BD3053A7fE09B6DaE69",
-                ])),
+                ]))
+                // Test Contract with non-standard ERC20 balance storage
+                .add_contract(Contract::new("NonStandardERC20Balances")),
         )
         .add_submodule(
             Submodule::new("support") // support module

crates/contracts/solidity/tests/NonStandardERC20Balances.sol

@@ -0,0 +1,47 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.17;
+
+contract NonStandardERC20Balances {
+
+    struct UserData {
+        // hypothetical other field
+        uint256 epoch;
+        mapping(address => uint256) approvals;
+        uint256 balance;
+    }
+
+    // In this token, the user's balance is stored at an offset position in the mapping -> cannot be detected by standard methods
+    mapping(address => UserData) users;
+
+    function mint(address user, uint256 amount) external {
+        users[user].epoch = 1;
+        users[user].balance = amount;
+    }
+
+    function balanceOf(address user) external view returns (uint256) {
+        return users[user].balance;
+    }
+
+    function allowance(address user, address spender) external view returns (uint256) {
+        return users[user].approvals[spender];
+    }
+
+    function transfer(address to, uint256 amount) external {
+        users[msg.sender].balance -= amount;
+        users[msg.sender].epoch++;
+        users[to].balance += amount;
+        users[to].epoch++;
+    }
+
+    function transferFrom(address from, address to, uint256 amount) external {
+        users[from].approvals[msg.sender] -= amount;
+        users[from].balance -= amount;
+        users[from].epoch++;
+        users[to].balance += amount;
+        users[to].epoch++;
+    }
+
+    function approve(address spender, uint256 amount) external {
+        users[msg.sender].approvals[spender] = amount;
+    }
+}

crates/e2e/tests/e2e/quote_verification.rs

@@ -8,7 +8,10 @@ use {
     ethcontract::H160,
     ethrpc::{
         Web3,
-        alloy::{CallBuilderExt, conversions::IntoAlloy},
+        alloy::{
+            CallBuilderExt,
+            conversions::{IntoAlloy, IntoLegacy},
+        },
     },
     model::{
         interaction::InteractionData,
@@ -25,7 +28,7 @@ use {
                 PriceQuery,
                 TradeVerifier,
                 TradeVerifying,
-                balance_overrides::BalanceOverrides,
+                balance_overrides::{BalanceOverrides, BalanceOverriding, Strategy},
             },
         },
         trade_finding::{Interaction, LegacyTrade, QuoteExecution, TradeKind},
@@ -69,6 +72,12 @@ async fn local_node_verified_quote_with_simulated_balance() {
     run_test(verified_quote_with_simulated_balance).await;
 }
 
+#[tokio::test]
+#[ignore]
+async fn local_node_trace_based_balance_detection() {
+    run_test(trace_based_balance_detection).await;
+}
+
 #[tokio::test]
 #[ignore]
 async fn forked_node_mainnet_usdt_quote() {
@@ -549,3 +558,141 @@ async fn usdt_quote_verification(web3: Web3) {
         .unwrap();
     assert!(quote.verified);
 }
+
+/// Tests that balance override detection works with debug_traceCall.
+/// This test verifies the trace-based detection strategy that's similar to
+/// Foundry's `deal`.
+async fn trace_based_balance_detection(web3: Web3) {
+    use shared::price_estimation::trade_verifier::balance_overrides::detector::Detector;
+
+    tracing::info!("Setting up chain state.");
+    let mut onchain = OnchainComponents::deploy(web3.clone()).await;
+
+    let [solver] = onchain.make_solvers(eth(10)).await;
+    let [trader] = onchain.make_accounts(eth(1)).await;
+
+    // Deploy the NonStandardERC20Balances token - this has balances stored at an
+    // offset within a struct mapping, making it undetectable by standard slot
+    // calculation methods
+    let non_standard_token =
+        contracts::alloy::test::NonStandardERC20Balances::Instance::deploy(web3.alloy.clone())
+            .await
+            .unwrap();
+
+    // Mint some tokens to the trader (so the contract has non-zero state)
+    non_standard_token
+        .mint(trader.address(), eth(100))
+        .from(solver.address())
+        .send_and_watch()
+        .await
+        .unwrap();
+
+    let detector = Detector::new(web3.clone(), 60);
+
+    let test_account = address!("0000000000000000000000000000000000000042");
+    let test_balance = U256::from(123_456_789_u64);
+
+    // Test with WETH (standard ERC20 with mapping at slot 3)
+    let weth = *onchain.contracts().weth.address();
+    tracing::info!("Testing WETH balance detection...");
+    let weth_strategy = detector
+        .detect(weth.into_legacy(), test_account.into_legacy())
+        .await;
+    assert!(
+        matches!(weth_strategy, Ok(Strategy::DirectSlot { .. })),
+        "Should detect WETH balance slot via trace"
+    );
+    tracing::info!("✓ WETH strategy detected: {:?}", weth_strategy);
+
+    // Test with NonStandardERC20Balances - this is the key test case
+    // The balance is at offset 2 within the UserData struct (epoch=0, approvals
+    // mapping=1, balance=2)
+    tracing::info!("Testing NonStandardERC20Balances detection...");
+    let non_standard_strategy = detector
+        .detect(
+            non_standard_token.address().into_legacy(),
+            test_account.into_legacy(),
+        )
+        .await;
+    assert!(
+        matches!(non_standard_strategy, Ok(Strategy::DirectSlot { .. })),
+        "Should detect non-standard token balance slot via trace-based detection"
+    );
+    tracing::info!(
+        "✓ NonStandardERC20Balances strategy detected: {:?}",
+        non_standard_strategy
+    );
+
+    // Verify that the detected strategies actually work by testing balance
+    // overrides
+    use {contracts::alloy::ERC20, ethrpc::alloy::conversions::IntoAlloy};
+
+    async fn test_balance_override(
+        web3: &Web3,
+        token: Address,
+        strategy: Strategy,
+        test_account: Address,
+        test_balance: U256,
+    ) {
+        use {
+            shared::price_estimation::trade_verifier::balance_overrides::BalanceOverrideRequest,
+            std::collections::HashMap,
+        };
+
+        let balance_overrides = BalanceOverrides {
+            hardcoded: HashMap::from([(token.into_legacy(), strategy)]),
+            detector: None,
+        };
+
+        let override_result = balance_overrides
+            .state_override(BalanceOverrideRequest {
+                token: token.into_legacy(),
+                holder: test_account.into_legacy(),
+                amount: test_balance.into_legacy(),
+            })
+            .await;
+
+        assert!(override_result.is_some(), "Should produce state override");
+        let (override_token, state_override) = override_result.unwrap();
+
+        // Call balanceOf with the state override to verify it works
+        let token_contract = ERC20::Instance::new(token, web3.alloy.clone());
+        let balance = token_contract
+            .balanceOf(test_account)
+            .state(HashMap::from([(override_token, state_override)]).into_alloy())
+            .call()
+            .await
+            .unwrap();
+
+        assert_eq!(
+            balance, test_balance,
+            "Balance override should work for token {:?}",
+            token
+        );
+
+        tracing::info!(
+            "✓ Balance override verified for token {:?}: set {} and got {}",
+            token,
+            test_balance,
+            balance
+        );
+    }
+
+    // Test each detected strategy
+    test_balance_override(
+        &web3,
+        weth,
+        weth_strategy.unwrap(),
+        test_account,
+        test_balance,
+    )
+    .await;
+    test_balance_override(
+        &web3,
+        *non_standard_token.address(),
+        non_standard_strategy.unwrap(),
+        test_account,
+        test_balance,
+    )
+    .await;
+}