Skip to content

Commit

Permalink
give the timelock the same nonce treatment
Browse files Browse the repository at this point in the history
  • Loading branch information
@moodysalem
moodysalem committed May 6, 2024
1 parent c901128 commit 8ab9882
Show file tree
Hide file tree
Showing 2 changed files with 70 additions and 45 deletions.
95 changes: 60 additions & 35 deletions src/timelock.cairo
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ pub trait ITimelock<TContractState> {
fn cancel(ref self: TContractState, id: felt252);

// Execute a list of calls that have previously been queued. Anyone may call this.
fn execute(ref self: TContractState, calls: Span<Call>) -> Array<Span<felt252>>;
fn execute(ref self: TContractState, id: felt252, calls: Span<Call>) -> Span<Span<felt252>>;

// Return the execution window, i.e. the start and end timestamp in which the call can be executed
fn get_execution_window(self: @TContractState, id: felt252) -> ExecutionWindow;
Expand Down Expand Up @@ -102,11 +102,18 @@ pub mod Timelock {
Executed: Executed,
}

#[derive(Copy, Drop, starknet::Store)]
struct BatchState {
calls_hash: felt252,
execution_state: ExecutionState,
}

#[storage]
struct Storage {
owner: ContractAddress,
config: Config,
execution_state: LegacyMap<felt252, ExecutionState>,
batch_state: LegacyMap<felt252, BatchState>,
nonce: u64,
}

#[constructor]
Expand All @@ -115,16 +122,21 @@ pub mod Timelock {
self.config.write(config);
}

// Take a list of calls and convert it to a unique identifier for the execution
// Two lists of calls will always have the same ID if they are equivalent
// A list of calls can only be queued and executed once. To make 2 different calls, add an empty call.
pub fn to_calls_id(mut calls: Span<Call>) -> felt252 {
pub fn hash_calls(mut calls: Span<Call>) -> felt252 {
PoseidonTrait::new()
.update(selector!("governance::timelock::Timelock::to_calls_id"))
.update(selector!("governance::timelock::Timelock::hash_calls"))
.update_with(@calls)
.finalize()
}

pub fn get_batch_id(address: ContractAddress, nonce: u64) -> felt252 {
PoseidonTrait::new()
.update(selector!("governance::timelock::Timelock::get_batch_id"))
.update_with(address)
.update_with(nonce)
.finalize()
}

#[generate_trait]
impl TimelockInternal of TimelockInternalTrait {
fn check_owner(self: @ContractState) {
Expand All @@ -141,17 +153,21 @@ pub mod Timelock {
fn queue(ref self: ContractState, calls: Span<Call>) -> felt252 {
self.check_owner();

let id = to_calls_id(calls);
let execution_state = self.execution_state.read(id);
let nonce = self.nonce.read();
self.nonce.write(nonce + 1);

assert(execution_state.executed.is_zero(), 'ALREADY_EXECUTED');
assert(execution_state.canceled.is_zero(), 'HAS_BEEN_CANCELED');
assert(execution_state.created.is_zero(), 'ALREADY_QUEUED');
let id = get_batch_id(get_contract_address(), nonce);

self
.execution_state
.batch_state
.write(
id, ExecutionState { created: get_block_timestamp(), executed: 0, canceled: 0 }
id,
BatchState {
calls_hash: hash_calls(calls),
execution_state: ExecutionState {
created: get_block_timestamp(), executed: 0, canceled: 0
}
}
);

self.emit(Queued { id, calls });
Expand All @@ -162,32 +178,38 @@ pub mod Timelock {
fn cancel(ref self: ContractState, id: felt252) {
self.check_owner();

let execution_state = self.execution_state.read(id);
assert(execution_state.created.is_non_zero(), 'DOES_NOT_EXIST');
assert(execution_state.executed.is_zero(), 'ALREADY_EXECUTED');
assert(execution_state.canceled.is_zero(), 'ALREADY_CANCELED');
let batch_state = self.batch_state.read(id);
assert(batch_state.execution_state.created.is_non_zero(), 'DOES_NOT_EXIST');
assert(batch_state.execution_state.executed.is_zero(), 'ALREADY_EXECUTED');
assert(batch_state.execution_state.canceled.is_zero(), 'ALREADY_CANCELED');

self
.execution_state
.batch_state
.write(
id,
ExecutionState {
created: execution_state.created,
executed: execution_state.executed,
canceled: get_block_timestamp()
BatchState {
calls_hash: batch_state.calls_hash,
execution_state: ExecutionState {
created: batch_state.execution_state.created,
executed: 0,
canceled: get_block_timestamp()
}
}
);

self.emit(Canceled { id });
}

fn execute(ref self: ContractState, mut calls: Span<Call>) -> Array<Span<felt252>> {
let id = to_calls_id(calls);
fn execute(
ref self: ContractState, id: felt252, mut calls: Span<Call>
) -> Span<Span<felt252>> {
let batch_state = self.batch_state.read(id);

let execution_state = self.execution_state.read(id);
let calls_hash = hash_calls(calls);

assert(execution_state.executed.is_zero(), 'ALREADY_EXECUTED');
assert(execution_state.canceled.is_zero(), 'HAS_BEEN_CANCELED');
assert(batch_state.calls_hash == calls_hash, 'CALLS_HASH_MISMATCH');
assert(batch_state.execution_state.executed.is_zero(), 'ALREADY_EXECUTED');
assert(batch_state.execution_state.canceled.is_zero(), 'HAS_BEEN_CANCELED');

let execution_window = self.get_execution_window(id);
let time_current = get_block_timestamp();
Expand All @@ -196,13 +218,16 @@ pub mod Timelock {
assert(time_current < execution_window.latest, 'TOO_LATE');

self
.execution_state
.batch_state
.write(
id,
ExecutionState {
created: execution_state.created,
executed: time_current,
canceled: execution_state.canceled
BatchState {
calls_hash,
execution_state: ExecutionState {
created: batch_state.execution_state.created,
executed: time_current,
canceled: batch_state.execution_state.canceled
}
}
);

Expand All @@ -214,11 +239,11 @@ pub mod Timelock {

self.emit(Executed { id });

results
results.span()
}

fn get_execution_window(self: @ContractState, id: felt252) -> ExecutionWindow {
let created = self.execution_state.read(id).created;
let created = self.batch_state.read(id).execution_state.created;

// this prevents the 0 timestamp for created from being considered valid and also executed
assert(created.is_non_zero(), 'DOES_NOT_EXIST');
Expand Down
20 changes: 10 additions & 10 deletions src/timelock_test.cairo
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ use governance::execution_state::{ExecutionState};
use governance::interfaces::erc20::{IERC20Dispatcher, IERC20DispatcherTrait};
use governance::test::test_token::{deploy as deploy_token};
use governance::timelock::{
ITimelockDispatcher, ITimelockDispatcherTrait, Timelock, Config, Timelock::{to_calls_id}
ITimelockDispatcher, ITimelockDispatcherTrait, Timelock, Config, Timelock::{hash_calls}
};
use starknet::account::{Call};
use starknet::{
Expand Down Expand Up @@ -69,7 +69,7 @@ fn test_queue_execute() {

set_block_timestamp(86401);

timelock.execute(single_call(transfer_call(token, recipient, 500_u256)));
timelock.execute(id, single_call(transfer_call(token, recipient, 500_u256)));
assert(token.balanceOf(recipient) == 500_u256, 'balance');
}

Expand All @@ -89,7 +89,7 @@ fn test_queue_cancel() {
set_block_timestamp(86401);
timelock.cancel(id);

timelock.execute(single_call(transfer_call(token, recipient, 500_u256)));
timelock.execute(id, single_call(transfer_call(token, recipient, 500_u256)));
}

#[test]
Expand Down Expand Up @@ -120,12 +120,12 @@ fn test_queue_execute_twice() {

let recipient = contract_address_const::<12345>();

timelock.queue(single_call(transfer_call(token, recipient, 500_u256)));
let id = timelock.queue(single_call(transfer_call(token, recipient, 500_u256)));

set_block_timestamp(86401);

timelock.execute(single_call(transfer_call(token, recipient, 500_u256)));
timelock.execute(single_call(transfer_call(token, recipient, 500_u256)));
timelock.execute(id, single_call(transfer_call(token, recipient, 500_u256)));
timelock.execute(id, single_call(transfer_call(token, recipient, 500_u256)));
}

#[test]
Expand All @@ -143,7 +143,7 @@ fn test_queue_executed_too_early() {

let execution_window = timelock.get_execution_window(id);
set_block_timestamp(execution_window.earliest - 1);
timelock.execute(single_call(transfer_call(token, recipient, 500_u256)));
timelock.execute(id, single_call(transfer_call(token, recipient, 500_u256)));
}

#[test]
Expand All @@ -161,12 +161,12 @@ fn test_queue_executed_too_late() {

let execution_window = timelock.get_execution_window(id);
set_block_timestamp(execution_window.latest);
timelock.execute(single_call(transfer_call(token, recipient, 500_u256)));
timelock.execute(id, single_call(transfer_call(token, recipient, 500_u256)));
}


#[test]
fn test_no_collision_to_calls_id() {
fn test_no_collision_hash_calls() {
let call_a1 = Call {
to: contract_address_const::<1>(), selector: 2, calldata: array![3, 4].span()
};
Expand All @@ -184,5 +184,5 @@ fn test_no_collision_to_calls_id() {
let calls_a: Span<Call> = array![call_a1, call_a2].span();
let calls_b: Span<Call> = array![call_b1, call_b2].span();

assert_ne!(to_calls_id(calls_a), to_calls_id(calls_b));
assert_ne!(hash_calls(calls_a), hash_calls(calls_b));
}

0 comments on commit 8ab9882

Please sign in to comment.