Move Smart Contract Access Control: How Capabilities Leak and Authorization Fails

// VULNERABLE: Anyone can call this and receive admin privileges
public fun get_admin_cap(): AdminCap {
    AdminCap {}
}

// VULNERABLE: TreasuryCap returned from a public initializer (unlimited minting)
public fun initialize(ctx: &mut TxContext): TreasuryCap<TOKEN> {
    let (treasury_cap, metadata) = coin::create_currency(
        TOKEN {}, 9, b"TKN", b"Token", b"", option::none(), ctx
    );
    transfer::public_freeze_object(metadata);
    treasury_cap  // BUG: Returns to any caller (they can mint unlimited tokens)
}

// SAFE: Send TreasuryCap only to the deployer via init()
fun init(ctx: &mut TxContext) {
    let (treasury_cap, metadata) = coin::create_currency(
        TOKEN {}, 9, b"TKN", b"Token", b"", option::none(), ctx
    );
    transfer::public_freeze_object(metadata);
    transfer::public_transfer(treasury_cap, tx_context::sender(ctx));
    // Only the deployer receives minting authority
}

// Context: a public getter that exposes &mut to the inner privileged
// Account is the same bug class as returning AdminCap by value, but at
// the reference level. Even if downstream privileged operations are not
// public, the &mut handle lets any caller mutate fields directly or
// pass it into any other public function that accepts &mut Account<C>.

module amm::vault {
    use sui::dynamic_object_field as dof;
    use sui::tx_context::TxContext;

    public struct Account<C> has store { collateral: u64, debt: u64 }
    public struct Vault<L, C> has key { id: UID }
    fun account_key(): vector<u8> { b"account" }

    // VULNERABLE: public getter hands out &mut Account<C> to anyone
    // holding a reference to the shared Vault. Turns every field of
    // Account<C> (collateral, debt, risk params) into a public API.
    public fun account_mut<L, C>(vault: &mut Vault<L, C>): &mut Account<C> {
        dof::borrow_mut(&mut vault.id, account_key())
    }

    // SAFE: restrict the &mut capability to the owning package, OR
    // don't expose the reference at all. Expose narrow, capability-
    // gated methods that mutate only what is intended.
    public(package) fun account_mut_safe<L, C>(vault: &mut Vault<L, C>): &mut Account<C> {
        dof::borrow_mut(&mut vault.id, account_key())
    }

    public(package) fun update_position<L, C>(vault: &mut Vault<L, C>, /* params */) {
        let acct = dof::borrow_mut<Account<C>>(&mut vault.id, account_key());
        // mutate only what is intended
    }
}

// Context: the witness pattern is meant to ensure only the defining
// module can mint a Cap<T>. `T: drop` is too loose a bound, and many
// types are droppable, so any attacker module can declare its own
// drop-only struct. The fix is `types::is_one_time_witness(&witness)`,
// which checks the witness was constructed in the same module that
// defined T (one-time witnesses are uppercase, drop-only, and produced
// exactly once at module init).

module enclave::caps {
    use sui::object;
    use sui::tx_context::TxContext;
    use sui::types;

    public struct Cap<T> has key { id: UID }

    // VULNERABLE: any droppable T satisfies the bound. An attacker
    // module can declare `struct W has drop {}` and mint unbounded
    // Cap<W> by calling new_cap(W {}, ctx).
    public fun new_cap<T: drop>(_: T, ctx: &mut TxContext): Cap<T> {
        Cap { id: object::new(ctx) }
    }

    // SAFE: enforce one-time witness so only the defining module's OTW
    // can pass the gate.
    public fun new_cap_safe<T: drop>(witness: T, ctx: &mut TxContext): Cap<T> {
        assert!(types::is_one_time_witness(&witness), ENotOneTimeWitness);
        Cap { id: object::new(ctx) }
    }
}

