Skip to content

Commit

Permalink
Add Access Control to wager contract (#66)
Browse files Browse the repository at this point in the history 
* feat: integrate access control and SRC5 components in StrkWager contract

* utils refactored to meet ADMIN role requirements in wager contract

* Admin role negative test

* test_wager refactoring
  • Loading branch information
@beeguy74
beeguy74 authored Feb 17, 2025
1 parent 824f2a8 commit e535a6b
Show file tree
Hide file tree
Showing 3 changed files with 88 additions and 19 deletions.
61 changes: 47 additions & 14 deletions contracts/src/tests/test_wager.cairo
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,22 +5,38 @@ use contracts::wager::wager::StrkWager;

use contracts::wager::interface::{IStrkWagerDispatcher, IStrkWagerDispatcherTrait};
use contracts::escrow::interface::IEscrowDispatcherTrait;
use contracts::tests::utils::{deploy_wager, create_wager, deploy_mock_erc20, deploy_escrow, OWNER};
use contracts::tests::utils::{
deploy_wager, create_wager, deploy_mock_erc20, deploy_escrow, OWNER, ADMIN
};
use openzeppelin::token::erc20::interface::IERC20DispatcherTrait;

use snforge_std::{
declare, ContractClassTrait, DeclareResultTrait, start_cheat_caller_address,
stop_cheat_caller_address, spy_events, EventSpyAssertionsTrait
};


#[test]
#[should_panic(expected: 'Caller is missing role')]
fn test_set_escrow_address_fail() {
let admin_address = ADMIN();
let new_address = contract_address_const::<'new_address'>();
let (wager, contract_address) = deploy_wager(admin_address);
start_cheat_caller_address(contract_address, new_address);
wager.set_escrow_address(new_address);
stop_cheat_caller_address(contract_address);
}

#[test]
fn test_set_escrow_address() {
let (wager, contract_address) = deploy_wager();
let admin_address = ADMIN();
let (wager, contract_address) = deploy_wager(admin_address);
let mut spy = spy_events();

let new_address = contract_address_const::<'new_address'>();

start_cheat_caller_address(wager.contract_address, admin_address);
wager.set_escrow_address(new_address);
stop_cheat_caller_address(wager.contract_address);
spy
.assert_emitted(
@array![
Expand All @@ -46,25 +62,33 @@ fn test_set_escrow_address() {
#[test]
#[should_panic(expected: ('Invalid address',))]
fn test_set_escrow_address_zero_address_fails() {
let (wager, _) = deploy_wager();
let admin_address = ADMIN();
let (wager, _) = deploy_wager(admin_address);
let zero_address: ContractAddress = contract_address_const::<0>();

start_cheat_caller_address(wager.contract_address, admin_address);
wager.set_escrow_address(zero_address);
stop_cheat_caller_address(wager.contract_address);
}

#[test]
fn test_get_escrow_address() {
let (wager, _) = deploy_wager();
let admin_address = ADMIN();
let (wager, _) = deploy_wager(admin_address);
let first_address = contract_address_const::<'new_address'>();

start_cheat_caller_address(wager.contract_address, admin_address);
wager.set_escrow_address(first_address);
stop_cheat_caller_address(wager.contract_address);
let initial_address: ContractAddress = wager.get_escrow_address();

// Check if the initial address is as expected
assert!(initial_address == first_address, "Initial escrow address is not being returned");

let second_address = contract_address_const::<'second_address'>();
start_cheat_caller_address(wager.contract_address, admin_address);
wager.set_escrow_address(second_address);
stop_cheat_caller_address(wager.contract_address);
let final_address: ContractAddress = wager.get_escrow_address();

// Check if the updated address is as expected
Expand All @@ -73,27 +97,29 @@ fn test_get_escrow_address() {

#[test]
fn test_create_wager_success() {
create_wager(3000, 2000);
let admin_address = ADMIN();
create_wager(3000, 2000, admin_address);
}

#[test]
#[should_panic(expected: 'Insufficient balance')]
fn test_create_wager_insufficient_balance() {
create_wager(2000, 2200);
let admin_address = ADMIN();
create_wager(2000, 2200, admin_address);
}

#[test]
fn test_fund_wallet_success() {
// Deploy contracts
let (wager, wager_address) = deploy_wager();
let admin_address = ADMIN();
let (wager, wager_address) = deploy_wager(admin_address);
let (escrow, strk_dispatcher) = deploy_escrow(wager_address);

// Configure wager with escrow
start_cheat_caller_address(wager.contract_address, admin_address);
wager.set_escrow_address(escrow.contract_address);

stop_cheat_caller_address(wager.contract_address);
let amount = 50_u256;
let owner = OWNER();

// Approve tokens from OWNER for escrow
start_cheat_caller_address(strk_dispatcher.contract_address, owner);
strk_dispatcher.approve(escrow.contract_address, amount);
Expand All @@ -109,7 +135,8 @@ fn test_fund_wallet_success() {
#[should_panic(expected: ('Escrow not configured',))]
fn test_fund_wallet_no_escrow() {
// Deploy wager without setting escrow
let (wager, _) = deploy_wager();
let admin_address = ADMIN();
let (wager, _) = deploy_wager(admin_address);

// Try to fund wallet without escrow configured
wager.fund_wallet(100_u256);
Expand All @@ -118,13 +145,16 @@ fn test_fund_wallet_no_escrow() {
#[test]
#[should_panic(expected: ('Amount must be positive',))]
fn test_fund_wallet_zero_amount() {
let (wager, wager_address) = deploy_wager();
let admin_address = ADMIN();
let (wager, wager_address) = deploy_wager(admin_address);

// Deploy escrow - using the correct signature without arguments
let (escrow, strk_dispatcher) = deploy_escrow(wager_address);

// Set escrow address
start_cheat_caller_address(wager.contract_address, admin_address);
wager.set_escrow_address(escrow.contract_address);
stop_cheat_caller_address(wager.contract_address);

// Test with zero amount - should panic
wager.fund_wallet(0_u256);
Expand All @@ -133,10 +163,13 @@ fn test_fund_wallet_zero_amount() {
#[test]
#[should_panic(expected: ('ERC20: insufficient allowance',))]
fn test_fund_wallet_without_approval() {
let (wager, wager_address) = deploy_wager();
let admin_address = ADMIN();
let (wager, wager_address) = deploy_wager(admin_address);
let (escrow, strk_dispatcher) = deploy_escrow(wager_address);

start_cheat_caller_address(wager.contract_address, admin_address);
wager.set_escrow_address(escrow.contract_address);
stop_cheat_caller_address(wager.contract_address);

// Try to fund without any token approval
start_cheat_caller_address(wager.contract_address, OWNER());
Expand Down
13 changes: 10 additions & 3 deletions contracts/src/tests/utils.cairo
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,10 @@ pub fn OWNER() -> ContractAddress {
'owner'.try_into().unwrap()
}

pub fn ADMIN() -> ContractAddress {
'admin'.try_into().unwrap()
}

pub fn WAGER_ADDRESS() -> ContractAddress {
'wager'.try_into().unwrap()
}
Expand Down Expand Up @@ -48,23 +52,26 @@ pub fn deploy_escrow(wager_address: ContractAddress) -> (IEscrowDispatcher, IERC
(IEscrowDispatcher { contract_address }, strk_dispatcher)
}

pub fn deploy_wager() -> (IStrkWagerDispatcher, ContractAddress) {
pub fn deploy_wager(admin_address: ContractAddress) -> (IStrkWagerDispatcher, ContractAddress) {
let contract = declare("StrkWager").unwrap().contract_class();
let mut calldata = array![];
admin_address.serialize(ref calldata);

let (contract_address, _) = contract.deploy(@calldata).unwrap();
let dispatcher = IStrkWagerDispatcher { contract_address };

(dispatcher, contract_address)
}

pub fn create_wager(deposit: u256, stake: u256) {
let (wager, wager_contract) = deploy_wager();
pub fn create_wager(deposit: u256, stake: u256, admin_address: ContractAddress) {
let (wager, wager_contract) = deploy_wager(admin_address);
let (escrow, strk_dispatcher) = deploy_escrow(wager_contract);
let creator = OWNER();
let mut spy = spy_events();

start_cheat_caller_address(wager.contract_address, admin_address);
wager.set_escrow_address(escrow.contract_address);
stop_cheat_caller_address(wager.contract_address);

cheat_caller_address(strk_dispatcher.contract_address, creator, CheatSpan::TargetCalls(1));
strk_dispatcher.approve(escrow.contract_address, deposit);
Expand Down
33 changes: 31 additions & 2 deletions contracts/src/wager/wager.cairo
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,20 @@ pub mod StrkWager {

use contracts::wager::interface::IStrkWager;
use contracts::wager::types::{Wager, Category, Mode};
use openzeppelin::introspection::src5::SRC5Component;
use openzeppelin_access::accesscontrol::{AccessControlComponent};

component!(path: AccessControlComponent, storage: accesscontrol, event: AccessControlEvent);
component!(path: SRC5Component, storage: src5, event: SRC5Event);

#[abi(embed_v0)]
impl SRC5Impl = SRC5Component::SRC5Impl<ContractState>;

#[abi(embed_v0)]
impl AccessControlImpl =
AccessControlComponent::AccessControlImpl<ContractState>;

impl AccessControlInternalImpl = AccessControlComponent::InternalImpl<ContractState>;

#[storage]
struct Storage {
Expand All @@ -19,13 +33,21 @@ pub mod StrkWager {
wager_participants: Map<u64, Map<u64, ContractAddress>>, // wager_id -> idx -> participants
wager_participants_count: Map<u64, u64>, // wager_id -> count
escrow_address: ContractAddress,
#[substorage(v0)]
accesscontrol: AccessControlComponent::Storage,
#[substorage(v0)]
src5: SRC5Component::Storage,
}

#[event]
#[derive(Drop, starknet::Event)]
pub enum Event {
EscrowAddressUpdated: EscrowAddressEvent,
WagerCreated: WagerCreatedEvent
WagerCreated: WagerCreatedEvent,
#[flat]
AccessControlEvent: AccessControlComponent::Event,
#[flat]
SRC5Event: SRC5Component::Event,
}

#[derive(Drop, starknet::Event)]
Expand All @@ -45,8 +67,14 @@ pub mod StrkWager {
pub mode: Mode,
}

const ADMIN_ROLE: felt252 = selector!("ADMIN_ROLE"); // Unique identifier for the role


#[constructor]
fn constructor(ref self: ContractState) {}
fn constructor(ref self: ContractState, admin_contract: ContractAddress) {
self.accesscontrol.initializer();
self.accesscontrol._grant_role(ADMIN_ROLE, admin_contract);
}

#[abi(embed_v0)]
impl StrkWagerImpl of IStrkWager<ContractState> {
Expand Down Expand Up @@ -134,6 +162,7 @@ pub mod StrkWager {
self.escrow_address.read()
}
fn set_escrow_address(ref self: ContractState, new_address: ContractAddress) {
self.accesscontrol.assert_only_role(ADMIN_ROLE);
assert(!new_address.is_zero(), 'Invalid address');

let old_address = self.escrow_address.read();
Expand Down

0 comments on commit e535a6b

Please sign in to comment.