tempo_transaction_pool/

paused.rs

//! Pool for transactions whose fee token is temporarily paused.
//!
//! When a TIP20 fee token emits `PauseStateUpdate(isPaused=true)`, transactions
//! using that fee token are moved here instead of being evicted entirely.
//! When the token is unpaused, transactions are moved back to the main pool
//! and re-validated.

use crate::{RevokedKeys, SpendingLimitUpdates, transaction::TempoPooledTransaction};
use alloy_primitives::{
    Address, TxHash,
    map::{AddressMap, B256Set},
};
use reth_transaction_pool::ValidPoolTransaction;
use std::{sync::Arc, time::Instant};

/// Duration after which paused transactions are expired and removed.
/// If a token isn't unpaused within this time, we clear all pending transactions.
pub const PAUSED_TX_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30 * 60); // 30 minutes

/// Global cap on the total number of paused transactions across all tokens.
///
/// Without this cap, an attacker could repeatedly fill the main pool, trigger a pause, and shift
/// transactions into the paused pool indefinitely. This bounds memory usage regardless of how many
/// tokens are paused or how frequently pause events occur.
pub const PAUSED_POOL_GLOBAL_CAP: usize = 10_000;

/// Entry in the paused pool.
#[derive(Debug, Clone)]
pub struct PausedEntry {
    /// The valid pool transaction that was paused (Arc to avoid expensive clones).
    pub tx: Arc<ValidPoolTransaction<TempoPooledTransaction>>,
    /// The `valid_before` timestamp, if any (for expiry tracking).
    pub valid_before: Option<u64>,
}

/// Metadata for a paused fee token.
#[derive(Debug, Clone)]
struct PausedTokenMeta {
    /// When this token was paused.
    paused_at: Instant,
    /// Transactions waiting for this token to be unpaused.
    entries: Vec<PausedEntry>,
}

/// Pool for transactions whose fee token is temporarily paused.
///
/// Transactions are indexed by fee token address for efficient batch operations.
/// Since all transactions for a token are paused/unpaused together, we track
/// the pause timestamp at the token level rather than per-transaction.
#[derive(Debug, Default)]
pub struct PausedFeeTokenPool {
    /// Fee token -> metadata including pause time and entries
    by_token: AddressMap<PausedTokenMeta>,
}

impl PausedFeeTokenPool {
    /// Creates a new empty paused pool.
    pub fn new() -> Self {
        Self::default()
    }

    /// Returns the total number of paused transactions across all tokens.
    pub fn len(&self) -> usize {
        self.by_token.values().map(|m| m.entries.len()).sum()
    }

    /// Returns true if there are no paused transactions.
    pub fn is_empty(&self) -> bool {
        self.by_token.is_empty()
    }

    /// Inserts transactions for a fee token into the paused pool.
    ///
    /// Takes the full batch at once since all transactions for a token
    /// are paused together. The pause timestamp is recorded at insertion time.
    ///
    /// Enforces [`PAUSED_POOL_GLOBAL_CAP`]: if adding the batch would exceed the cap,
    /// the oldest-paused tokens are evicted first to make room. If the batch itself
    /// exceeds the cap, it is truncated.
    ///
    /// Returns the number of existing entries that were evicted to make room.
    pub fn insert_batch(&mut self, fee_token: Address, entries: Vec<PausedEntry>) -> usize {
        if entries.is_empty() {
            return 0;
        }

        let current = self.len();
        let incoming = entries.len();
        let available = PAUSED_POOL_GLOBAL_CAP.saturating_sub(current);
        let mut evicted = 0;

        if incoming > available {
            let need = incoming - available;
            evicted = self.evict_oldest(need);
        }

        let remaining_capacity = PAUSED_POOL_GLOBAL_CAP.saturating_sub(self.len());
        let to_insert = if incoming > remaining_capacity {
            entries.into_iter().take(remaining_capacity).collect()
        } else {
            entries
        };

        self.by_token
            .entry(fee_token)
            .or_insert_with(|| PausedTokenMeta {
                paused_at: Instant::now(),
                entries: Vec::new(),
            })
            .entries
            .extend(to_insert);

        evicted
    }