// The defining module supplies the OTW type (conventionally uppercase, drop-only).
module enclave::witness {
    public struct ENCLAVE has drop {}
}

// VULNERABLE: accepts &signer but never checks ownership of the order
public entry fun cancel_order<CoinType>(
    user: &signer, order_id: u64
) acquires OrderStore {
    let order = get_order(order_id);
    // Missing: assert!(order.user_address == signer::address_of(user), ERR_PERMISSION_DENIED);
    deposit_funds<CoinType>(order_store, signer::address_of(user), order.base);
    // Attacker cancels someone else's order and receives the funds
}

// SAFE: Verify caller owns the order before releasing funds
public entry fun cancel_order<CoinType>(
    user: &signer, order_id: u64
) acquires OrderStore {
    let order = get_order(order_id);
    assert!(order.base_type == type_info::type_of<CoinType>(), ERR_WRONG_COIN_TYPE);
    assert!(order.user_address == signer::address_of(user), ERR_PERMISSION_DENIED);
    deposit_funds<CoinType>(order_store, signer::address_of(user), order.base);
}

// Context: the resource account is derived from the *caller's* address,
// but every other function in the module looks up resources at the
// fixed @movedrop-derived address. With no signer assert, any user-
// initiated airdrop permanently strands its funds at an unreachable
// address. The funds are not stolen, but they are inaccessible to the
// real protocol code.

// VULNERABLE: 'admin' signer accepted but never validated.
public entry fun initialize_airdrop(admin: &signer, seed: vector<u8>, total_amount: u64) {
    // Missing: assert!(signer::address_of(admin) == @movedrop, ENotAuthorized);
    let (resource_signer, _cap) = account::create_resource_account(admin, seed);
    let resource_addr = signer::address_of(&resource_signer);
    aptos_account::transfer(admin, resource_addr, total_amount);
}

// SAFE: assert the signer matches the expected protocol admin address.
public entry fun initialize_airdrop(admin: &signer, seed: vector<u8>, total_amount: u64) {
    assert!(signer::address_of(admin) == @movedrop, ENotAuthorized);
    let (resource_signer, _cap) = account::create_resource_account(admin, seed);
    let resource_addr = signer::address_of(&resource_signer);
    aptos_account::transfer(admin, resource_addr, total_amount);
}

// Context: the deadline check is unconditional, but the driver-identity
// assert lives inside the penalty-period branch. Outside the penalty
// window, the function returns a fulfill ticket without checking the
// caller, so anyone can fulfill the order and sidestep the auction.
// Cognitive bias: a reviewer skims, sees `assert!`, marks it safe;
// always check the *control-flow position* of every authorization assert.

// VULNERABLE: auth assert only fires inside the penalty-period branch.
public fun prepare_fulfill_winner(
    order_item: &Order,
    msg_driver: address,
    clock: &Clock,
    ctx: &TxContext,
): FulfillTicket {
    let clock_now_s = clock.timestamp_ms() / 1000;
    assert!(order_item.deadline() > clock_now_s, EDeadlineIsPassed);
    if (clock_now_s >= order_item.deadline() - (order_item.penalty_period() as u64)) {
        assert!(msg_driver == ctx.sender(), EInvalidDriver);
    };
    build_ticket(order_item, msg_driver)
}

// SAFE: auth assert runs unconditionally, before any branching.
public fun prepare_fulfill_winner(
    order_item: &Order,
    msg_driver: address,
    clock: &Clock,
    ctx: &TxContext,
): FulfillTicket {
    let clock_now_s = clock.timestamp_ms() / 1000;
    assert!(order_item.deadline() > clock_now_s, EDeadlineIsPassed);
    assert!(msg_driver == ctx.sender(), EInvalidDriver);
    build_ticket(order_item, msg_driver)
}

// VULNERABLE: Test module left as friend in production
module protocol::vault {
    friend protocol::vault_tests;  // <-- Should be removed before deploy!
    friend protocol::router;

