Skip to main content

transaction_fuzzer/
transaction_data_gen.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use iota_sdk_types::{
6    Address, GasPayment, ObjectDigest, ObjectId, ObjectReference, Transaction,
7    TransactionExpiration, TransactionKind, TransactionV1, Version,
8};
9use move_core_types::account_address::AccountAddress;
10use proptest::{arbitrary::*, collection::vec, prelude::*};
11
12use crate::{
13    account_universe::{gas_budget_selection_strategy, gas_price_selection_strategy},
14    type_arg_fuzzer::{gen_type_tag, pt_for_tags},
15};
16
17const MAX_NUM_GAS_OBJS: usize = 1024_usize;
18
19pub fn gen_transaction_expiration_with_bound(
20    max_epoch: u64,
21) -> impl Strategy<Value = TransactionExpiration> {
22    prop_oneof![
23        Just(TransactionExpiration::None),
24        (0u64..=max_epoch).prop_map(TransactionExpiration::Epoch),
25    ]
26}
27
28pub fn gen_transaction_expiration() -> impl Strategy<Value = TransactionExpiration> {
29    prop_oneof![
30        Just(TransactionExpiration::None),
31        (0u64..=u64::MAX).prop_map(TransactionExpiration::Epoch),
32    ]
33}
34
35pub fn gen_object_ref() -> impl Strategy<Value = ObjectReference> {
36    (any::<AccountAddress>(), any::<Version>(), any::<[u8; 32]>()).prop_map(
37        move |(addr, seq, seed)| {
38            ObjectReference::new(
39                ObjectId::new(addr.into_bytes()),
40                seq,
41                ObjectDigest::new(seed),
42            )
43        },
44    )
45}
46
47pub fn gen_gas_data(sender: Address) -> impl Strategy<Value = GasPayment> {
48    (
49        vec(gen_object_ref(), 0..MAX_NUM_GAS_OBJS),
50        gas_price_selection_strategy(),
51        gas_budget_selection_strategy(),
52    )
53        .prop_map(move |(obj_refs, price, budget)| GasPayment {
54            objects: obj_refs,
55            owner: sender,
56            price,
57            budget,
58        })
59}
60
61pub fn gen_transaction_kind() -> impl Strategy<Value = TransactionKind> {
62    (vec(gen_type_tag(), 0..10))
63        .prop_map(pt_for_tags)
64        .prop_map(TransactionKind::Programmable)
65}
66
67pub fn transaction_data_gen(sender: Address) -> impl Strategy<Value = Transaction> {
68    TransactionDataGenBuilder::new(sender)
69        .kind(gen_transaction_kind())
70        .gas_data(gen_gas_data(sender))
71        .expiration(gen_transaction_expiration())
72        .finish()
73}
74
75pub struct TransactionDataGenBuilder<
76    K: Strategy<Value = TransactionKind>,
77    G: Strategy<Value = GasPayment>,
78    E: Strategy<Value = TransactionExpiration>,
79> {
80    pub kind: Option<K>,
81    pub sender: Address,
82    pub gas_data: Option<G>,
83    pub expiration: Option<E>,
84}
85
86impl<
87    K: Strategy<Value = TransactionKind>,
88    G: Strategy<Value = GasPayment>,
89    E: Strategy<Value = TransactionExpiration>,
90> TransactionDataGenBuilder<K, G, E>
91{
92    pub fn new(sender: Address) -> Self {
93        Self {
94            kind: None,
95            sender,
96            gas_data: None,
97            expiration: None,
98        }
99    }
100
101    pub fn kind(mut self, kind: K) -> Self {
102        self.kind = Some(kind);
103        self
104    }
105
106    pub fn gas_data(mut self, gas_data: G) -> Self {
107        self.gas_data = Some(gas_data);
108        self
109    }
110
111    pub fn expiration(mut self, expiration: E) -> Self {
112        self.expiration = Some(expiration);
113        self
114    }
115
116    pub fn finish(self) -> impl Strategy<Value = Transaction> {
117        (
118            self.kind.expect("kind must be set"),
119            Just(self.sender),
120            self.gas_data.expect("gas_data must be set"),
121            self.expiration.expect("expiration must be set"),
122        )
123            .prop_map(|(kind, sender, gas_data, expiration)| TransactionV1 {
124                kind,
125                sender,
126                gas_payment: gas_data,
127                expiration,
128            })
129            .prop_map(Transaction::V1)
130    }
131}