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

ab_system_contract_simple_wallet_base/
utils.rs