    public(friend) fun mint_lp(amount: u64): Coin<LP> {
        // vault_tests can call this in production
        // Anyone who deploys a package that imports vault_tests
        // gains indirect access to LP minting
    }
}

// SAFE: Only production modules as friends, minimal surface
module protocol::vault {
    friend protocol::router;  // Only the router needs mint access

    public(friend) fun mint_lp(amount: u64): Coin<LP> {
        // Only router can call this
    }
}

// Context: `new` is correctly gated to the owning package, but `parse`
// reconstructs the same Ticket from arbitrary bytes at `public` visibility.
// Anyone outside the package can fabricate a Ticket, bypassing the `new`
// restriction entirely. Audit every sibling of a restricted constructor:
// parse, from_bytes, deserialize, unpack, decode.

module nft::ticket {
    public struct Ticket has key, store {
        id: UID,
        event_id: u64,
        seat: u64,
    }

    // Canonical constructor (correctly restricted).
    public(package) fun new(event_id: u64, seat: u64, ctx: &mut TxContext): Ticket {
        Ticket { id: object::new(ctx), event_id, seat }
    }

    // VULNERABLE: deserializer at wider visibility is a second construction
    // path that the `new` restriction doesn't cover.
    public fun parse(bytes: vector<u8>, ctx: &mut TxContext): Ticket {
        let (event_id, seat) = bcs::peel_u64_pair(bytes);
        Ticket { id: object::new(ctx), event_id, seat }
    }

    // SAFE: match `new`'s visibility AND route through `new` so there is
    // exactly one construction path.
    public(package) fun parse_safe(bytes: vector<u8>, ctx: &mut TxContext): Ticket {
        let (event_id, seat) = bcs::peel_u64_pair(bytes);
        new(event_id, seat, ctx)
    }
}

// Context: `insert` (enter the active set) correctly asserts the stake
// is at or above threshold. `update` (modify in place) was supposed to
// enforce the same invariant but skips the assert. Attacker stakes
// above threshold to enter, then calls update to drop below it,
// retaining active-set membership while dodging the stake requirement.

module committee::active_set {
    const EBelowThreshold: u64 = 0;

    struct Set has key {
        id: UID,
        threshold_stake: u64,
        // active: Table<address, u64>
    }

    // SAFE: "enter the set" path enforces the threshold.
    public fun insert(set: &mut Set, node: address, staked_amount: u64) {
        assert!(staked_amount >= set.threshold_stake, EBelowThreshold);
        // table::add(&mut set.active, node, staked_amount);
    }

    // VULNERABLE: "update" sibling skips the threshold check.
    public fun update(set: &mut Set, node: address, new_staked_amount: u64) {
        // Missing: assert!(new_staked_amount >= set.threshold_stake, EBelowThreshold);
        // table::borrow_mut(&mut set.active, node) = new_staked_amount;
    }
}

// VULNERABLE: AdminCap inside a shared Vault (anyone can access)
public struct Vault has key {
    id: UID,
    balance: Balance<USDC>,
    admin_cap: AdminCap,  // Accessible to anyone interacting with the shared vault
}

fun init(ctx: &mut TxContext) {
    let vault = Vault {
        id: object::new(ctx),
        balance: balance::zero(),
        admin_cap: AdminCap { id: object::new(ctx) },
    };
    transfer::share_object(vault);  // Now EVERYONE can reach admin_cap
}

// SAFE: AdminCap held by the admin's owned address, separate from shared state
public struct Vault has key {
    id: UID,
    balance: Balance<USDC>,
}

public struct AdminCap has key {
    id: UID,
    vault_id: ID,  // Links cap to vault without embedding it
}

fun init(ctx: &mut TxContext) {
    let vault = Vault { id: object::new(ctx), balance: balance::zero() };
    let vault_id = object::id(&vault);
    transfer::share_object(vault);
    // AdminCap goes ONLY to deployer, not inside the shared object
    transfer::transfer(AdminCap { id: object::new(ctx), vault_id }, ctx.sender());
}

