nonce_validation - [mainnet]

public(friend) fun initialize(aptos_framework: &signer) {
    initialize_nonce_table(aptos_framework);
}

public entry fun initialize_nonce_table(aptos_framework: &signer) {
    system_addresses::assert_aptos_framework(aptos_framework);
    if (!exists<NonceHistory>(@aptos_framework)) {
        let table = table::new();
        let nonce_history = NonceHistory {
            nonce_table: table,
            next_key: 0,
        };
        move_to<NonceHistory>(aptos_framework, nonce_history);
    };
}

fun empty_bucket(pre_allocate_slots: bool): Bucket {
    let bucket = Bucket {
        nonces_ordered_by_exp_time: big_ordered_map::new_with_reusable(),
        nonce_to_exp_time_map: big_ordered_map::new_with_reusable(),
    };

    if (pre_allocate_slots) {
        // Initiating big ordered maps with 5 pre-allocated storage slots.
        // (expiration time, address, nonce) is together 48 bytes.
        // A 4 KB storage slot can store 80+ such tuples.
        // The 5 slots should be more than enough for the current use case.
        bucket.nonces_ordered_by_exp_time.allocate_spare_slots(5);
        bucket.nonce_to_exp_time_map.allocate_spare_slots(5);
    };
    bucket
}

public entry fun add_nonce_buckets(count: u64) acquires NonceHistory {
    assert!(exists<NonceHistory>(@aptos_framework), error::invalid_state(E_NONCE_HISTORY_DOES_NOT_EXIST));
    let nonce_history = &mut NonceHistory[@aptos_framework];
    for (i in 0..count) {
        if (nonce_history.next_key <= NUM_BUCKETS) {
            if (!nonce_history.nonce_table.contains(nonce_history.next_key)) {
                nonce_history.nonce_table.add(
                    nonce_history.next_key,
                    empty_bucket(true)
                );
            };
            nonce_history.next_key = nonce_history.next_key + 1;
        }
    }
}

public(friend) fun check_and_insert_nonce(
    sender_address: address,
    nonce: u64,
    txn_expiration_time: u64,
): bool acquires NonceHistory {
    assert!(exists<NonceHistory>(@aptos_framework), error::invalid_state(E_NONCE_HISTORY_DOES_NOT_EXIST));
    // Check if the transaction expiration time is too far in the future.
    assert!(txn_expiration_time <= timestamp::now_seconds() + NONCE_REPLAY_PROTECTION_OVERLAP_INTERVAL_SECS, error::invalid_argument(ETRANSACTION_EXPIRATION_TOO_FAR_IN_FUTURE));
    let nonce_history = &mut NonceHistory[@aptos_framework];
    let nonce_key = NonceKey {
        sender_address,
        nonce,
    };
    let bucket_index = sip_hash_from_value(&nonce_key) % NUM_BUCKETS;
    let current_time = timestamp::now_seconds();
    if (!nonce_history.nonce_table.contains(bucket_index)) {
        nonce_history.nonce_table.add(
            bucket_index,
            empty_bucket(false)
        );
    };
    let bucket = table::borrow_mut(&mut nonce_history.nonce_table, bucket_index);

    let existing_exp_time = bucket.nonce_to_exp_time_map.get(&nonce_key);
    if (existing_exp_time.is_some()) {
        let existing_exp_time = existing_exp_time.extract();

        // If the existing (address, nonce) pair has not expired, return false.
        if (existing_exp_time >= current_time) {
            return false;
        };

        // We maintain an invariant that two transaction with the same (address, nonce) pair cannot be stored
        // in the nonce history if their transaction expiration times are less than `NONCE_REPLAY_PROTECTION_OVERLAP_INTERVAL_SECS`
        // seconds apart.
        if (txn_expiration_time <= existing_exp_time + NONCE_REPLAY_PROTECTION_OVERLAP_INTERVAL_SECS) {
            return false;
        };

        // If the existing (address, nonce) pair has expired, garbage collect it.
        bucket.nonce_to_exp_time_map.remove(&nonce_key);
        bucket.nonces_ordered_by_exp_time.remove(&NonceKeyWithExpTime {
            txn_expiration_time: existing_exp_time,
            sender_address,
            nonce,
        });
    };

    // Garbage collect upto MAX_ENTRIES_GARBAGE_COLLECTED_PER_CALL expired nonces in the bucket.
    let i = 0;
    while (i < MAX_ENTRIES_GARBAGE_COLLECTED_PER_CALL && !bucket.nonces_ordered_by_exp_time.is_empty()) {
        let (front_k, _) = bucket.nonces_ordered_by_exp_time.borrow_front();
        // We garbage collect a nonce after it has expired and the NONCE_REPLAY_PROTECTION_OVERLAP_INTERVAL_SECS
        // seconds have passed.
        if (front_k.txn_expiration_time + NONCE_REPLAY_PROTECTION_OVERLAP_INTERVAL_SECS < current_time) {
            bucket.nonces_ordered_by_exp_time.pop_front();
            bucket.nonce_to_exp_time_map.remove(&NonceKey {
                sender_address: front_k.sender_address,
                nonce: front_k.nonce,
            });
        } else {
            break;
        };
        i = i + 1;
    };

    // Insert the (address, nonce) pair in the bucket.
    let nonce_key_with_exp_time = NonceKeyWithExpTime {
        txn_expiration_time,
        sender_address,
        nonce,
    };
    bucket.nonces_ordered_by_exp_time.add(nonce_key_with_exp_time, true);
    bucket.nonce_to_exp_time_map.add(nonce_key, txn_expiration_time);
    true
}