ab_system_contract_simple_wallet_base/
utils.rs

1use crate::{SimpleWalletBaseExt, WalletState};
2use ab_contracts_common::ContractError;
3use ab_contracts_common::env::{Env, MethodContext};
4use ab_contracts_standards::tx_handler::{TxHandlerPayload, TxHandlerSeal, TxHandlerSlots};
5use ab_core_primitives::address::Address;
6use ab_core_primitives::transaction::TransactionHeader;
7use ab_io_type::IoType;
8use ab_io_type::trivial_type::TrivialType;
9use ab_io_type::variable_bytes::VariableBytes;
10use ab_system_contract_state::StateExt;
11use core::mem::MaybeUninit;
12
13/// Utility function to initialize the state of the wallet in a typical setup
14#[inline(always)]
15pub fn initialize_state(env: &mut Env<'_>, public_key: &[u8; 32]) -> Result<(), ContractError> {
16    let state =
17        env.simple_wallet_base_initialize(Address::SYSTEM_SIMPLE_WALLET_BASE, public_key)?;
18
19    env.state_initialize(
20        MethodContext::Reset,
21        Address::SYSTEM_STATE,
22        &env.own_address(),
23        &VariableBytes::from_buffer(state.as_bytes(), &state.size()),
24    )
25}
26
27/// Utility function to authorize transaction with the wallet in a typical setup
28#[inline(always)]
29pub fn authorize(
30    env: &Env<'_>,
31    header: &TransactionHeader,
32    read_slots: &TxHandlerSlots,
33    write_slots: &TxHandlerSlots,
34    payload: &TxHandlerPayload,
35    seal: &TxHandlerSeal,
36) -> Result<(), ContractError> {
37    let state = load_current_state(env)?;
38
39    env.simple_wallet_base_authorize(
40        Address::SYSTEM_SIMPLE_WALLET_BASE,
41        &state,
42        header,
43        read_slots,
44        write_slots,
45        payload,
46        seal,
47    )
48}
49
50/// Utility function to execute transaction with the wallet in a typical setup and increase nonce
51/// afterward
52#[inline(always)]
53pub fn execute(
54    env: &mut Env<'_>,
55    header: &TransactionHeader,
56    read_slots: &TxHandlerSlots,
57    write_slots: &TxHandlerSlots,
58    payload: &TxHandlerPayload,
59    seal: &TxHandlerSeal,
60) -> Result<(), ContractError> {
61    // Only execution environment can make a direct call here
62    if env.caller() != Address::NULL {
63        return Err(ContractError::Forbidden);
64    }
65
66    // Read existing state
67    let old_state = load_current_state(env)?;
68
69    env.simple_wallet_base_execute(
70        MethodContext::Replace,
71        Address::SYSTEM_SIMPLE_WALLET_BASE,
72        header,
73        read_slots,
74        write_slots,
75        payload,
76        seal,
77    )?;
78
79    // Manual state management due to the possibility that one of the calls during execution above
80    // may update the state too (like changing public key)
81    {
82        // Fill `new_state` with updated `old_state` containing increased nonce
83        let new_state =
84            env.simple_wallet_base_increase_nonce(Address::SYSTEM_SIMPLE_WALLET_BASE, &old_state)?;
85        // Write new state of the contract, this can only be done by the direct owner
86        env.state_compare_and_write(
87            MethodContext::Reset,
88            Address::SYSTEM_STATE,
89            &env.own_address(),
90            &VariableBytes::from_buffer(old_state.as_bytes(), &old_state.size()),
91            &VariableBytes::from_buffer(new_state.as_bytes(), &new_state.size()),
92        )
93        .map(|_| ())
94    }
95}
96
97/// Utility function to change public key of the wallet in a typical setup
98#[inline(always)]
99pub fn change_public_key(env: &mut Env<'_>, public_key: &[u8; 32]) -> Result<(), ContractError> {
100    // Only the system simple wallet base contract under the context of this contract is allowed
101    // to change public key
102    if !(env.context() == env.own_address() && env.caller() == Address::SYSTEM_SIMPLE_WALLET_BASE) {
103        return Err(ContractError::Forbidden);
104    }
105
106    // Read existing state
107    let old_state = load_current_state(env)?;
108    // Fill `new_state` with updated `old_state` containing new public key
109    let new_state = env.simple_wallet_base_change_public_key(
110        Address::SYSTEM_SIMPLE_WALLET_BASE,
111        &old_state,
112        public_key,
113    )?;
114    // Write new state of the contract, this can only be done by the direct owner
115    env.state_write(
116        MethodContext::Reset,
117        Address::SYSTEM_STATE,
118        &env.own_address(),
119        &VariableBytes::from_buffer(new_state.as_bytes(), &new_state.size()),
120    )
121}
122
123#[inline(always)]
124fn load_current_state(env: &Env<'_>) -> Result<WalletState, ContractError> {
125    let current_state = {
126        let mut current_state = MaybeUninit::<WalletState>::uninit();
127        let mut current_state_size = 0;
128        env.state_read(
129            Address::SYSTEM_STATE,
130            &env.own_address(),
131            &mut VariableBytes::from_uninit(current_state.as_bytes_mut(), &mut current_state_size),
132        )?;
133        if current_state_size != WalletState::SIZE {
134            return Err(ContractError::BadOutput);
135        }
136        // Just initialized
137        unsafe { current_state.assume_init() }
138    };
139
140    Ok(current_state)
141}
ab_system_contract_simple_wallet_base/utils.rs

ab_system_contract_simple_wallet_base/
utils.rs
