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

ab_system_contract_simple_wallet_base/
lib.rs
