ab_system_contract_simple_wallet_base/
lib.rs

1//! A simple wallet contract base contract to be used by other contracts
2//!
3//! It includes the core logic, making contracts using it much more compact. The implementation is
4//! based on [`schnorrkel`] crate and its SR25519 signature scheme.
5//!
6//! It abstracts away its inner types in the public API to allow it to evolve over time.
7//!
8//! The general workflow is:
9//! * [`SimpleWalletBase::initialize`] is used for wallet initialization
10//! * [`SimpleWalletBase::authorize`] is used for authorization
11//! * [`SimpleWalletBase::execute`] is used for executing method calls contained in the payload,
12//!   followed by [`SimpleWalletBase::increase_nonce`]
13//! * [`SimpleWalletBase::change_public_key`] is used for change public key to a different one
14
15#![feature(non_null_from_ref, ptr_as_ref_unchecked, try_blocks, unchecked_shifts)]
16#![no_std]
17
18pub mod payload;
19pub mod seal;
20
21use crate::payload::{TransactionMethodContext, TransactionPayloadDecoder};
22use crate::seal::hash_and_verify;
23use ab_contracts_common::ContractError;
24use ab_contracts_common::env::{Env, MethodContext, TransactionHeader};
25use ab_contracts_io_type::trivial_type::TrivialType;
26use ab_contracts_io_type::variable_bytes::VariableBytes;
27use ab_contracts_macros::contract;
28use ab_contracts_standards::tx_handler::{TxHandlerPayload, TxHandlerSeal, TxHandlerSlots};
29use core::mem::MaybeUninit;
30use core::ptr;
31use schnorrkel::PublicKey;
32
33/// Context for transaction signatures, see [`SigningContext`].
34///
35/// [`SigningContext`]: schnorrkel::context::SigningContext
36///
37/// This constant is helpful for frontend/hardware wallet implementations.
38pub const SIGNING_CONTEXT: &[u8] = b"system-simple-wallet";
39/// Size of the buffer in pointers that is used for `ExternalArgs` pointers.
40///
41/// This constant is helpful for transaction generation to check whether a created transaction
42/// doesn't exceed this limit.
43///
44/// `#[slot]` argument using one pointer, `#[input]` two pointers and `#[output]` three pointers
45/// each.
46pub const EXTERNAL_ARGS_BUFFER_SIZE: usize = 32 * 1024;
47/// Size of the buffer in bytes that is used as a stack for storing outputs.
48///
49/// This constant is helpful for transaction generation to check whether a created transaction
50/// doesn't exceed this limit.
51///
52/// This defines how big the total size of `#[output]` arguments and return values could be in all
53/// methods of the payload together.
54///
55/// Overflow will result in an error.
56pub const OUTPUT_BUFFER_SIZE: usize = 32 * 1024;
57/// Size of the buffer in entries that is used to store buffer offsets.
58///
59/// This constant is helpful for transaction generation to check whether a created transaction
60/// doesn't exceed this limit.
61///
62/// This defines how many `#[output]` arguments and return values could exist in all methods of the
63/// payload together.
64///
65/// Overflow will result in an error.
66pub const OUTPUT_BUFFER_OFFSETS_SIZE: usize = 16;
67
68/// Transaction seal.
69///
70/// Contains signature and nonce, this is necessary to produce a correctly sealed transaction.
71#[derive(Debug, Copy, Clone, TrivialType)]
72#[repr(C)]
73pub struct Seal {
74    pub signature: [u8; 64],
75    pub nonce: u64,
76}
77
78/// State of the wallet.
79///
80/// Shouldn't be necessary to use directly.
81#[derive(Debug, Copy, Clone, TrivialType)]
82#[repr(C)]
83pub struct WalletState {
84    pub public_key: [u8; 32],
85    pub nonce: u64,
86}
87
88/// A simple wallet contract base contract to be used by other contracts.
89///
90/// See the module description for details.
91#[derive(Debug, Copy, Clone, TrivialType)]
92#[repr(C)]
93pub struct SimpleWalletBase;
94
95#[contract]
96impl SimpleWalletBase {
97    /// Returns initial state with a provided public key
98    #[view]
99    pub fn initialize(
100        #[input] &public_key: &[u8; 32],
101        #[output] state: &mut VariableBytes,
102    ) -> Result<(), ContractError> {
103        // TODO: Storing some lower-level representation of the public key might reduce the cost of
104        //  verification in `Self::authorize()` method
105        // Ensure public key is valid
106        PublicKey::from_bytes(&public_key).map_err(|_error| ContractError::BadInput)?;
107
108        if !state.copy_from(&WalletState {
109            public_key,
110            nonce: 0,
111        }) {
112            return Err(ContractError::BadInput);
113        }
114
115        Ok(())
116    }
117
118    /// Reads state of `owner` and returns `Ok(())` if authorization succeeds
119    #[view]
120    pub fn authorize(
121        #[input] state: &VariableBytes,
122        #[input] header: &TransactionHeader,
123        #[input] read_slots: &TxHandlerSlots,
124        #[input] write_slots: &TxHandlerSlots,
125        #[input] payload: &TxHandlerPayload,
126        #[input] seal: &TxHandlerSeal,
127    ) -> Result<(), ContractError> {
128        let Some(state) = state.read_trivial_type::<WalletState>() else {
129            return Err(ContractError::BadInput);
130        };
131        let Some(seal) = seal.read_trivial_type::<Seal>() else {
132            return Err(ContractError::BadInput);
133        };
134
135        let expected_nonce = state.nonce;
136        // Check if max nonce value was already reached
137        if expected_nonce.checked_add(1).is_none() {
138            return Err(ContractError::Forbidden);
139        };
140
141        let public_key = PublicKey::from_bytes(state.public_key.as_ref())
142            .expect("Guaranteed by constructor; qed");
143        hash_and_verify(
144            &public_key,
145            expected_nonce,
146            header,
147            read_slots.get_initialized(),
148            write_slots.get_initialized(),
149            payload.get_initialized(),
150            &seal,
151        )
152    }
153
154    /// Executes provided transactions in the payload.
155    ///
156    /// IMPORTANT:
157    /// * *must only be called with trusted input*, for example, successful signature verification
158    ///   in [`SimpleWalletBase::authorize()`] implies transaction was seen and verified by the user
159    /// * *remember to also [`SimpleWalletBase::increase_nonce()`] afterward* unless there is a very
160    ///   good reason not to (like when wallet was replaced with another implementation containing a
161    ///   different state)
162    ///
163    /// The caller must set themselves as a context or else error will be returned.
164    #[update]
165    pub fn execute(
166        #[env] env: &mut Env<'_>,
167        #[input] header: &TransactionHeader,
168        #[input] read_slots: &TxHandlerSlots,
169        #[input] write_slots: &TxHandlerSlots,
170        #[input] payload: &TxHandlerPayload,
171        #[input] seal: &TxHandlerSeal,
172    ) -> Result<(), ContractError> {
173        let _ = header;
174        let _ = read_slots;
175        let _ = write_slots;
176        let _ = seal;
177
178        // Only allow direct calls by context owner
179        if env.caller() != env.context() {
180            return Err(ContractError::Forbidden);
181        }
182
183        let mut external_args_buffer = [ptr::null_mut(); EXTERNAL_ARGS_BUFFER_SIZE];
184        let mut output_buffer = [MaybeUninit::uninit(); OUTPUT_BUFFER_SIZE];
185        let mut output_buffer_offsets = [MaybeUninit::uninit(); OUTPUT_BUFFER_OFFSETS_SIZE];
186
187        let mut payload_decoder = TransactionPayloadDecoder::new(
188            payload.get_initialized(),
189            &mut external_args_buffer,
190            &mut output_buffer,
191            &mut output_buffer_offsets,
192            |method_context| match method_context {
193                TransactionMethodContext::Null => MethodContext::Reset,
194                TransactionMethodContext::Wallet => MethodContext::Keep,
195            },
196        );
197
198        while let Some(prepared_method) = payload_decoder
199            .decode_next_method()
200            .map_err(|_error| ContractError::BadInput)?
201        {
202            env.call_prepared(prepared_method)?;
203        }
204
205        Ok(())
206    }
207
208    /// Returns state with increased nonce
209    #[view]
210    pub fn increase_nonce(
211        #[input] state: &VariableBytes,
212        #[input] seal: &TxHandlerSeal,
213        #[output] new_state: &mut VariableBytes,
214    ) -> Result<(), ContractError> {
215        let _ = seal;
216
217        let Some(mut state) = state.read_trivial_type::<WalletState>() else {
218            return Err(ContractError::BadInput);
219        };
220
221        state.nonce = state.nonce.checked_add(1).ok_or(ContractError::Forbidden)?;
222
223        if !new_state.copy_from(&state) {
224            return Err(ContractError::BadInput);
225        }
226
227        Ok(())
228    }
229
230    /// Returns state with a changed public key
231    #[view]
232    pub fn change_public_key(
233        #[input] state: &VariableBytes,
234        #[input] &public_key: &[u8; 32],
235        #[output] new_state: &mut VariableBytes,
236    ) -> Result<(), ContractError> {
237        let Some(mut state) = state.read_trivial_type::<WalletState>() else {
238            return Err(ContractError::BadInput);
239        };
240        // Ensure public key is valid
241        PublicKey::from_bytes(&public_key).map_err(|_error| ContractError::BadInput)?;
242
243        state.public_key = public_key;
244
245        if !new_state.copy_from(&state) {
246            return Err(ContractError::BadInput);
247        }
248
249        Ok(())
250    }
251}
ab_system_contract_simple_wallet_base/lib.rs

ab_system_contract_simple_wallet_base/
lib.rs
