ab_system_contract_simple_wallet_base/payload/
builder.rs

1//! Transaction payload creation utilities
2
3#[cfg(test)]
4mod tests;
5
6extern crate alloc;
7
8use crate::payload::{TransactionInput, TransactionMethodContext};
9use ab_contracts_common::Address;
10use ab_contracts_common::metadata::decode::{
11    ArgumentKind, MetadataDecodingError, MethodMetadataDecoder, MethodsContainerKind,
12};
13use ab_contracts_common::method::{ExternalArgs, MethodFingerprint};
14use ab_contracts_io_type::MAX_ALIGNMENT;
15use ab_contracts_io_type::trivial_type::TrivialType;
16use alloc::vec::Vec;
17use core::ffi::c_void;
18use core::num::NonZeroU8;
19use core::ptr::NonNull;
20use core::{ptr, slice};
21
22/// Errors for [`TransactionPayloadBuilder`]
23#[derive(Debug, thiserror::Error)]
24pub enum TransactionPayloadBuilderError<'a> {
25    /// Metadata decoding error
26    #[error("Metadata decoding error: {0}")]
27    MetadataDecodingError(MetadataDecodingError<'a>),
28    /// Invalid alignment
29    #[error("Invalid alignment: {0}")]
30    InvalidAlignment(NonZeroU8),
31    /// Invalid output index
32    #[error("Invalid output index: {0}")]
33    InvalidOutputIndex(u8),
34}
35
36/// Builder for payload to be used with [`TxHandler`] (primarily for [`SimpleWallet`]).
37///
38/// Decoding can be done with [`TransactionPayloadDecoder`]
39///
40/// [`TxHandler`]: ab_contracts_standards::tx_handler::TxHandler
41/// [`SimpleWallet`]: crate::SimpleWalletBase
42/// [`TransactionPayloadDecoder`]: crate::payload::TransactionPayloadDecoder
43#[derive(Debug, Clone)]
44pub struct TransactionPayloadBuilder {
45    payload: Vec<u8>,
46}
47
48impl Default for TransactionPayloadBuilder {
49    fn default() -> Self {
50        Self {
51            payload: Vec::with_capacity(1024),
52        }
53    }
54}
55
56impl TransactionPayloadBuilder {
57    /// Add method call to the payload.
58    ///
59    /// The wallet will call this method in addition order.
60    ///
61    /// `input_output_index` is used for referencing earlier outputs as inputs of this method,
62    /// its values are optional, see [`TransactionInput`] for more details.
63    pub fn with_method_call<Args>(
64        &mut self,
65        contract: &Address,
66        external_args: &Args,
67        method_context: TransactionMethodContext,
68        input_output_index: &[Option<u8>],
69    ) -> Result<(), TransactionPayloadBuilderError<'static>>
70    where
71        Args: ExternalArgs,
72    {
73        let external_args = NonNull::from_ref(external_args).cast::<*const c_void>();
74
75        // SAFETY: Called with statically valid data
76        unsafe {
77            self.with_method_call_untyped(
78                contract,
79                &external_args,
80                Args::METADATA,
81                &Args::FINGERPRINT,
82                method_context,
83                input_output_index,
84            )
85        }
86    }
87
88    /// Other than unsafe API, this method is identical to [`Self::with_method_call()`].
89    ///
90    /// # Safety
91    /// `external_args` must correspond to `method_metadata` and `method_fingerprint`. Outputs are
92    /// never read from `external_args` and inputs that have corresponding `input_output_index`
93    /// are not read either.
94    pub unsafe fn with_method_call_untyped<'a>(
95        &mut self,
96        contract: &Address,
97        external_args: &NonNull<*const c_void>,
98        mut method_metadata: &'a [u8],
99        method_fingerprint: &MethodFingerprint,
100        method_context: TransactionMethodContext,
101        input_output_index: &[Option<u8>],
102    ) -> Result<(), TransactionPayloadBuilderError<'a>> {
103        let mut external_args = *external_args;
104
105        let (mut metadata_decoder, _method_metadata_item) =
106            MethodMetadataDecoder::new(&mut method_metadata, MethodsContainerKind::Unknown)
107                .decode_next()
108                .map_err(TransactionPayloadBuilderError::MetadataDecodingError)?;
109
110        self.extend_payload_with_alignment(contract.as_bytes(), align_of_val(contract));
111        self.extend_payload_with_alignment(
112            method_fingerprint.as_bytes(),
113            align_of_val(method_fingerprint),
114        );
115        self.payload.push(method_context as u8);
116
117        // Remember the position to update later
118        let num_slots_index = self.payload.len();
119        self.payload.push(0);
120        // Remember the position to update later
121        let num_inputs_index = self.payload.len();
122        self.payload.push(0);
123        // Remember the position to update later
124        let num_outputs_index = self.payload.len();
125        self.payload.push(0);
126
127        while let Some(item) = metadata_decoder
128            .decode_next()
129            .transpose()
130            .map_err(TransactionPayloadBuilderError::MetadataDecodingError)?
131        {
132            match item.argument_kind {
133                ArgumentKind::EnvRo
134                | ArgumentKind::EnvRw
135                | ArgumentKind::TmpRo
136                | ArgumentKind::TmpRw => {
137                    // Not represented in external args
138                }
139                ArgumentKind::SlotRo | ArgumentKind::SlotRw => {
140                    self.payload[num_slots_index] += 1;
141
142                    // SAFETY: Method description requires the layout to correspond to metadata
143                    let address = unsafe {
144                        let address = external_args.cast::<NonNull<Address>>().read().as_ref();
145                        external_args = external_args.offset(1);
146                        address
147                    };
148                    self.extend_payload_with_alignment(address.as_bytes(), align_of_val(address));
149                }
150                ArgumentKind::Input => {
151                    let input_offset = usize::from(self.payload[num_inputs_index]);
152                    self.payload[num_inputs_index] += 1;
153
154                    let type_details = &item
155                        .type_details
156                        .expect("Always present for `#[input]`; qed");
157
158                    let maybe_output_index =
159                        input_output_index.get(input_offset).copied().flatten();
160                    let input_type = match maybe_output_index {
161                        Some(output_index) => TransactionInput::new_output_index(output_index)
162                            .ok_or(TransactionPayloadBuilderError::InvalidOutputIndex(
163                                output_index,
164                            ))?,
165                        None => TransactionInput::new_value(type_details.alignment).ok_or(
166                            TransactionPayloadBuilderError::InvalidAlignment(
167                                type_details.alignment,
168                            ),
169                        )?,
170                    };
171                    self.payload.push(input_type.into_u8());
172
173                    if maybe_output_index.is_none() {
174                        // SAFETY: Method description requires the layout to correspond to metadata
175                        let (size, data) = unsafe {
176                            let data = external_args.cast::<NonNull<u8>>().read();
177                            external_args = external_args.offset(1);
178                            let size = external_args.cast::<NonNull<u32>>().read().read();
179                            external_args = external_args.offset(1);
180
181                            let data =
182                                slice::from_raw_parts(data.as_ptr().cast_const(), size as usize);
183
184                            (size, data)
185                        };
186
187                        self.extend_payload_with_alignment(
188                            &size.to_le_bytes(),
189                            align_of_val(&size),
190                        );
191                        self.extend_payload_with_alignment(
192                            data,
193                            type_details.alignment.get() as usize,
194                        );
195                    }
196                }
197                ArgumentKind::Output => {
198                    self.payload[num_outputs_index] += 1;
199
200                    // May be skipped for `#[init]`, see `ContractMetadataKind::Init` for details
201                    if let Some(type_details) = &item.type_details {
202                        self.extend_payload_with_alignment(
203                            &type_details.recommended_capacity.to_le_bytes(),
204                            align_of_val(&type_details.recommended_capacity),
205                        );
206                        self.extend_payload_with_alignment(
207                            &[type_details.alignment.ilog2() as u8],
208                            align_of::<u8>(),
209                        );
210                    }
211                }
212            }
213        }
214
215        Ok(())
216    }
217
218    /// Returns 16-byte aligned bytes.
219    ///
220    /// The contents is a concatenated sequence of method calls with their arguments. All data
221    /// structures are correctly aligned in the returned byte buffer with `0` used as padding when
222    /// necessary.
223    ///
224    /// Each method is serialized in the following way:
225    /// * Contract to call: [`Address`]
226    /// * Fingerprint of the method to call: [`MethodFingerprint`]
227    /// * Method context: [`TransactionMethodContext`]
228    /// * Number of slot arguments: `u8`
229    /// * Number of `#[input]` arguments: `u8`
230    /// * Number of `#[output]` arguments: `u8`
231    /// * Concatenated sequence of arguments
232    ///
233    /// Each argument is serialized in the following way (others are skipped):
234    /// * `#[slot]`: [`Address`]
235    /// * `#[input]`: [`TransactionInput`] as `u8`
236    ///     * If [`TransactionInput::new_value()`] then input size as little-endian `u32`
237    ///       followed by the input itself
238    /// * `#[output]`: recommended capacity as little-endian `u32` followed by alignment power as
239    ///   `u8` (`NonZeroU8::ilog2(alignment)`)
240    pub fn into_aligned_bytes(mut self) -> Vec<u128> {
241        // Fill bytes to make it multiple of `u128` before creating `u128`-based vector
242        self.ensure_alignment(usize::from(MAX_ALIGNMENT));
243
244        let output_len = self.payload.len() / size_of::<u128>();
245        let mut output = Vec::<u128>::with_capacity(output_len);
246
247        // SAFETY: Pointers are valid for reads/writes, aligned and not overlapping
248        unsafe {
249            ptr::copy_nonoverlapping(
250                self.payload.as_ptr(),
251                output.as_mut_ptr().cast::<u8>(),
252                self.payload.len(),
253            );
254            output.set_len(output_len);
255        }
256
257        debug_assert_eq!(align_of_val(output.as_slice()), usize::from(MAX_ALIGNMENT));
258
259        output
260    }
261
262    fn extend_payload_with_alignment(&mut self, bytes: &[u8], alignment: usize) {
263        self.ensure_alignment(alignment);
264
265        self.payload.extend_from_slice(bytes);
266    }
267
268    fn ensure_alignment(&mut self, alignment: usize) {
269        debug_assert!(alignment <= usize::from(MAX_ALIGNMENT));
270
271        // Optimized version of the following that expects `alignment` to be a power of 2:
272        // let unaligned_by = self.payload.len() % alignment;
273        let unaligned_by = self.payload.len() & (alignment - 1);
274        if unaligned_by > 0 {
275            self.payload
276                .resize(self.payload.len() + (alignment - unaligned_by), 0);
277        }
278    }
279}