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(
16    maybe_uninit_as_bytes,
17    maybe_uninit_slice,
18    ptr_as_ref_unchecked,
19    slice_as_array,
20    try_blocks,
21    unchecked_shifts
22)]
23#![no_std]
24
25pub mod payload;
26pub mod seal;
27pub mod utils;
28
29use crate::payload::{TransactionMethodContext, TransactionPayloadDecoder};
30use crate::seal::hash_and_verify;
31use ab_contracts_common::env::{Env, MethodContext};
32use ab_contracts_common::{ContractError, MAX_TOTAL_METHOD_ARGS};
33use ab_contracts_macros::contract;
34use ab_contracts_standards::tx_handler::{TxHandlerPayload, TxHandlerSeal, TxHandlerSlots};
35use ab_core_primitives::transaction::TransactionHeader;
36use ab_io_type::trivial_type::TrivialType;
37use core::mem::MaybeUninit;
38use core::ptr;
39use schnorrkel::PublicKey;
40
41/// Context for transaction signatures, see [`SigningContext`].
42///
43/// [`SigningContext`]: schnorrkel::context::SigningContext
44///
45/// This constant is helpful for frontend/hardware wallet implementations.
46pub const SIGNING_CONTEXT: &[u8] = b"system-simple-wallet";
47/// Size of the buffer in pointers that is used for `ExternalArgs` pointers.
48///
49/// This constant is helpful for transaction generation to check whether a created transaction
50/// doesn't exceed this limit.
51///
52/// `#[slot]` argument using one pointer, `#[input]` two pointers and `#[output]` three pointers
53/// each.
54pub const EXTERNAL_ARGS_BUFFER_SIZE: usize = 3 * MAX_TOTAL_METHOD_ARGS as usize;
55/// Size of the buffer in `u128` elements that is used as a stack for storing outputs.
56///
57/// This constant is helpful for transaction generation to check whether a created transaction
58/// doesn't exceed this limit.
59///
60/// This defines how big the total size of `#[output]` arguments and return values could be in all
61/// methods of the payload together.
62///
63/// Overflow will result in an error.
64pub const OUTPUT_BUFFER_SIZE: usize = 32 * 1024 / size_of::<u128>();
65/// Size of the buffer in entries that is used to store buffer offsets.
66///
67/// This constant is helpful for transaction generation to check whether a created transaction
68/// doesn't exceed this limit.
69///
70/// This defines how many `#[output]` arguments and return values could exist in all methods of the
71/// payload together.
72///
73/// Overflow will result in an error.
74pub const OUTPUT_BUFFER_OFFSETS_SIZE: usize = 16;
75
76/// Transaction seal.
77///
78/// Contains signature and nonce, this is necessary to produce a correctly sealed transaction.
79#[derive(Debug, Copy, Clone, TrivialType)]
80#[repr(C)]
81pub struct Seal {
82    pub signature: [u8; 64],
83    pub nonce: u64,
84}
85
86/// State of the wallet.
87///
88/// Shouldn't be necessary to use directly.
89#[derive(Debug, Copy, Clone, Eq, PartialEq, TrivialType)]
90#[repr(C)]
91pub struct WalletState {
92    pub public_key: [u8; 32],
93    pub nonce: u64,
94}
95
96/// A simple wallet contract base contract to be used by other contracts.
97///
98/// See the module description for details.
99#[derive(Debug, Copy, Clone, TrivialType)]
100#[repr(C)]
101pub struct SimpleWalletBase;
102
103#[contract]
104impl SimpleWalletBase {
105    /// Returns initial state with a provided public key
106    #[view]
107    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
108    pub fn initialize(#[input] &public_key: &[u8; 32]) -> Result<WalletState, ContractError> {
109        // TODO: Storing some lower-level representation of the public key might reduce the cost of
110        //  verification in `Self::authorize()` method
111        // Ensure public key is valid
112        PublicKey::from_bytes(&public_key).map_err(|_error| ContractError::BadInput)?;
113
114        Ok(WalletState {
115            public_key,
116            nonce: 0,
117        })
118    }
119
120    /// Reads state of `owner` and returns `Ok(())` if authorization succeeds
121    #[view]
122    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
123    pub fn authorize(
124        #[input] state: &WalletState,
125        #[input] header: &TransactionHeader,
126        #[input] read_slots: &TxHandlerSlots,
127        #[input] write_slots: &TxHandlerSlots,
128        #[input] payload: &TxHandlerPayload,
129        #[input] seal: &TxHandlerSeal,
130    ) -> Result<(), ContractError> {
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    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
166    pub fn execute(
167        #[env] env: &mut Env<'_>,
168        #[input] header: &TransactionHeader,
169        #[input] read_slots: &TxHandlerSlots,
170        #[input] write_slots: &TxHandlerSlots,
171        #[input] payload: &TxHandlerPayload,
172        #[input] seal: &TxHandlerSeal,
173    ) -> Result<(), ContractError> {
174        let _ = header;
175        let _ = read_slots;
176        let _ = write_slots;
177        let _ = seal;
178
179        // Only allow direct calls by context owner
180        if env.caller() != env.context() {
181            return Err(ContractError::Forbidden);
182        }
183
184        let mut external_args_buffer = [ptr::null_mut(); EXTERNAL_ARGS_BUFFER_SIZE];
185        let mut output_buffer = [MaybeUninit::uninit(); OUTPUT_BUFFER_SIZE];
186        let mut output_buffer_offsets = [MaybeUninit::uninit(); OUTPUT_BUFFER_OFFSETS_SIZE];
187
188        let mut payload_decoder = TransactionPayloadDecoder::new(
189            payload.get_initialized(),
190            &mut external_args_buffer,
191            &mut output_buffer,
192            &mut output_buffer_offsets,
193            |method_context| match method_context {
194                TransactionMethodContext::Null => MethodContext::Reset,
195                TransactionMethodContext::Wallet => MethodContext::Keep,
196            },
197        );
198
199        while let Some(prepared_method) = payload_decoder
200            .decode_next_method()
201            .map_err(|_error| ContractError::BadInput)?
202        {
203            env.call_prepared(prepared_method)?;
204        }
205
206        Ok(())
207    }
208
209    /// Returns state with increased nonce
210    #[view]
211    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
212    pub fn increase_nonce(#[input] state: &WalletState) -> Result<WalletState, ContractError> {
213        let nonce = state.nonce.checked_add(1).ok_or(ContractError::Forbidden)?;
214
215        Ok(WalletState {
216            public_key: state.public_key,
217            nonce,
218        })
219    }
220
221    /// Returns a new state with a changed public key
222    #[view]
223    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
224    pub fn change_public_key(
225        #[input] state: &WalletState,
226        #[input] &public_key: &[u8; 32],
227    ) -> Result<WalletState, ContractError> {
228        // Ensure public key is valid
229        PublicKey::from_bytes(&public_key).map_err(|_error| ContractError::BadInput)?;
230
231        Ok(WalletState {
232            public_key,
233            nonce: state.nonce,
234        })
235    }
236}
ab_system_contract_simple_wallet_base/lib.rs

ab_system_contract_simple_wallet_base/
lib.rs
