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

ab_system_contract_simple_wallet_base/
utils.rs
