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

ab_system_contract_simple_wallet_base/
lib.rs
