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

ab_system_contract_simple_wallet_base/
lib.rs