    /// Evicts at least `need` entries from the oldest-paused tokens.
    ///
    /// Returns the total number of entries evicted.
    fn evict_oldest(&mut self, need: usize) -> usize {
        let mut tokens_by_age: Vec<_> = self
            .by_token
            .iter()
            .map(|(addr, meta)| (*addr, meta.paused_at))
            .collect();
        tokens_by_age.sort_unstable_by_key(|(_, paused_at)| *paused_at);

        let mut evicted = 0;
        for (token, _) in tokens_by_age {
            if evicted >= need {
                break;
            }
            if let Some(meta) = self.by_token.remove(&token) {
                evicted += meta.entries.len();
            }
        }
        evicted
    }

    /// Drains all transactions for a given fee token.
    ///
    /// Returns the list of paused entries for that token.
    pub fn drain_token(&mut self, fee_token: &Address) -> Vec<PausedEntry> {
        self.by_token
            .remove(fee_token)
            .map(|m| m.entries)
            .unwrap_or_default()
    }

    /// Returns the number of transactions paused for a given fee token.
    pub fn count_for_token(&self, fee_token: &Address) -> usize {
        self.by_token.get(fee_token).map_or(0, |m| m.entries.len())
    }

    /// Returns true if a transaction with the given hash is in the paused pool.
    pub fn contains(&self, tx_hash: &TxHash) -> bool {
        self.by_token
            .values()
            .any(|m| m.entries.iter().any(|e| e.tx.hash() == tx_hash))
    }

    /// Evicts expired transactions based on `valid_before` timestamp.
    ///
    /// Returns the number of transactions removed.
    pub fn evict_expired(&mut self, tip_timestamp: u64) -> usize {
        let mut count = 0;
        for meta in self.by_token.values_mut() {
            let before = meta.entries.len();
            meta.entries
                .retain(|e| e.valid_before.is_none_or(|vb| vb > tip_timestamp));
            count += before - meta.entries.len();
        }
        // Clean up empty token entries
        self.by_token.retain(|_, m| !m.entries.is_empty());
        count
    }

    /// Evicts all transactions for tokens that have been paused for too long (timeout).
    ///
    /// Since all transactions for a token are paused together, we evict the entire
    /// token's transactions when the token-level timeout expires.
    ///
    /// Returns the number of transactions removed.
    pub fn evict_timed_out(&mut self) -> usize {
        let now = Instant::now();
        let mut count = 0;
        self.by_token.retain(|_, meta| {
            if now.duration_since(meta.paused_at) >= PAUSED_TX_TIMEOUT {
                count += meta.entries.len();
                false
            } else {
                true
            }
        });
        count
    }

