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

ab_system_contract_simple_wallet_base/
lib.rs
