Skip to content

Support Aave v3 aToken balance overrides#4345

Merged
squadgazzz merged 17 commits intomainfrom
feat/aave-atoken-balance-override
Apr 22, 2026
Merged

Support Aave v3 aToken balance overrides#4345
squadgazzz merged 17 commits intomainfrom
feat/aave-atoken-balance-override

Conversation

@squadgazzz
Copy link
Copy Markdown
Contributor

@squadgazzz squadgazzz commented Apr 20, 2026

Description

Quote verification was silently failing for any Aave v3 aToken as the sell token (e.g. aEthWETH → WETH) because the balance-override mechanism assumes the value at the balance storage slot is what balanceOf returns. Aave v3 aTokens break both assumptions: balanceOf applies a rayMul(scaled, liquidityIndex) scaling, and storage is packed UserState { uint128 balance; uint128 additionalData }. The auto-detector's verify step therefore never matched, returned NotFound, and the trade verifier silently skipped the override, producing the production revert execution reverted: trader does not have enough sell token visible in barn logs for aEthWETH quotes today (reproduced on Tenderly: https://dashboard.tenderly.co/cow-protocol/barn/simulator/1e69fa91-9496-44d1-9aec-a0e34166f9df).

Changes

  • New Strategy::AaveV3AToken { target_contract, pool, underlying, map_slot } variant in configs::balance_overrides::Strategy and a corresponding async resolver on BalanceOverrides that fetches the current getReserveNormalizedIncome from the Aave v3 Pool, rayDivs the requested amount (round-half-up, bit-for-bit compatible with WadRayMath.rayDiv), and writes it into the low 128 bits of the packed _userState slot.
  • Auto-detector is now Aave-aware: the detector probes the token with UNDERLYING_ASSET_ADDRESS() + POOL() and then calls pool.getReserveData(underlying) and verifies the returned aTokenAddress equals the probed token — an identity check that only passes for tokens the pool itself has registered as an aToken for their underlying, preventing rogue contracts that merely implement the aToken selectors from being accepted. When the probe succeeds, mapping-style slots are also offered as AaveV3AToken candidates and verified with a 1-wei round-trip tolerance instead of strict equality. No hardcoded per-token list is needed. aEthWETH and every other Aave v3 aToken on every chain (mainnet, Arbitrum, Base, Gnosis, Polygon, BNB, Linea, Plasma, Ink, Avalanche, etc.) are picked up automatically the first time they're quoted.
  • Shared Aave math lives in a new balance-overrides::aave module so the production override builder and the detector probe/verify use identical arithmetic.
  • trade_verifier::prepare_state_overrides now emits tracing::warn! when the spardose balance override for the sell token can't be resolved — mirroring the existing warn on the buy-side path so future missing overrides surface in logs instead of only showing up as downstream reverts.

Aave v4 note

Aave v4's user-facing deposit receipt is TokenizationSpoke, an ERC-4626 vault built on OpenZeppelin's ERC20Upgradeable — not a scaled aToken. balanceOf is the plain OZ _balances[user] mapping (no rayMul scaling), and the contract doesn't expose UNDERLYING_ASSET_ADDRESS() or POOL(). Our probe therefore correctly rejects v4 Spokes and the detector falls through to the existing SolidityMapping / DirectSlot heuristics, which handle them as standard ERC-20s. No v4-specific work is needed here.
Source: https://github.com/aave/aave-v4/blob/main/src/spoke/TokenizationSpoke.sol

How to test

Unit tests:

  • balance-overrides::tests::a_token_balance_override_bug_reproduction — pins the arithmetic: writing the raw amount (what the old strategies do) makes balanceOf return rayMul(amount, index) (off by ~6e16 wei at aEthWETH's current index); writing rayDiv(amount, index) round-trips within 1 wei.
  • balance-overrides::aave::tests — mock-provider tests for the probe: accepts a valid aToken, rejects when either selector reverts, rejects when the claimed pool doesn't look like Aave v3, rejects when the pool registers a different aToken for the declared underlying.
  • balance-overrides::tests::aave_v3_a_token_override_scales_amount_and_writes_low_128 — mock-provider integration test for the override builder.

e2e local + forked node tests:

  • balance-overrides::detector::tests::detects_aave_v3_a_token_mainnet — asserts the detector returns AaveV3AToken { pool=0x87870bca…4fa4e2, underlying=WETH, map_slot=52 } for aEthWETH.
  • balance-overrides::tests::aave_v3_a_token_override_mainnet_roundtrip — applies the override via eth_call against real aEthWETH and asserts balanceOf(holder) ≈ amount.
  • e2e::quote_verification::forked_node_mainnet_aave_atoken_quote — same round-trip against an anvil mainnet fork, exercising the full BalanceOverrides::state_override path.

Follow-up items

  • Registry-based detection. The current probe costs 3 eth_calls per cold-cache token. An alternative is to enumerate all reserves once per chain from Aave's AaveProtocolDataProvider (or equivalent), cache the full list of aTokens, and do a pure HashMap lookup on each quote. That would make detection amortised O(1) and provide the strongest identity guarantee (taken from Aave's own registry), at the cost of introducing a per-chain constant for the DataProvider address, a refresh cadence for new markets, and handling multiple Aave v3 markets per chain (e.g. mainnet has Ethereum + Prime). Probably worth doing once aToken quote volume warrants the optimisation.
  • Cache the accrued liquidity index for ~1 block. Each aToken quote currently re-fetches getReserveNormalizedIncome. The index drifts slowly (fractions of a wei per block), so caching for 6–12 s would drop the per-quote cost to zero without meaningful accuracy loss.
  • Aave v2 support. Same scaling + packed-storage shape; different LendingPool interface. v2 is deprecated but still has markets; worth revisiting if any start to attract volume.