// Context: the type system can prove a value has type
// &AuthorityCap<PACKAGE, Role>, but it cannot prove the cap is *still*
// authorized. Allowlists are mutated; roles are revoked; multisigs are
// rotated. If the runtime allowlist isn't consulted on every call, a
// revoked operator who still holds the cap object continues to act as
// a current one. The cap doesn't know it was revoked.

// VULNERABLE: typed cap checked at compile time, runtime allowlist never consulted.
public fun new<Role>(
    wrapper: &mut OracleAggregatorStorkIntegration,
    cap: &AuthorityCap<PACKAGE, Role>,
    feed_id: vector<u8>,
): FeedInfoObject {
    let wrapper_id = wrapper.borrow_mut_id(cap);
    let key = FeedInfoObjectKey(feed_id);
    // Missing: assert!(wrapper.active_assistants.contains(&object::id(cap)), EAssistantRevoked);
    assert!(!derived_object::exists(wrapper_id, key), EFeedInfoObjectAlreadyCreated);
    FeedInfoObject { id: derived_object::claim(wrapper_id, key), feed_id }
}

// SAFE: consult the runtime allowlist before honoring the typed cap.
public fun new<Role>(
    wrapper: &mut OracleAggregatorStorkIntegration,
    cap: &AuthorityCap<PACKAGE, Role>,
    feed_id: vector<u8>,
): FeedInfoObject {
    assert!(wrapper.active_assistants.contains(&object::id(cap)), EAssistantRevoked);
    let wrapper_id = wrapper.borrow_mut_id(cap);
    let key = FeedInfoObjectKey(feed_id);
    assert!(!derived_object::exists(wrapper_id, key), EFeedInfoObjectAlreadyCreated);
    FeedInfoObject { id: derived_object::claim(wrapper_id, key), feed_id }
}

// VULNERABLE: ConstructorRef leaked (holder can add resources or re-transfer)
public fun mint(creator: &signer): ConstructorRef {
    let constructor_ref = token::create_named_token(creator, ...);
    constructor_ref // Any caller can now:
    // 1. Generate TransferRef → steal the token back after selling it
    // 2. Generate DeleteRef → destroy the token
    // 3. Add arbitrary resources to the token object
}

// SAFE: Use ConstructorRef internally; never return or store it
public fun mint(creator: &signer) {
    let constructor_ref = token::create_named_token(creator, ...);
    let transfer_ref = object::generate_transfer_ref(&constructor_ref);
    // Store only what's needed; constructor_ref is dropped at function end
    // No external caller ever touches it
}

// Context: Aptos derives resource and object addresses deterministically
// from seeds, so the target address of `create_primary_store` is known
// in advance. `create_primary_store` aborts if a store already exists,
// so an attacker who pre-creates one at the predictable address bricks
// every subsequent reward addition. A companion finding does the same
// with `aptos_account::create_account` to brick lending-pair creation.

// VULNERABLE: non-idempotent create at deterministic address.
public entry fun new_reward_fa(manager: &signer, asset_metadata: Object<Metadata>) acquires Farming {
    assert!(manager::is_manager(manager), ERR_FARMING_UNAUTHORIZED);
    primary_fungible_store::create_primary_store(package::package_address(), asset_metadata);
    new_reward_for_farming_internal(fungible_asset::name(asset_metadata));
}

// SAFE: idempotent API tolerates a pre-existing store.
public entry fun new_reward_fa(manager: &signer, asset_metadata: Object<Metadata>) acquires Farming {
    assert!(manager::is_manager(manager), ERR_FARMING_UNAUTHORIZED);
    primary_fungible_store::ensure_primary_store_exists(package::package_address(), asset_metadata);
    new_reward_for_farming_internal(fungible_asset::name(asset_metadata));
}