    /// Removes transactions matching invalidation criteria from the paused pool.
    ///
    /// This handles hard keychain invalidations in a single pass. Keychain
    /// `IAccountKeychain::AccessKeySpend` events are intentionally omitted: they only prove
    /// partial spending-limit consumption, and paused transactions are fully revalidated when
    /// their fee token unpauses.
    /// Uses account-keyed indexes for O(1) account lookup per transaction.
    /// Returns the number of transactions removed.
    pub fn evict_invalidated(
        &mut self,
        revoked_keys: &RevokedKeys,
        key_authorization_target_changes: &RevokedKeys,
        spending_limit_updates: &SpendingLimitUpdates,
        key_authorization_witness_burns: &AddressMap<B256Set>,
    ) -> usize {
        if revoked_keys.is_empty()
            && key_authorization_target_changes.is_empty()
            && spending_limit_updates.is_empty()
            && key_authorization_witness_burns.is_empty()
        {
            return 0;
        }

        let mut count = 0;
        let has_keychain_subject_updates =
            !revoked_keys.is_empty() || !spending_limit_updates.is_empty();
        let has_key_authorization_target_updates = !key_authorization_target_changes.is_empty();
        for meta in self.by_token.values_mut() {
            let before = meta.entries.len();
            meta.entries.retain(|entry| {
                let key_authorization_subject = (!revoked_keys.is_empty())
                    .then(|| entry.tx.transaction.key_authorization_signer_subject())
                    .flatten();
                let key_authorization_target = has_key_authorization_target_updates
                    .then(|| entry.tx.transaction.key_authorization_target_subject())
                    .flatten();

                let keychain_subject = has_keychain_subject_updates
                    .then(|| entry.tx.transaction.keychain_subject())
                    .flatten();
                let Some(subject) = keychain_subject else {
                    let Some(witness_subject) =
                        entry.tx.transaction.key_authorization_witness_subject()
                    else {
                        return !key_authorization_subject
                            .as_ref()
                            .is_some_and(|subject| subject.matches_revoked(revoked_keys))
                            && !key_authorization_target.as_ref().is_some_and(|subject| {
                                subject.matches_key_update(key_authorization_target_changes)
                            });
                    };

                    return !key_authorization_subject
                        .as_ref()
                        .is_some_and(|subject| subject.matches_revoked(revoked_keys))
                        && !key_authorization_target.as_ref().is_some_and(|subject| {
                            subject.matches_key_update(key_authorization_target_changes)
                        })
                        && !key_authorization_witness_burns
                            .get(&witness_subject.account)
                            .is_some_and(|witnesses| witnesses.contains(&witness_subject.witness));
                };

                let matches_limit_update =
                    subject.matches_spending_limit_update(spending_limit_updates);
                let sender_paid = matches_limit_update && entry.tx.transaction.is_sender_paid_fee();

                if subject.matches_revoked(revoked_keys)
                    || key_authorization_subject
                        .as_ref()
                        .is_some_and(|subject| subject.matches_revoked(revoked_keys))
                    || key_authorization_target.as_ref().is_some_and(|subject| {
                        subject.matches_key_update(key_authorization_target_changes)
                    })
                    || (sender_paid && matches_limit_update)
                {
                    return false;
                }

                let Some(witness_subject) =
                    entry.tx.transaction.key_authorization_witness_subject()
                else {
                    return true;
                };

                !key_authorization_witness_burns
                    .get(&witness_subject.account)
                    .is_some_and(|witnesses| witnesses.contains(&witness_subject.witness))
            });
            count += before - meta.entries.len();
        }
        // Clean up empty token entries
        self.by_token.retain(|_, m| !m.entries.is_empty());
        count
    }

