#[test_only]
module battle_addr::timestamp_manipulation_test {
use std::signer;
use aptos_framework::account;
use aptos_framework::timestamp;
use aptos_framework::coin;
use aptos_framework::object;
use aptos_token_v2::token::Token;
use battle_addr::rap_battle;
use battle_addr::one_shot;
use battle_addr::cred_token::{Self, CRED};
#[test(framework = @0x1, module_owner = @battle_addr, defender = @0x100, attacker = @0x200)]
fun test_timestamp_manipulation_exploit(
framework: &signer,
module_owner: &signer,
defender: &signer,
attacker: &signer
) {
timestamp::set_time_has_started_for_testing(framework);
account::create_account_for_test(signer::address_of(module_owner));
account::create_account_for_test(signer::address_of(defender));
account::create_account_for_test(signer::address_of(attacker));
rap_battle::init_module(module_owner);
one_shot::init_module(module_owner);
cred_token::init_module(module_owner);
cred_token::register(defender);
cred_token::register(attacker);
cred_token::mint(module_owner, signer::address_of(defender), 1000);
cred_token::mint(module_owner, signer::address_of(attacker), 1000);
one_shot::mint_rapper(module_owner, signer::address_of(defender));
one_shot::mint_rapper(module_owner, signer::address_of(attacker));
let defender_token = create_test_token(defender);
let attacker_token = create_test_token(attacker);
rap_battle::go_on_stage_or_battle(defender, defender_token, 100);
let defender_skill = 50;
let challenger_skill = 50;
let total_skill = 100;
let current_time = timestamp::now_seconds();
let target_time = if (current_time % 2 == 0) {
current_time
} else {
current_time + 1
};
timestamp::update_global_time_for_test(target_time * 1000000);
rap_battle::go_on_stage_or_battle(attacker, attacker_token, 100);
let attacker_balance = coin::balance<CRED>(signer::address_of(attacker));
assert!(attacker_balance == 1100, 1);
}
#[test(framework = @0x1, module_owner = @battle_addr, defender = @0x100, attacker = @0x200)]
fun test_statistical_bias_over_multiple_battles(
framework: &signer,
module_owner: &signer,
defender: &signer,
attacker: &signer
) {
timestamp::set_time_has_started_for_testing(framework);
account::create_account_for_test(signer::address_of(module_owner));
account::create_account_for_test(signer::address_of(defender));
account::create_account_for_test(signer::address_of(attacker));
rap_battle::init_module(module_owner);
one_shot::init_module(module_owner);
cred_token::init_module(module_owner);
cred_token::register(defender);
cred_token::register(attacker);
cred_token::mint(module_owner, signer::address_of(defender), 10000);
cred_token::mint(module_owner, signer::address_of(attacker), 10000);
let wins = 0;
let i = 0;
while (i < 10) {
one_shot::mint_rapper(module_owner, signer::address_of(defender));
one_shot::mint_rapper(module_owner, signer::address_of(attacker));
let defender_token = create_test_token(defender);
let attacker_token = create_test_token(attacker);
rap_battle::go_on_stage_or_battle(defender, defender_token, 100);
let current = timestamp::now_seconds();
let favorable_time = if (current % 100 >= 50) {
current
} else {
current + (50 - (current % 100))
};
timestamp::update_global_time_for_test(favorable_time * 1000000);
let balance_before = coin::balance<CRED>(signer::address_of(attacker));
rap_battle::go_on_stage_or_battle(attacker, attacker_token, 100);
let balance_after = coin::balance<CRED>(signer::address_of(attacker));
if (balance_after > balance_before) {
wins = wins + 1;
};
i = i + 1;
};
assert!(wins >= 8, 1);
}
fun create_test_token(owner: &signer): object::Object<Token> {
abort 99
}
}
module battle_addr::rap_battle {
use std::signer;
use aptos_framework::event;
use aptos_framework::object::{Self as object, Object};
use aptos_framework::coin::{Self as coin, Coin};
- use aptos_framework::timestamp;
+ use aptos_framework::timestamp;
+ use aptos_std::random;
@@
- public entry fun go_on_stage_or_battle(
+ #[randomness]
+ public entry fun go_on_stage_or_battle(
player: &signer,
rapper_token: Object<Token>,
bet_amount: u64
) acquires BattleArena {
@@
- let rnd = timestamp::now_seconds() % (if (total_skill == 0) { 1 } else { total_skill });
+ // The #[randomness] attribute is required for aptos_std::random.
+ let rnd = random::rand_u64() % (if (total_skill == 0) { 1 } else { total_skill });
let winner = if (rnd < defender_skill) { defender_addr } else { chall_addr };