// VULNERABLE: update_v2 has no authorization check
public entry fun update_v2(
    oracle: &mut Oracle,
    token_type: u8,
    price: u64,
    timestamp: u64,
    _ctx: &TxContext,
) {
    // Missing: assert!(is_whitelisted(oracle, ctx.sender()), ENotAuthorized);
    // Anyone can set any price for any token
    oracle.prices[token_type] = price;
    oracle.timestamps[token_type] = timestamp;
}

// FIXED: Enforce whitelist before allowing price update
public entry fun update_v2(
    oracle: &mut Oracle,
    token_type: u8,
    price: u64,
    timestamp: u64,
    ctx: &TxContext,
) {
    assert!(table::contains(&oracle.whitelist, ctx.sender()), ENotAuthorized);
    assert!(timestamp > oracle.timestamps[token_type], EStaleUpdate);
    oracle.prices[token_type] = price;
    oracle.timestamps[token_type] = timestamp;
}

// VULNERABLE: SetRevokerReq is owner-level. The gate must be check_owner, not check_operator
entry fun revoke_set_revoker<T>(state: &State<T>, req: SetRevokerReq, ctx: &TxContext) {
    check_version(state);
    check_operator(state, ctx); // BUG: copy-pasted from an operator-gated sibling
    let SetRevokerReq { id, .. } = req;
    id.delete();
}

// SAFE: the gate matches the role that issued the request
entry fun revoke_set_revoker<T>(state: &State<T>, req: SetRevokerReq, ctx: &TxContext) {
    check_version(state);
    check_owner(state, ctx);
    let SetRevokerReq { id, .. } = req;
    id.delete();
}

// VULNERABLE: single-step push (destination never proves control)
entry fun execute_transfer_ownership<T>(
    state: &mut State<T>,
    req: TransferOwnershipReq,
    ctx: &TxContext,
) {
    check_owner(state, ctx);
    let TransferOwnershipReq { id, upgrade_cap, new_owner } = req;
    transfer::public_transfer(upgrade_cap, new_owner); // typo here ⇒ UpgradeCap lost forever
    state.owner = new_owner;
    id.delete();
}

// SAFE: two-step pending/claim (recipient must prove control before takeover)
entry fun request_transfer_ownership<T>(
    state: &mut State<T>,
    new_owner: address,
    upgrade_cap: UpgradeCap,
    ctx: &TxContext,
) {
    check_owner(state, ctx);
    state.pending_owner = option::some(new_owner);
    state.pending_upgrade_cap = option::some(upgrade_cap);
}

entry fun accept_ownership<T>(state: &mut State<T>, ctx: &TxContext) {
    let new_owner = option::extract(&mut state.pending_owner);
    let upgrade_cap = option::extract(&mut state.pending_upgrade_cap);
    assert!(ctx.sender() == new_owner, ENotPendingOwner); // destination signs to claim
    transfer::public_transfer(upgrade_cap, new_owner);
    state.owner = new_owner;
}

// VULNERABLE: Anyone can mutate any NFT
public entry fun update_metadata(nft: &mut NFT, new_name: String) {
    nft.name = new_name;
}

// SAFE: Explicit owner validation
public entry fun update_metadata(
    caller: &signer,
    nft: &mut NFT,
    new_name: String
) {
    assert!(
        object::owner(nft) == signer::address_of(caller),
        E_NOT_OWNER
    );
    nft.name = new_name;
}

// Context: Walrus stake reward share derives from activation_epoch.
// `join` correctly checks node_id, withdrawing state, and withdraw_epoch,
// but omits activation_epoch. Merging a small old stake (epoch 10)
// with a fresh large stake (epoch 50) reports activation_epoch = 10 on
// the combined principal, so the entire new principal earns rewards as
// if it had been staked since epoch 10. Move's resource-safety
// guarantees hold throughout (no resource is duplicated, no balance
// disappears). The bug is purely at the protocol level: the equality
// assertion was incomplete.