    /// Returns an iterator over all paused entries across all tokens.
    pub fn all_entries(&self) -> impl Iterator<Item = &PausedEntry> {
        self.by_token.values().flat_map(|m| &m.entries)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_utils::{TxBuilder, wrap_valid_tx};
    use alloy_primitives::B256;
    use alloy_signer::SignerSync;
    use alloy_signer_local::PrivateKeySigner;
    use reth_primitives_traits::Recovered;
    use reth_transaction_pool::TransactionOrigin;
    use tempo_primitives::{
        SignatureType, TempoTxEnvelope,
        transaction::{KeyAuthorization, PrimitiveSignature, tt_signed::AASigned},
    };

    fn create_valid_tx(sender: Address) -> Arc<ValidPoolTransaction<TempoPooledTransaction>> {
        let pooled = TxBuilder::aa(sender).build();
        Arc::new(wrap_valid_tx(pooled, TransactionOrigin::External))
    }

    fn create_valid_keychain_tx(
        sender: Address,
        fee_token: Address,
        sponsored: bool,
    ) -> Arc<ValidPoolTransaction<TempoPooledTransaction>> {
        let access_key_signer = PrivateKeySigner::random();
        let pooled = TxBuilder::aa(sender)
            .fee_token(fee_token)
            .build_keychain(sender, &access_key_signer);

        let pooled = if sponsored {
            let sponsor = PrivateKeySigner::random();
            let aa = pooled
                .inner()
                .as_aa()
                .expect("builder should produce AA tx");
            let mut tx = aa.tx().clone();
            tx.fee_payer_signature = Some(alloy_primitives::Signature::new(
                alloy_primitives::U256::ZERO,
                alloy_primitives::U256::ZERO,
                false,
            ));
            let fee_payer_hash = tx.fee_payer_signature_hash(sender);
            tx.fee_payer_signature = Some(
                sponsor
                    .sign_hash_sync(&fee_payer_hash)
                    .expect("sponsor signing should succeed"),
            );

            let aa_signed = AASigned::new_unhashed(tx, aa.signature().clone());
            let envelope: TempoTxEnvelope = aa_signed.into();
            TempoPooledTransaction::new(Recovered::new_unchecked(envelope, sender))
        } else {
            pooled
        };

        Arc::new(wrap_valid_tx(pooled, TransactionOrigin::External))
    }

    #[test]
    fn test_insert_and_drain() {
        let mut pool = PausedFeeTokenPool::new();
        let fee_token = Address::random();

        let entries: Vec<_> = (0..3)
            .map(|_| PausedEntry {
                tx: create_valid_tx(Address::random()),
                valid_before: None,
            })
            .collect();

        assert!(pool.is_empty());
        pool.insert_batch(fee_token, entries);

        assert_eq!(pool.len(), 3);
        assert_eq!(pool.count_for_token(&fee_token), 3);

        let drained = pool.drain_token(&fee_token);
        assert_eq!(drained.len(), 3);
        assert!(pool.is_empty());
    }

    #[test]
    fn test_evict_expired() {
        let mut pool = PausedFeeTokenPool::new();
        let fee_token = Address::random();

        let entries = vec![
            PausedEntry {
                tx: create_valid_tx(Address::random()),
                valid_before: Some(100), // Will expire
            },
            PausedEntry {
                tx: create_valid_tx(Address::random()),
                valid_before: Some(200), // Won't expire
            },
            PausedEntry {
                tx: create_valid_tx(Address::random()),
                valid_before: None, // No expiry
            },
        ];

        pool.insert_batch(fee_token, entries);
        assert_eq!(pool.len(), 3);

        let evicted = pool.evict_expired(150);
        assert_eq!(evicted, 1);
        assert_eq!(pool.len(), 2);
    }

    #[test]
    fn test_global_cap_evicts_oldest() {
        let mut pool = PausedFeeTokenPool::new();

        let token_a = Address::random();
        let token_b = Address::random();

        let make_entries = |n: usize| -> Vec<PausedEntry> {
            (0..n)
                .map(|_| PausedEntry {
                    tx: create_valid_tx(Address::random()),
                    valid_before: None,
                })
                .collect()
        };

        // Fill to cap
        let evicted = pool.insert_batch(token_a, make_entries(PAUSED_POOL_GLOBAL_CAP));
        assert_eq!(evicted, 0);
        assert_eq!(pool.len(), PAUSED_POOL_GLOBAL_CAP);

        // Inserting more should evict token_a (oldest) to make room
        let evicted = pool.insert_batch(token_b, make_entries(100));
        assert!(evicted > 0);
        assert!(pool.len() <= PAUSED_POOL_GLOBAL_CAP);
        assert_eq!(pool.count_for_token(&token_b), 100);
    }

    #[test]
    fn test_global_cap_truncates_oversized_batch() {
        let mut pool = PausedFeeTokenPool::new();
        let token = Address::random();

        let entries: Vec<_> = (0..PAUSED_POOL_GLOBAL_CAP + 500)
            .map(|_| PausedEntry {
                tx: create_valid_tx(Address::random()),
                valid_before: None,
            })
            .collect();

        let evicted = pool.insert_batch(token, entries);
        assert_eq!(evicted, 0);
        assert_eq!(pool.len(), PAUSED_POOL_GLOBAL_CAP);
    }

    #[test]
    fn test_evict_invalidated_with_spending_limit_updates() {
        let mut pool = PausedFeeTokenPool::new();
        let user_address = Address::random();
        let fee_token = Address::random();

        // Create a keychain-signed tx using the builder
        let access_key_signer = alloy_signer_local::PrivateKeySigner::random();
        let key_id = alloy_signer::Signer::address(&access_key_signer);
        let tx = TxBuilder::aa(user_address)
            .fee_token(fee_token)
            .build_keychain(user_address, &access_key_signer);
        let tx = Arc::new(wrap_valid_tx(
            tx,
            reth_transaction_pool::TransactionOrigin::External,
        ));

        // Also add a non-keychain tx that should NOT be evicted
        let other_tx = create_valid_tx(Address::random());

        pool.insert_batch(
            fee_token,
            vec![
                PausedEntry {
                    tx,
                    valid_before: None,
                },
                PausedEntry {
                    tx: other_tx,
                    valid_before: None,
                },
            ],
        );
        assert_eq!(pool.len(), 2);

        let mut updates = SpendingLimitUpdates::new();
        updates.insert(user_address, key_id, Some(fee_token));

        let evicted = pool.evict_invalidated(
            &RevokedKeys::new(),
            &RevokedKeys::new(),
            &updates,
            &AddressMap::default(),
        );

        assert_eq!(
            evicted, 1,
            "Should evict the keychain tx matching the spending limit update"
        );
        assert_eq!(pool.len(), 1, "Non-keychain tx should remain");
    }

    #[test]
    fn test_evict_invalidated_keeps_sponsored_keychain_for_spending_limit_updates() {
        let mut pool = PausedFeeTokenPool::new();
        let user_address = Address::random();
        let fee_token = Address::random();

        let sponsored_keychain_tx = create_valid_keychain_tx(user_address, fee_token, true);
        pool.insert_batch(
            fee_token,
            vec![PausedEntry {
                tx: sponsored_keychain_tx,
                valid_before: None,
            }],
        );

        let key_id = pool
            .all_entries()
            .next()
            .and_then(|entry| entry.tx.transaction.keychain_subject())
            .map(|subject| subject.key_id)
            .expect("sponsored keychain tx should have keychain subject");

        let mut updates = SpendingLimitUpdates::new();
        updates.insert(user_address, key_id, Some(fee_token));

        let evicted = pool.evict_invalidated(
            &RevokedKeys::new(),
            &RevokedKeys::new(),
            &updates,
            &AddressMap::default(),
        );

        assert_eq!(evicted, 0, "Sponsored keychain tx should not be evicted");
        assert_eq!(pool.len(), 1);
    }

    #[test]
    fn test_evict_invalidated_with_key_authorization_witness_burn() {
        let mut pool = PausedFeeTokenPool::new();
        let user_address = Address::random();
        let fee_token = Address::random();
        let burned_witness = B256::random();
        let other_witness = B256::random();

        let key_authorization = |witness| {
            KeyAuthorization::unrestricted(42431, SignatureType::Secp256k1, Address::random())
                .with_witness(witness)
                .into_signed(PrimitiveSignature::Secp256k1(
                    alloy_primitives::Signature::test_signature(),
                ))
        };

        let matching = Arc::new(wrap_valid_tx(
            TxBuilder::aa(user_address)
                .fee_token(fee_token)
                .key_authorization(key_authorization(burned_witness))
                .build(),
            TransactionOrigin::External,
        ));
        let untouched = Arc::new(wrap_valid_tx(
            TxBuilder::aa(user_address)
                .nonce(1)
                .fee_token(fee_token)
                .key_authorization(key_authorization(other_witness))
                .build(),
            TransactionOrigin::External,
        ));

        pool.insert_batch(
            fee_token,
            vec![
                PausedEntry {
                    tx: matching,
                    valid_before: None,
                },
                PausedEntry {
                    tx: untouched,
                    valid_before: None,
                },
            ],
        );

        let mut burned = AddressMap::default();
        burned
            .entry(user_address)
            .or_insert_with(B256Set::default)
            .insert(burned_witness);

        let evicted = pool.evict_invalidated(
            &RevokedKeys::new(),
            &RevokedKeys::new(),
            &SpendingLimitUpdates::new(),
            &burned,
        );

        assert_eq!(evicted, 1);
        assert_eq!(pool.len(), 1);
        assert_eq!(
            pool.all_entries()
                .next()
                .and_then(|entry| entry.tx.transaction.key_authorization_witness_subject())
                .map(|subject| subject.witness),
            Some(other_witness)
        );
    }

    #[test]
    fn test_evict_invalidated_with_revoked_key_authorization_signer() {
        let mut pool = PausedFeeTokenPool::new();
        let user_address = Address::random();
        let fee_token = Address::random();
        let admin_signer = PrivateKeySigner::random();
        let admin_key = alloy_signer::Signer::address(&admin_signer);
        let other_signer = PrivateKeySigner::random();

        let key_authorization = |signer: &PrivateKeySigner| {
            let authorization =
                KeyAuthorization::unrestricted(42431, SignatureType::Secp256k1, Address::random())
                    .with_account(user_address);
            let signature = signer
                .sign_hash_sync(&authorization.signature_hash())
                .expect("key authorization signing should succeed");
            authorization.into_signed(PrimitiveSignature::Secp256k1(signature))
        };

        let matching = Arc::new(wrap_valid_tx(
            TxBuilder::aa(user_address)
                .fee_token(fee_token)
                .key_authorization(key_authorization(&admin_signer))
                .build(),
            TransactionOrigin::External,
        ));
        let untouched = Arc::new(wrap_valid_tx(
            TxBuilder::aa(user_address)
                .nonce(1)
                .fee_token(fee_token)
                .key_authorization(key_authorization(&other_signer))
                .build(),
            TransactionOrigin::External,
        ));

        pool.insert_batch(
            fee_token,
            vec![
                PausedEntry {
                    tx: matching,
                    valid_before: None,
                },
                PausedEntry {
                    tx: untouched,
                    valid_before: None,
                },
            ],
        );

        let mut revoked_keys = RevokedKeys::new();
        revoked_keys.insert(user_address, admin_key);

        let evicted = pool.evict_invalidated(
            &revoked_keys,
            &RevokedKeys::new(),
            &SpendingLimitUpdates::new(),
            &AddressMap::default(),
        );

        assert_eq!(evicted, 1);
        assert_eq!(pool.len(), 1);
    }

    #[test]
    fn test_evict_invalidated_with_key_authorization_target_change() {
        let mut pool = PausedFeeTokenPool::new();
        let user_address = Address::random();
        let fee_token = Address::random();
        let signer = PrivateKeySigner::random();
        let target_key = Address::random();
        let other_key = Address::random();

        let key_authorization = |key_id| {
            let authorization =
                KeyAuthorization::unrestricted(42431, SignatureType::Secp256k1, key_id)
                    .with_account(user_address);
            let signature = signer
                .sign_hash_sync(&authorization.signature_hash())
                .expect("key authorization signing should succeed");
            authorization.into_signed(PrimitiveSignature::Secp256k1(signature))
        };

        let matching = Arc::new(wrap_valid_tx(
            TxBuilder::aa(user_address)
                .fee_token(fee_token)
                .key_authorization(key_authorization(target_key))
                .build(),
            TransactionOrigin::External,
        ));
        let untouched = Arc::new(wrap_valid_tx(
            TxBuilder::aa(user_address)
                .nonce(1)
                .fee_token(fee_token)
                .key_authorization(key_authorization(other_key))
                .build(),
            TransactionOrigin::External,
        ));

        pool.insert_batch(
            fee_token,
            vec![
                PausedEntry {
                    tx: matching,
                    valid_before: None,
                },
                PausedEntry {
                    tx: untouched,
                    valid_before: None,
                },
            ],
        );

        let mut target_changes = RevokedKeys::new();
        target_changes.insert(user_address, target_key);

        let evicted = pool.evict_invalidated(
            &RevokedKeys::new(),
            &target_changes,
            &SpendingLimitUpdates::new(),
            &AddressMap::default(),
        );

        assert_eq!(evicted, 1);
        assert_eq!(pool.len(), 1);
    }

    #[test]
    fn test_contains() {
        let mut pool = PausedFeeTokenPool::new();
        let fee_token = Address::random();

        let tx = create_valid_tx(Address::random());
        let tx_hash = *tx.hash();

        let entry = PausedEntry {
            tx,
            valid_before: None,
        };

        assert!(!pool.contains(&tx_hash));
        pool.insert_batch(fee_token, vec![entry]);
        assert!(pool.contains(&tx_hash));
    }
}