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

ab_system_contract_simple_wallet_base/
utils.rs