// VULNERABLE: activation_epoch missing from equality check.
public fun join(sw: &mut StakedWal, other: StakedWal) {
    assert!(sw.node_id == other.node_id, EMetadataMismatch);
    assert!(sw.is_withdrawing() && other.is_withdrawing(), EMetadataMismatch);
    assert!(sw.withdraw_epoch() == other.withdraw_epoch(), EMetadataMismatch);
    // Missing: assert!(sw.activation_epoch == other.activation_epoch, EMetadataMismatch);
    let StakedWal { id, principal, .. } = other;
    sw.principal.join(principal);
    id.delete();
}

// SAFE: every field that affects downstream accounting is in the equality check.
public fun join(sw: &mut StakedWal, other: StakedWal) {
    assert!(sw.node_id == other.node_id, EMetadataMismatch);
    assert!(sw.is_withdrawing() && other.is_withdrawing(), EMetadataMismatch);
    assert!(sw.withdraw_epoch() == other.withdraw_epoch(), EMetadataMismatch);
    assert!(sw.activation_epoch == other.activation_epoch, EMetadataMismatch);
    let StakedWal { id, principal, .. } = other;
    sw.principal.join(principal);
    id.delete();
}

// VULNERABLE: Internal mint callable indirectly through a public function
fun internal_mint(treasury: &mut TreasuryCap<TOKEN>, amount: u64, recipient: address) {
    coin::mint_and_transfer(treasury, amount, recipient);
    // No check against supply cap
}

// Called from:
public entry fun claim_reward(caller: &signer, treasury: &mut TreasuryCap<TOKEN>) {
    // Missing: is this user actually entitled to a reward?
    // Missing: has this user already claimed?
    internal_mint(treasury, REWARD_AMOUNT, signer::address_of(caller));
}

// Context: `unlocked_tokens` is the *cumulative* amount unlocked by the
// schedule, not the amount remaining for this call. The per-call bound
// `amount <= unlocked_tokens` is therefore vacuous: tokens_transferred
// is incremented *after* the assert, so the user can call
// `claim(amount = unlocked_tokens)` on every transaction and drain the
// vesting position in O(1) calls. The missing line is the cumulative
// bound: `amount + self.tokens_transferred <= unlocked_tokens`.

// VULNERABLE: per-call bound correct, cumulative total never enforced.
public fun claim<T>(
    self: &mut Timelock<T>,
    amount: Option<u64>,
    clock: &Clock,
    ctx: &mut TxContext,
): SplitRequest<T> {
    let unlocked_tokens = self.unlocked_at(clock);
    let amount = amount.destroy_or!(unlocked_tokens);
    assert!(amount <= unlocked_tokens, ENotEnoughBalanceUnlocked);
    // Missing: assert!(amount + self.tokens_transferred <= unlocked_tokens, ...)
    self.tokens_transferred = self.tokens_transferred + amount;
    let coin = self.left_balance.split(amount).into_coin(ctx);
    let (token, request) = shared_token::from_coin(coin, ctx);
    token.share();
    request
}

// SAFE: bound the new amount against the cumulative total already claimed.
public fun claim<T>(
    self: &mut Timelock<T>,
    amount: Option<u64>,
    clock: &Clock,
    ctx: &mut TxContext,
): SplitRequest<T> {
    let unlocked_tokens = self.unlocked_at(clock);
    let amount = amount.destroy_or!(unlocked_tokens - self.tokens_transferred);
    assert!(amount + self.tokens_transferred <= unlocked_tokens, ENotEnoughBalanceUnlocked);
    self.tokens_transferred = self.tokens_transferred + amount;
    let coin = self.left_balance.split(amount).into_coin(ctx);
    let (token, request) = shared_token::from_coin(coin, ctx);
    token.share();
    request
}

