ab_example_contract_wallet/
lib.rs

1#![no_std]
2
3use ab_contracts_common::env::{Env, MethodContext, TransactionHeader};
4use ab_contracts_common::{Address, ContractError};
5use ab_contracts_io_type::trivial_type::TrivialType;
6use ab_contracts_macros::contract;
7use ab_contracts_standards::tx_handler::{
8    TxHandler, TxHandlerPayload, TxHandlerSeal, TxHandlerSlots,
9};
10use ab_system_contract_simple_wallet_base::SimpleWalletBaseExt;
11use ab_system_contract_state::{StateExt, with_state_buffer};
12
13#[derive(Debug, Copy, Clone, TrivialType)]
14#[repr(C)]
15pub struct ExampleWallet;
16
17#[contract]
18impl TxHandler for ExampleWallet {
19    #[view]
20    fn authorize(
21        #[env] env: &Env<'_>,
22        #[input] header: &TransactionHeader,
23        #[input] read_slots: &TxHandlerSlots,
24        #[input] write_slots: &TxHandlerSlots,
25        #[input] payload: &TxHandlerPayload,
26        #[input] seal: &TxHandlerSeal,
27    ) -> Result<(), ContractError> {
28        with_state_buffer(|state| {
29            env.state_read(Address::SYSTEM_STATE, &env.own_address(), state)?;
30
31            env.simple_wallet_base_authorize(
32                Address::SYSTEM_SIMPLE_WALLET_BASE,
33                state.cast_ref(),
34                header,
35                read_slots,
36                write_slots,
37                payload,
38                seal,
39            )
40        })
41    }
42
43    #[update]
44    fn execute(
45        #[env] env: &mut Env<'_>,
46        #[input] header: &TransactionHeader,
47        #[input] read_slots: &TxHandlerSlots,
48        #[input] write_slots: &TxHandlerSlots,
49        #[input] payload: &TxHandlerPayload,
50        #[input] seal: &TxHandlerSeal,
51    ) -> Result<(), ContractError> {
52        // Only execution environment is allowed to make this call
53        if env.caller() != Address::NULL {
54            return Err(ContractError::Forbidden);
55        }
56
57        env.simple_wallet_base_execute(
58            MethodContext::Replace,
59            Address::SYSTEM_SIMPLE_WALLET_BASE,
60            header,
61            read_slots,
62            write_slots,
63            payload,
64            seal,
65        )?;
66
67        // Manual state management due to the possibility that one of the calls during execution
68        // above may update the state.
69        with_state_buffer(|old_state| {
70            env.state_read(Address::SYSTEM_STATE, &env.own_address(), old_state)?;
71
72            with_state_buffer(|new_state| {
73                // Fill `state` with updated state containing increased nonce
74                env.simple_wallet_base_increase_nonce(
75                    Address::SYSTEM_SIMPLE_WALLET_BASE,
76                    old_state.cast_ref(),
77                    seal,
78                    new_state.cast_mut(),
79                )?;
80                // Write new state of the contract, this can only be done by the direct owner
81                env.state_write(
82                    MethodContext::Reset,
83                    Address::SYSTEM_STATE,
84                    &env.own_address(),
85                    new_state,
86                )
87            })
88        })
89    }
90}
91
92/// TODO: Support upgrading wallet to a different implementation
93#[contract]
94impl ExampleWallet {
95    /// Initialize a wallet with specified public key
96    #[update]
97    pub fn initialize(
98        #[env] env: &mut Env<'_>,
99        #[input] public_key: &[u8; 32],
100    ) -> Result<(), ContractError> {
101        with_state_buffer(|state| {
102            // Fill `state` with initialized wallet state
103            env.simple_wallet_base_initialize(
104                Address::SYSTEM_SIMPLE_WALLET_BASE,
105                public_key,
106                state.cast_mut(),
107            )?;
108            // Initialize state of the contract, this can only be done by the direct owner
109            env.state_initialize(
110                MethodContext::Reset,
111                Address::SYSTEM_STATE,
112                &env.own_address(),
113                state,
114            )
115        })
116    }
117
118    /// Change public key to a different one
119    #[update]
120    pub fn change_public_key(
121        #[env] env: &mut Env<'_>,
122        #[input] public_key: &[u8; 32],
123    ) -> Result<(), ContractError> {
124        // Only the system simple wallet base contract under the context of this contract is allowed
125        // to change public key
126        if !(env.context() == env.own_address()
127            && env.caller() == Address::SYSTEM_SIMPLE_WALLET_BASE)
128        {
129            return Err(ContractError::Forbidden);
130        }
131
132        with_state_buffer(|old_state| {
133            env.state_read(Address::SYSTEM_STATE, &env.own_address(), old_state)?;
134
135            with_state_buffer(|new_state| {
136                // Fill `state` with updated state containing changed public key
137                env.simple_wallet_base_change_public_key(
138                    Address::SYSTEM_SIMPLE_WALLET_BASE,
139                    old_state.cast_ref(),
140                    public_key,
141                    new_state.cast_mut(),
142                )?;
143                // Write new state of the contract, this can only be done by the direct owner
144                env.state_write(
145                    MethodContext::Reset,
146                    Address::SYSTEM_STATE,
147                    &env.own_address(),
148                    new_state,
149                )
150            })
151        })
152    }
153}
ab_example_contract_wallet/lib.rs

ab_example_contract_wallet/
lib.rs
