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 to change a public key to a different one
14
15#![feature(
16    const_block_items,
17    const_convert,
18    const_trait_impl,
19    maybe_uninit_as_bytes,
20    mut_restriction,
21    slice_ptr_get
22)]
23#![cfg_attr(feature = "no-panic", feature(const_closures))]
24#![no_std]
25
26pub mod payload;
27pub mod seal;
28pub mod utils;
29
30use crate::payload::{
31    TransactionMethodContext, TransactionPayloadCallError, TransactionPayloadCaller,
32};
33use crate::seal::hash_and_verify;
34use ab_contracts_common::env::{Env, MethodContext};
35use ab_contracts_common::{ContractError, MAX_TOTAL_METHOD_ARGS};
36use ab_contracts_macros::contract;
37use ab_contracts_standards::tx_handler::{TxHandlerPayload, TxHandlerSeal, TxHandlerSlots};
38use ab_core_primitives::transaction::TransactionHeader;
39use ab_io_type::trivial_type::TrivialType;
40use core::ffi::c_void;
41use core::mem::MaybeUninit;
42use core::ptr;
43use schnorrkel::PublicKey;
44
45/// Context for transaction signatures, see [`SigningContext`].
46///
47/// [`SigningContext`]: schnorrkel::context::SigningContext
48///
49/// This constant is helpful for frontend/hardware wallet implementations.
50pub const SIGNING_CONTEXT: &[u8] = b"system-simple-wallet";
51/// Size of the buffer (in pointers) that is used for `ExternalArgs` pointers.
52///
53/// This constant is helpful for transaction generation to check whether a created transaction
54/// doesn't exceed this limit.
55///
56/// `#[slot]` argument using one pointer, `#[input]` and `#[output]` use one pointer and size +
57/// capacity each.
58pub const EXTERNAL_ARGS_BUFFER_SIZE: usize = MAX_TOTAL_METHOD_ARGS as usize
59    * (size_of::<*mut c_void>() + size_of::<u32>() * 2)
60    / size_of::<*mut c_void>();
61/// Size of the buffer in `u128` elements that holds the memory of the outputs that a payload
62/// declares.
63///
64/// This constant is helpful for transaction generation to check whether a created transaction
65/// doesn't exceed this limit.
66///
67/// This defines how big the outputs that a payload declares could be together. Each of them starts
68/// at the next multiple of 16 bytes ([`Alignment::MAX`]) after the previous one.
69///
70/// Overflow will result in an error before any method is called.
71///
72/// [`Alignment::MAX`]: ab_io_type::alignment::Alignment::MAX
73pub const OUTPUT_BUFFER_SIZE: usize = 32 * 1024 / size_of::<u128>();
74/// The largest number of outputs that a payload can declare.
75///
76/// This constant is helpful for transaction generation to check whether a created transaction
77/// doesn't exceed this limit.
78///
79/// Every `#[output]` argument and return value of the payload's methods writes one of the declared
80/// outputs, and several methods can write the same one. Slots, inputs and outputs refer to a
81/// declared output by its index, which is below this limit.
82///
83/// Overflow will result in an error before any method is called.
84pub const MAX_OUTPUTS: usize = 16;
85
86/// Transaction seal.
87///
88/// Contains signature and nonce, this is necessary to produce a correctly sealed transaction.
89#[derive(Debug, Copy, Clone, TrivialType)]
90#[repr(C)]
91pub struct Seal {
92    pub signature: [u8; 64],
93    pub nonce: u64,
94}
95
96/// State of the wallet.
97///
98/// Shouldn't be necessary to use directly.
99#[derive(Debug, Copy, Clone, Eq, PartialEq, TrivialType)]
100#[repr(C)]
101pub struct WalletState {
102    pub public_key: [u8; 32],
103    pub nonce: u64,
104}
105
106/// A simple wallet contract base contract to be used by other contracts.
107///
108/// See the module description for details.
109#[derive(Debug, Copy, Clone, TrivialType)]
110#[repr(C)]
111pub struct SimpleWalletBase;
112
113#[contract]
114impl SimpleWalletBase {
115    /// Returns initial state with a provided public key
116    #[view]
117    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
118    pub fn initialize(#[input] &public_key: &[u8; 32]) -> Result<WalletState, ContractError> {
119        // TODO: Storing some lower-level representation of the public key might reduce the cost of
120        //  verification in `Self::authorize()` method
121        // Ensure public key is valid
122        PublicKey::from_bytes(&public_key).map_err(|_error| ContractError::BadInput)?;
123
124        Ok(WalletState {
125            public_key,
126            nonce: 0,
127        })
128    }
129
130    /// Reads state of `owner` and returns `Ok(())` if authorization succeeds.
131    ///
132    /// `seal` must hold exactly the bytes of a [`Seal`].
133    #[view]
134    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
135    pub fn authorize(
136        #[input] state: &WalletState,
137        #[input] header: &TransactionHeader,
138        #[input] read_slots: &TxHandlerSlots,
139        #[input] write_slots: &TxHandlerSlots,
140        #[input] payload: &TxHandlerPayload,
141        #[input] seal: &TxHandlerSeal,
142    ) -> Result<(), ContractError> {
143        let Some(seal) = Seal::try_read_unaligned_exact(seal.get_initialized()) else {
144            return Err(ContractError::BadInput);
145        };
146
147        let expected_nonce = state.nonce;
148        // Check if max nonce value was already reached
149        if expected_nonce.checked_add(1).is_none() {
150            return Err(ContractError::Forbidden);
151        }
152
153        let public_key = PublicKey::from_bytes(state.public_key.as_ref())
154            .expect("Guaranteed by constructor; qed");
155        hash_and_verify(
156            &public_key,
157            expected_nonce,
158            header,
159            read_slots.get_initialized(),
160            write_slots.get_initialized(),
161            payload.get_initialized(),
162            &seal,
163        )
164    }
165
166    /// Executes provided transactions in the payload.
167    ///
168    /// IMPORTANT:
169    /// * *must only be called with trusted input*, for example, successful signature verification
170    ///   in [`SimpleWalletBase::authorize()`] implies transaction was seen and verified by the user
171    /// * *remember to also [`SimpleWalletBase::increase_nonce()`] afterward* unless there is a very
172    ///   good reason not to (like when wallet was replaced with another implementation containing a
173    ///   different state)
174    ///
175    /// The caller must set themselves as a context or else error will be returned.
176    #[update]
177    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
178    pub fn execute(
179        #[env] env: &mut Env<'_>,
180        #[input] header: &TransactionHeader,
181        #[input] read_slots: &TxHandlerSlots,
182        #[input] write_slots: &TxHandlerSlots,
183        #[input] payload: &TxHandlerPayload,
184        #[input] seal: &TxHandlerSeal,
185    ) -> Result<(), ContractError> {
186        let _: &_ = header;
187        let _: &_ = read_slots;
188        let _: &_ = write_slots;
189        let _: &_ = seal;
190
191        // Only allow direct calls by context owner
192        if env.caller() != env.context() {
193            return Err(ContractError::Forbidden);
194        }
195
196        let mut external_args_buffer = [ptr::null_mut(); EXTERNAL_ARGS_BUFFER_SIZE];
197        let mut output_buffer = [MaybeUninit::uninit(); OUTPUT_BUFFER_SIZE];
198        let mut declared_outputs = [MaybeUninit::uninit(); MAX_OUTPUTS];
199
200        // The outputs that the payload declares must fit, before any method is called
201        let Ok(mut payload_caller) = TransactionPayloadCaller::new(
202            payload.get_initialized(),
203            &mut external_args_buffer,
204            &mut output_buffer,
205            &mut declared_outputs,
206            |method_context| match method_context {
207                TransactionMethodContext::Null => MethodContext::Reset,
208                TransactionMethodContext::Wallet => MethodContext::Keep,
209            },
210        ) else {
211            return Err(ContractError::BadInput);
212        };
213
214        loop {
215            payload_caller = match payload_caller.call_next_method(env) {
216                Ok(Some(payload_caller)) => payload_caller,
217                Ok(None) => return Ok(()),
218                Err(TransactionPayloadCallError::Decoder(_error)) => {
219                    return Err(ContractError::BadInput);
220                }
221                Err(TransactionPayloadCallError::Method(error)) => return Err(error),
222            };
223        }
224    }
225
226    /// Returns state with increased nonce
227    #[view]
228    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
229    pub fn increase_nonce(#[input] state: &WalletState) -> Result<WalletState, ContractError> {
230        let nonce = state.nonce.checked_add(1).ok_or(ContractError::Forbidden)?;
231
232        Ok(WalletState {
233            public_key: state.public_key,
234            nonce,
235        })
236    }
237
238    /// Returns a new state with a changed public key
239    #[view]
240    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
241    pub fn change_public_key(
242        #[input] state: &WalletState,
243        #[input] &public_key: &[u8; 32],
244    ) -> Result<WalletState, ContractError> {
245        // Ensure a public key is valid
246        PublicKey::from_bytes(&public_key).map_err(|_error| ContractError::BadInput)?;
247
248        Ok(WalletState {
249            public_key,
250            nonce: state.nonce,
251        })
252    }
253}
ab_system_contract_simple_wallet_base/lib.rs

ab_system_contract_simple_wallet_base/
lib.rs