@squadgazzz squadgazzz marked this pull request as ready for review April 20, 2026 16:08
@squadgazzz squadgazzz requested a review from a team as a code owner April 20, 2026 16:08
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces a specialized balance override strategy for Aave v3 aTokens, addressing an issue where the standard storage-slot-based overrides failed due to Aave's internal scaling and packed storage layout. It adds a new aave module to handle the necessary math and RPC probes, integrates this into the Detector for automatic discovery, and updates the BalanceOverrides component to support asynchronous state overrides. A performance improvement was suggested to parallelize independent RPC calls during the aToken probing process to reduce latency.

Comment thread crates/balance-overrides/src/aave.rs Outdated
@squadgazzz
Copy link
Copy Markdown
Contributor Author

@claude review

@claude
Copy link
Copy Markdown

claude Bot commented Apr 20, 2026

Claude finished @squadgazzz's task in 2m 45s —— View job


Code Review

  • Read and analyze changed files
  • Review core Aave math module (aave.rs)
  • Review balance override changes (lib.rs)
  • Review detector changes (detector.rs)
  • Review config changes (balance_overrides.rs)
  • Review e2e and trade verifier changes
  • Post full review

This is a well-structured fix for a real production bug. The three-call identity check for auto-detection is the right approach — it prevents false positives without requiring a hardcoded registry. The math is correct and is tested with a round-trip reproduction. A few issues to address:

1. Performance: Parallelize independent eth_calls in probe_a_token (Important)

UNDERLYING_ASSET_ADDRESS() and POOL() are independent calls but currently sequential, adding an unnecessary RTT to every cold-cache token detection. Since this runs on every uncached quote, it matters:

// crates/balance-overrides/src/aave.rs:119-127 — current (sequential)
let underlying = call_address(web3, token, IAaveV3AToken::UNDERLYING_ASSET_ADDRESSCall {}.abi_encode()).await?;
let pool = call_address(web3, token, IAaveV3AToken::POOLCall {}.abi_encode()).await?;
let reserve = fetch_reserve_data(web3, pool, underlying).await?;

Should be:

let (underlying, pool) = tokio::join!(
    call_address(web3, token, IAaveV3AToken::UNDERLYING_ASSET_ADDRESSCall {}.abi_encode()),
    call_address(web3, token, IAaveV3AToken::POOLCall {}.abi_encode()),
);
let (underlying, pool) = (underlying?, pool?);
let reserve = fetch_reserve_data(web3, pool, underlying).await?;

Fix this →

2. Cache inefficiency: is_valid_for_all_holders not extended for AaveV3AToken (Important)

lib.rs:97-99 currently returns true only for DirectSlot:

fn is_valid_for_all_holders(&self) -> bool {
    matches!(self, Self::DirectSlot { .. })
}

AaveV3AToken stores pool, underlying, and map_slot — all token-level constants, not holder-specific. The current implementation caches the detected strategy as (token, Some(holder)), meaning each new holder for the same aToken triggers a fresh probe_a_token + getReserveNormalizedIncome round-trip. Since many traders quote aEthWETH → WETH, this will result in repeated detection for every unique from address.

SolidityMapping and SoladyMapping have the same pre-existing issue (their map_slot is token-level), but this PR is the right place to also fix it for AaveV3AToken. At minimum, AaveV3AToken should return true here since it's being introduced fresh.

Fix this →

3. Stale comment in the e2e test (Minor)

crates/e2e/tests/e2e/quote_verification.rs:87-89:

/// Quote verification for an Aave v3 aToken as the sell token. The
/// auto-detector cannot resolve aToken balance slots (scaled `balanceOf` +
/// packed `UserState`), so we rely on the `AaveV3AToken` hardcoded strategy.

This contradicts the PR's own goal — the auto-detector now can resolve aToken balance slots. The test exercises the hardcoded path specifically (with detector: None) but the reason is to isolate that code path, not because detection is impossible. The comment should reflect reality.

Fix this →

4. Silent ray_div overflow (Minor)

In aave.rs:188:

let scaled = ray_div(amount, index)?;

If amount * RAY overflows U256, ray_div returns None and build_override returns None silently. The fetch_normalized_income failure above it emits a warn!, but a math overflow here produces no diagnostic, making it harder to debug if it ever happens on a very large quote amount. A tracing::warn!(amount = %amount, index = %index, "ray_div overflow computing AaveV3AToken scaled balance") before the ? would help.

5. unreachable!() in StrategyExt::state_override (Design note)

lib.rs:84-87:

Self::AaveV3AToken { .. } => unreachable!(
    "AaveV3AToken strategy must be resolved asynchronously, not via StrategyExt"
),

This works and the comment is accurate — the dispatch in BalanceOverriding::state_override peels off AaveV3AToken before calling StrategyExt::state_override. But it's a runtime panic if the dispatch logic ever regresses. An alternative would be to split the enum: a SyncStrategy enum (for StrategyExt) and the full Strategy enum that includes AaveV3AToken. That said, this is a bigger refactor and the current approach is acceptable as long as the invariant stays local.


Summary

The overall approach is solid. The math is correct, the identity-check probe design is the right tradeoff between safety and simplicity (no per-chain registry needed), and the tests cover the key paths well. The two important items are the parallelization of probe_a_token calls (latency impact per quote) and cache key policy for AaveV3AToken (repeated detection overhead). The rest are minor cleanup items.

Comment thread crates/balance-overrides/src/aave.rs
Comment thread crates/balance-overrides/src/aave.rs Outdated
Comment thread crates/balance-overrides/src/aave.rs Outdated
Comment thread crates/balance-overrides/src/detector.rs Outdated
Comment thread crates/balance-overrides/src/detector.rs Outdated
Comment thread crates/balance-overrides/src/lib.rs Outdated
Comment thread crates/balance-overrides/src/lib.rs Outdated
Comment thread crates/balance-overrides/src/lib.rs Outdated
Copy link
Copy Markdown
Contributor

@jmg-duarte jmg-duarte left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

logic wise, everything looks sound to me

Comment thread crates/e2e/tests/e2e/quote_verification.rs Outdated
Comment thread crates/e2e/tests/e2e/quote_verification.rs Outdated
Comment thread crates/e2e/tests/e2e/quote_verification.rs Outdated
Comment thread crates/e2e/tests/e2e/quote_verification.rs Outdated
Comment thread crates/balance-overrides/src/detector.rs Outdated
Comment thread crates/balance-overrides/src/aave.rs Outdated
Comment thread crates/balance-overrides/src/aave.rs
Comment thread crates/balance-overrides/src/aave.rs Outdated
Comment thread crates/balance-overrides/src/aave.rs Outdated
Comment thread crates/balance-overrides/src/aave.rs Outdated
@squadgazzz squadgazzz added the hotfix Labels PRs that should be applied into production right away label Apr 20, 2026
Comment thread crates/balance-overrides/src/aave.rs Outdated
Comment thread crates/balance-overrides/src/aave.rs Outdated
Comment thread crates/balance-overrides/src/aave.rs Outdated
Comment thread crates/balance-overrides/src/aave.rs Outdated
Comment thread crates/balance-overrides/src/aave.rs Outdated
Comment thread crates/balance-overrides/src/detector.rs Outdated
Comment thread crates/balance-overrides/src/detector.rs Outdated
Comment on lines +56 to +58
/// holder and amount. `web3` is only consulted by strategies that need
/// on-chain reads at override time (currently `AaveV3AToken`); other
/// variants ignore it and complete synchronously.
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If web3 is only used by strategies could we not put the Web3 struct into the Strategy itself so that it has everything it needs?

Copy link
Copy Markdown
Contributor Author

@squadgazzz squadgazzz Apr 21, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Considered this and the tension is that Strategy lives in configs::balance_overrides as a serde-deserialized config type (loaded from TOML). A Web3 handle cannot sit inside a serde type, so carrying it on the enum would require splitting into a config variant and a runtime variant with an attachment pass on BalanceOverrides construction.

That's a fair chunk of plumbing for a one-off Option<&Web3> parameter on the trait. Leaving it as the extra arg for now.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see. Given that we don't use the hardcoded configurations in practice and that our detection capabilities should now automatically find anything the hardcoded list can cover what do you think about dropping the hardcoded configuration altogether?
That way we can drop the Deserialize on the strategy and have the cleaner implementation with the provider inside the Strategy.

Given that this change would be not super small it would probably make the most sense to merge this PR as is and follow up with another clean up PR.

Comment thread crates/balance-overrides/src/lib.rs Outdated
Comment thread crates/configs/src/balance_overrides.rs
@squadgazzz squadgazzz requested a review from MartinquaXD April 21, 2026 20:14
Copy link
Copy Markdown
Contributor

@MartinquaXD MartinquaXD left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM besides the 2 stray comments

Comment thread crates/balance-overrides/src/aave.rs Outdated
Comment thread crates/balance-overrides/src/lib.rs Outdated
@squadgazzz squadgazzz enabled auto-merge April 22, 2026 11:20
@squadgazzz squadgazzz added this pull request to the merge queue Apr 22, 2026
Merged via the queue into main with commit 8ce89a6 Apr 22, 2026
20 checks passed
@squadgazzz squadgazzz deleted the feat/aave-atoken-balance-override branch April 22, 2026 11:38
@github-actions github-actions Bot locked and limited conversation to collaborators Apr 22, 2026
Sign up for free to subscribe to this conversation on GitHub. Already have an account? Sign in.

Labels

hotfix Labels PRs that should be applied into production right away

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants