Skip to content

Commit

Permalink
refacto(onchain): remove useless components and protect initialize
Browse files Browse the repository at this point in the history
  • Loading branch information
@0xLucqs
0xLucqs committed Jun 27, 2024
1 parent 4c08cda commit 466997c
Show file tree
Hide file tree
Showing 9 changed files with 95 additions and 201 deletions.
129 changes: 26 additions & 103 deletions onchain/src/contracts/account/account.cairo
View file
Original file line number Diff line number Diff line change
@@ -1,41 +1,23 @@
#[starknet::contract(account)]
mod VaultAccount {
use core::box::BoxTrait;
use core::hash::Hash;
use core::option::OptionTrait;
use core::result::ResultTrait;
use core::starknet::{get_tx_info, SyscallResultTrait};
use openzeppelin::account::AccountComponent;
use openzeppelin::account::interface::ISRC6;
use openzeppelin::introspection::src5::SRC5Component;
use openzeppelin::token::erc20::interface::{IERC20Dispatcher, IERC20DispatcherTrait};
use openzeppelin::upgrades::UpgradeableComponent;
use openzeppelin::upgrades::interface::IUpgradeable;
use openzeppelin::utils::cryptography::snip12::{
OffchainMessageHashImpl, StructHash, SNIP12Metadata
use core::{box::BoxTrait, hash::Hash, option::OptionTrait, result::ResultTrait,};
use openzeppelin::{
account::AccountComponent, account::interface::ISRC6, introspection::src5::SRC5Component,
token::erc20::interface::{IERC20Dispatcher, IERC20DispatcherTrait},
upgrades::UpgradeableComponent, upgrades::interface::IUpgradeable,
utils::cryptography::snip12::{OffchainMessageHashImpl, StructHash, SNIP12Metadata}
};
use starknet::account::Call;
use starknet::secp256_trait::is_valid_signature;
use starknet::secp256r1::{Secp256r1Point, Secp256r1Impl};
use starknet::{ContractAddress, ClassHash};
use starknet::{get_caller_address, contract_address_const, get_contract_address};
use vault::components::spending_limit::weekly::interface::IWeeklyLimit;
use vault::components::{
WeeklyLimitComponent, WhitelistComponent, TransactionApprovalComponent,
OutsideExecutionComponent
use starknet::{
account::Call, get_tx_info, SyscallResultTrait, secp256_trait::is_valid_signature,
secp256r1::{Secp256r1Point, Secp256r1Impl}, {ContractAddress, ClassHash},
{get_caller_address, contract_address_const, get_contract_address}
};
use vault::{
components::OutsideExecutionComponent, contracts::account::interface::IVaultAccount,
};
use vault::contracts::account::interface::{IVaultAccount, IClaimLink};
use vault::utils::claim::Claim;

component!(path: AccountComponent, storage: account, event: AccountEvent);
component!(path: SRC5Component, storage: src5, event: SRC5Event);
component!(
path: TransactionApprovalComponent,
storage: transaction_approval,
event: TransactionApprovalEvent
);
component!(path: WeeklyLimitComponent, storage: weekly_limit, event: WeeklyLimitEvent);
component!(path: WhitelistComponent, storage: whitelist, event: WhitelistEvent);
component!(
path: OutsideExecutionComponent, storage: outside_execution, event: OutsideExecutionEvent
);
Expand All @@ -56,23 +38,6 @@ mod VaultAccount {
OutsideExecutionComponent::OutsideExecution_V2Impl<ContractState>;
impl OutsideExecution_V2InternalImpl = OutsideExecutionComponent::InternalImpl<ContractState>;

// Weekly Limit
impl WeeklyLimitInternalImpl = WeeklyLimitComponent::InternalImpl<ContractState>;

// Whitelisting
impl WhitelistContractsInternalImpl = WhitelistComponent::WhitelistContractsImpl<ContractState>;
impl WhitelistClassHashesInternalImpl =
WhitelistComponent::WhitelistClassHashesImpl<ContractState>;
impl WhitelistEntrypointsInternalImpl =
WhitelistComponent::WhitelistEntrypointsImpl<ContractState>;
impl WhitelistContractEntrypointInternalImpl =
WhitelistComponent::WhitelistContractEntrypointImpl<ContractState>;
impl WhitelistClassHashEntrypointInternalImpl =
WhitelistComponent::WhitelistClassHashEntrypointImpl<ContractState>;

// Transaction approval
impl TransactionApprovalInternalImpl =
TransactionApprovalComponent::InternalImpl<ContractState>;

// SRC5
#[abi(embed_v0)]
Expand All @@ -90,15 +55,10 @@ mod VaultAccount {
claims: LegacyMap<felt252, bool>,
usdc_address: ContractAddress,
public_key: (u256, u256),
admin_address: ContractAddress,
#[substorage(v0)]
src5: SRC5Component::Storage,
#[substorage(v0)]
transaction_approval: TransactionApprovalComponent::Storage,
#[substorage(v0)]
weekly_limit: WeeklyLimitComponent::Storage,
#[substorage(v0)]
whitelist: WhitelistComponent::Storage,
#[substorage(v0)]
account: AccountComponent::Storage,
#[substorage(v0)]
outside_execution: OutsideExecutionComponent::Storage,
Expand All @@ -118,11 +78,6 @@ mod VaultAccount {
#[flat]
SRC5Event: SRC5Component::Event,
#[flat]
TransactionApprovalEvent: TransactionApprovalComponent::Event,
#[flat]
WeeklyLimitEvent: WeeklyLimitComponent::Event,
#[flat]
WhitelistEvent: WhitelistComponent::Event,
#[flat]
OutsideExecutionEvent: OutsideExecutionComponent::Event,
#[flat]
Expand Down Expand Up @@ -150,12 +105,20 @@ mod VaultAccount {
ref self: ContractState,
pub_key_x: u256,
pub_key_y: u256,
approver: ContractAddress,
limit: u256,
admin_address: ContractAddress
) {
let contract_admin_address = self.admin_address.read();
if contract_admin_address != contract_address_const::<0>() {
assert!(
get_caller_address() == contract_admin_address, "Only admin can call initialize"
);
}
self.admin_address.write(admin_address);
// Verify public key validity
Secp256r1Impl::secp256_ec_new_syscall(pub_key_x, pub_key_y)
.unwrap()
.expect('Invalid public key');
self.public_key.write((pub_key_x, pub_key_y));
self.transaction_approval.initializer(:approver);
self.weekly_limit.initializer(:limit);
self.outside_execution.initializer();
self
.usdc_address
Expand All @@ -164,38 +127,9 @@ mod VaultAccount {
0x053b40a647cedfca6ca84f542a0fe36736031905a9639a7f19a3c1e66bfd5080
>()
);

// Verify public key validity
let _ = Secp256r1Impl::secp256_ec_new_syscall(pub_key_x, pub_key_y).unwrap().unwrap();
}
}

//
// ClaimLink impl
//

#[abi(embed_v0)]
impl ClaimLink of IClaimLink<ContractState> {
fn claim(ref self: ContractState, claim: Claim, signature: Array<felt252>) {
let hash = claim.get_message_hash(get_contract_address());

assert!(!self.claims.read(hash), "Link already used");
assert!(
self.is_valid_signature(hash, signature) == starknet::VALIDATED,
"Invalid signature for claim"
);

self.claims.write(hash, true);

let usdc = IERC20Dispatcher { contract_address: self.usdc_address.read() };
usdc.transfer(get_caller_address(), claim.amount);
}

#[cfg(test)]
fn set_usdc_address(ref self: ContractState, usdc_address: ContractAddress) {
self.usdc_address.write(usdc_address);
}
}

//
// SRC6 impl
Expand Down Expand Up @@ -246,17 +180,6 @@ mod VaultAccount {
}
}

//
// Weekly Limit impl
//

#[abi(embed_v0)]
impl WeeklyLimit of IWeeklyLimit<ContractState> {
fn get_weekly_limit(self: @ContractState) -> u256 {
self.weekly_limit.get_weekly_limit()
}
}

//
// Upgradeable impl
//
Expand Down
60 changes: 0 additions & 60 deletions onchain/src/contracts/account/account_test.cairo
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,71 +12,11 @@ use vault::tests::{utils, constants};
use vault::utils::claim::Claim;
use vault::utils::outside_execution::OutsideExecution;

// use vault::contracts::account::account::VaultAccount::VaultSNIP12Metadata;

//
// Claim link
//

#[test]
fn test_claim_link_valid_signature_not_already_claimed_works() {
let address = utils::setup_vault_account().contract_address;
let erc20 = utils::setup_erc20(address);
let claimer = VaultAccountABIDispatcher { contract_address: address };

claimer.set_usdc_address(erc20.contract_address);

testing::set_contract_address(contract_address_const::<0x1>());

claimer
.claim(
Claim { amount: constants::AMOUNT, nonce: 0 },
signature: constants::VALID_SIGNATURE_CLAIM()
);
}

#[test]
#[should_panic(expected: ("Invalid signature for claim", 'ENTRYPOINT_FAILED'))]
fn test_claim_link_invalid_signature_not_already_claimed_fails() {
let address = utils::setup_vault_account().contract_address;
let erc20 = utils::setup_erc20(address);
let claimer = VaultAccountABIDispatcher { contract_address: address };

claimer.set_usdc_address(erc20.contract_address);

testing::set_contract_address(contract_address_const::<0x1>());

// println!("hash: {}", Claim { amount: constants::AMOUNT, nonce: 0 }.get_message_hash(address));

claimer
.claim(
Claim { amount: constants::AMOUNT, nonce: 0 }, signature: constants::INVALID_SIGNATURE()
);
}

#[test]
#[should_panic(expected: ("Link already used", 'ENTRYPOINT_FAILED'))]
fn test_claim_link_valid_signature_already_claimed_fails() {
let address = utils::setup_vault_account().contract_address;
let erc20 = utils::setup_erc20(address);
let claimer = VaultAccountABIDispatcher { contract_address: address };

claimer.set_usdc_address(erc20.contract_address);

testing::set_contract_address(contract_address_const::<0x1>());

claimer
.claim(
Claim { amount: constants::AMOUNT, nonce: 0 },
signature: constants::VALID_SIGNATURE_CLAIM()
);
claimer
.claim(
Claim { amount: constants::AMOUNT, nonce: 0 },
signature: constants::VALID_SIGNATURE_CLAIM()
);
}

#[test]
fn test_execute_from_outside_multiple_erc20_transfers_works() {
let address = utils::setup_vault_account().contract_address;
Expand Down
12 changes: 1 addition & 11 deletions onchain/src/contracts/account/interface.cairo
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,23 +13,13 @@ trait VaultAccountABI<TState> {
fn initialize(
ref self: TState, pub_key_x: u256, pub_key_y: u256, approver: ContractAddress, limit: u256
);
fn claim(ref self: TState, claim: Claim, signature: Array<felt252>);

#[cfg(test)]
fn set_usdc_address(ref self: TState, usdc_address: ContractAddress);
}

#[starknet::interface]
trait IVaultAccount<TState> {
fn initialize(
ref self: TState, pub_key_x: u256, pub_key_y: u256, approver: ContractAddress, limit: u256
ref self: TState, pub_key_x: u256, pub_key_y: u256, admin_address: ContractAddress
);
}

#[starknet::interface]
pub trait IClaimLink<TState> {
fn claim(ref self: TState, claim: Claim, signature: Array<felt252>);

#[cfg(test)]
fn set_usdc_address(ref self: TState, usdc_address: ContractAddress);
}
10 changes: 3 additions & 7 deletions onchain/src/contracts/factory/factory.cairo
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
mod VaultFactory {
use core::starknet::SyscallResultTrait;
use openzeppelin::access::ownable::OwnableComponent;
use openzeppelin::access::ownable::interface::OwnableABI;
use openzeppelin::upgrades::UpgradeableComponent;
use openzeppelin::upgrades::interface::{
IUpgradeable, IUpgradeableDispatcher, IUpgradeableDispatcherTrait
Expand Down Expand Up @@ -100,12 +101,7 @@ mod VaultFactory {
}

fn deploy_account(
ref self: ContractState,
salt: felt252,
pub_key_x: u256,
pub_key_y: u256,
approver: ContractAddress,
limit: u256
ref self: ContractState, salt: felt252, pub_key_x: u256, pub_key_y: u256,
) {
// Owner only
self.ownable.assert_only_owner();
Expand All @@ -128,7 +124,7 @@ mod VaultFactory {

// Step 3: Set up the account
IVaultAccountDispatcher { contract_address: address }
.initialize(:pub_key_x, :pub_key_y, :approver, :limit);
.initialize(:pub_key_x, :pub_key_y, admin_address: self.owner());

// Emit event
self.emit(VaultAccountDeployed { address, salt });
Expand Down
9 changes: 1 addition & 8 deletions onchain/src/contracts/factory/interface.cairo
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,12 +3,5 @@ use starknet::{ClassHash, ContractAddress};
#[starknet::interface]
trait IFactory<TState> {
fn set_account_class_hash(ref self: TState, account_class_hash: ClassHash);
fn deploy_account(
ref self: TState,
salt: felt252,
pub_key_x: u256,
pub_key_y: u256,
approver: ContractAddress,
limit: u256
);
fn deploy_account(ref self: TState, salt: felt252, pub_key_x: u256, pub_key_y: u256);
}
4 changes: 2 additions & 2 deletions onchain/src/tests/constants.cairo
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,8 @@ mod P256_SIGNATURE_CLAIM {
}

mod P256_SIGNATURE_EXECUTE_FROM_OUTSIDE {
const R: u256 = 0x308c7afff650ae588153e1e05de97c5d9a37303e84c1c7cdaf57f8e9864c3282;
const S: u256 = 0x5026173a080a8bcd28f125060d129b119b977df90eb5884e2c5fc0e00b93bbda;
const R: u256 = 0x28ec869afade84793d60c619b2b13ac7dcc3c98916f05b3e8e927ad1e530436f;
const S: u256 = 0xc29bc8aaabeae0e640b495e3ca864069a44c163683a875b1e8aab769718b5b7e;
}

const PUBLIC_KEY: felt252 = 0x1f3c942d7f492a37608cde0d77b884a5aa9e11d2919225968557370ddb5a5aa;
Expand Down
8 changes: 1 addition & 7 deletions onchain/src/tests/utils.cairo
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,12 +37,6 @@ fn setup_erc20(recipient: ContractAddress) -> ERC20UpgradeableABIDispatcher {
}

fn setup_vault_account() -> IVaultAccountDispatcher {
setup_custom_vault_account(
approver: contract_address_const::<0>(), limit: u256 { low: 2, high: 2 }
)
}

fn setup_custom_vault_account(approver: ContractAddress, limit: u256) -> IVaultAccountDispatcher {
let calldata = array![];

// deploy
Expand All @@ -52,7 +46,7 @@ fn setup_custom_vault_account(approver: ContractAddress, limit: u256) -> IVaultA

let vault_account = IVaultAccountDispatcher { contract_address: address };

vault_account.initialize(:pub_key_x, :pub_key_y, :approver, :limit);
vault_account.initialize(:pub_key_x, :pub_key_y, admin_address: contract_address_const::<0>());

vault_account
}
Expand Down
4 changes: 2 additions & 2 deletions onchain/src/utils/claim.cairo
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
mod claim;

#[cfg(test)]
mod claim_test;
// #[cfg(test)]
// mod claim_test;

use claim::Claim;
Loading

0 comments on commit 466997c

Please sign in to comment.