// VULNERABLE: Receipt can be discarded if it accidentally has `drop`
struct FlashLoanReceipt has drop {  // BAD: has drop means it can be discarded
    amount: u64,
}

// SAFE: No abilities means the receipt MUST be consumed
struct FlashLoanReceipt {
    amount: u64,
}

public fun borrow(pool: &mut Pool, amount: u64): (Coin<USDC>, FlashLoanReceipt) {
    // ...
}

public fun repay(pool: &mut Pool, payment: Coin<USDC>, receipt: FlashLoanReceipt) {
    let FlashLoanReceipt { amount } = receipt;  // Destructuring consumes the receipt
    assert!(coin::value(&payment) >= amount, E_INSUFFICIENT_REPAYMENT);
    // ...
}

// Context: ThalaSwap V2 metastable pools normalize balances into a peg
// scale before invariant math. The pre-flashloan `balances` are upscaled
// once. The post-flashloan vector is then built from those already-
// upscaled values, and upscaled a second time. With upscale factor k:
//   single-upscale (correct): asserts k*(b + r) >= k*b, requires r >= 0.
//   double-upscale (bug):     asserts k*(k*b + r) >= k*b, passes at r=0.
// The hot potato is consumed correctly, satisfying Move's resource
// discipline; the economic invariant is violated by the math layered on
// top. Hot-potato consumption is necessary but not sufficient.

// VULNERABLE: post-flashloan balances upscaled a second time.
public fun pay_flashloan(assets: vector<FungibleAsset>, loan: Flashloan) acquires Pool, MetaStablePool {
    let (pool_obj, mut borrow_amounts, repay_amounts, fee_amounts) = consume_loan(loan, &assets);
    let mut balances = pool_balances(pool_obj);
    if (pool_is_metastable(pool_obj)) {
        borrow_amounts = upscale_metastable_amounts(pool_obj, borrow_amounts);
        balances = upscale_metastable_amounts(pool_obj, balances); // first upscale
    };

    let mut balances_after_flashloan = vector::empty<u64>();
    let mut i = 0;
    let len = vector::length(&balances);
    while (i < len) {
        let repay_sub_fees = *vector::borrow(&repay_amounts, i) - *vector::borrow(&fee_amounts, i);
        let after = *vector::borrow(&balances, i) + repay_sub_fees;
        vector::push_back(&mut balances_after_flashloan, after);
        i = i + 1;
    };

    if (pool_is_metastable(pool_obj)) {
        // BUG: vector is already in upscaled space; second upscale inflates the invariant.
        balances_after_flashloan = upscale_metastable_amounts(pool_obj, balances_after_flashloan);
    };
    assert_invariant_preserved(pool_obj, balances_after_flashloan);
    deposit_assets(pool_obj, assets);
}

// SAFE: upscale exactly once (invariant checked in the same scale on both sides).
public fun pay_flashloan(assets: vector<FungibleAsset>, loan: Flashloan) acquires Pool, MetaStablePool {
    let (pool_obj, mut borrow_amounts, repay_amounts, fee_amounts) = consume_loan(loan, &assets);
    let mut balances = pool_balances(pool_obj);
    if (pool_is_metastable(pool_obj)) {
        borrow_amounts = upscale_metastable_amounts(pool_obj, borrow_amounts);
        balances = upscale_metastable_amounts(pool_obj, balances);
    };

    let mut balances_after_flashloan = vector::empty<u64>();
    let mut i = 0;
    let len = vector::length(&balances);
    while (i < len) {
        let repay_sub_fees = *vector::borrow(&repay_amounts, i) - *vector::borrow(&fee_amounts, i);
        let after = *vector::borrow(&balances, i) + repay_sub_fees;
        vector::push_back(&mut balances_after_flashloan, after);
        i = i + 1;
    };
    // Already in upscaled space (no second upscaling).
    assert_invariant_preserved(pool_obj, balances_after_flashloan);
    deposit_assets(pool_obj, assets);
}