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::TransactionMethodContext;
31use crate::payload::caller::{OutputChunk, TransactionPayloadCallError, TransactionPayloadCaller};
32use crate::seal::hash_and_verify;
33use ab_contracts_common::env::{Env, MethodContext};
34use ab_contracts_common::{ContractError, MAX_TOTAL_METHOD_ARGS};
35use ab_contracts_macros::contract;
36use ab_contracts_standards::tx_handler::{TxHandlerPayload, TxHandlerSeal, TxHandlerSlots};
37use ab_core_primitives::transaction::TransactionHeader;
38use ab_io_type::trivial_type::TrivialType;
39use ab_io_type::unaligned::Unaligned;
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 [`OutputChunk`] elements that holds the memory of the outputs that a
62/// payload 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::<OutputChunk>();
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 nonce and signature, this is necessary to produce a correctly sealed transaction.
89#[derive(Debug, Copy, Clone, TrivialType)]
90#[repr(C)]
91pub struct Seal {
92    pub nonce: u64,
93    pub signature: [u8; 64],
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    // TODO: `no_panic` can't prove lack of panics in `hash_and_verify()` and public key `expect()`
135    // #[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
136    pub fn authorize(
137        #[input] state: &WalletState,
138        #[input] header: &TransactionHeader,
139        #[input] read_slots: &TxHandlerSlots<'_>,
140        #[input] write_slots: &TxHandlerSlots<'_>,
141        #[input] payload: &TxHandlerPayload<'_>,
142        #[input] seal: &TxHandlerSeal<'_>,
143    ) -> Result<(), ContractError> {
144        let Some(seal) = Unaligned::<Seal>::from_bytes_exact(seal.get_initialized()) else {
145            return Err(ContractError::BadInput);
146        };
147
148        let expected_nonce = state.nonce;
149        // Check if max nonce value was already reached
150        if expected_nonce.checked_add(1).is_none() {
151            return Err(ContractError::Forbidden);
152        }
153
154        let public_key = PublicKey::from_bytes(state.public_key.as_ref())
155            .expect("Guaranteed by constructor; qed");
156        hash_and_verify(
157            &public_key,
158            expected_nonce,
159            header,
160            read_slots.get_initialized(),
161            write_slots.get_initialized(),
162            payload.get_initialized(),
163            &seal.as_inner(),
164        )
165    }
166
167    /// Executes provided transactions in the payload.
168    ///
169    /// IMPORTANT:
170    /// * *must only be called with trusted input*, for example, successful signature verification
171    ///   in [`SimpleWalletBase::authorize()`] implies transaction was seen and verified by the user
172    /// * *remember to also [`SimpleWalletBase::increase_nonce()`] afterward* unless there is a very
173    ///   good reason not to (like when wallet was replaced with another implementation containing a
174    ///   different state)
175    ///
176    /// The caller must set themselves as a context or else error will be returned.
177    #[update]
178    // TODO: `no_panic` can't prove lack of panics in methods of the executor called through `Env`
179    // #[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
180    pub fn execute(
181        #[env] env: &mut Env<'_>,
182        #[input] header: &TransactionHeader,
183        #[input] read_slots: &TxHandlerSlots<'_>,
184        #[input] write_slots: &TxHandlerSlots<'_>,
185        #[input] payload: &TxHandlerPayload<'_>,
186        #[input] seal: &TxHandlerSeal<'_>,
187    ) -> Result<(), ContractError> {
188        let _: &_ = header;
189        let _: &_ = read_slots;
190        let _: &_ = write_slots;
191        let _: &_ = seal;
192
193        // Only allow direct calls by context owner
194        if env.caller() != env.context() {
195            return Err(ContractError::Forbidden);
196        }
197
198        let mut external_args_buffer = [ptr::null_mut(); EXTERNAL_ARGS_BUFFER_SIZE];
199        let mut output_buffer = [MaybeUninit::uninit(); OUTPUT_BUFFER_SIZE];
200        let mut declared_outputs = [MaybeUninit::uninit(); MAX_OUTPUTS];
201
202        // The outputs that the payload declares must fit, before any method is called
203        let Ok(mut payload_caller) = TransactionPayloadCaller::new(
204            payload.get_initialized(),
205            &mut external_args_buffer,
206            &mut output_buffer,
207            &mut declared_outputs,
208            |method_context| match method_context {
209                TransactionMethodContext::Null => MethodContext::Reset,
210                TransactionMethodContext::Wallet => MethodContext::Keep,
211            },
212        ) else {
213            return Err(ContractError::BadInput);
214        };
215
216        loop {
217            payload_caller = match payload_caller.call_next_method(env) {
218                Ok(Some(payload_caller)) => payload_caller,
219                Ok(None) => return Ok(()),
220                Err(TransactionPayloadCallError::Decoder(_error)) => {
221                    return Err(ContractError::BadInput);
222                }
223                Err(TransactionPayloadCallError::Method(error)) => return Err(error),
224            };
225        }
226    }
227
228    /// Returns state with increased nonce
229    #[view]
230    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
231    pub fn increase_nonce(#[input] state: &WalletState) -> Result<WalletState, ContractError> {
232        let nonce = state.nonce.checked_add(1).ok_or(ContractError::Forbidden)?;
233
234        Ok(WalletState {
235            public_key: state.public_key,
236            nonce,
237        })
238    }
239
240    /// Returns a new state with a changed public key
241    #[view]
242    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
243    pub fn change_public_key(
244        #[input] state: &WalletState,
245        #[input] &public_key: &[u8; 32],
246    ) -> Result<WalletState, ContractError> {
247        // Ensure a public key is valid
248        PublicKey::from_bytes(&public_key).map_err(|_error| ContractError::BadInput)?;
249
250        Ok(WalletState {
251            public_key,
252            nonce: state.nonce,
253        })
254    }
255}
ab_system_contract_simple_wallet_base/lib.rs

ab_system_contract_simple_wallet_base/
lib.rs
