iota_surfer/
surf_strategy.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::time::Duration;
6
7use iota_types::{
8    base_types::ObjectRef,
9    transaction::{CallArg, ObjectArg},
10};
11use move_binary_format::normalized::Type;
12use move_core_types::language_storage::StructTag;
13use rand::{Rng, seq::SliceRandom};
14use tokio::time::Instant;
15use tracing::debug;
16
17use crate::surfer_state::{EntryFunction, SurferState};
18
19enum InputObjectPassKind {
20    Value,
21    ByRef,
22    MutRef,
23}
24
25#[derive(Clone, Default)]
26pub struct SurfStrategy {
27    min_tx_interval: Duration,
28}
29
30impl SurfStrategy {
31    pub fn new(min_tx_interval: Duration) -> Self {
32        Self { min_tx_interval }
33    }
34
35    /// Given a state and a list of callable Move entry functions,
36    /// explore them for a while, and eventually return. This function may
37    /// not return in some situations, so its important to call it with a
38    /// timeout or select! to ensure the task doesn't block forever.
39    pub async fn surf_for_a_while(
40        &mut self,
41        state: &mut SurferState,
42        mut entry_functions: Vec<EntryFunction>,
43    ) {
44        entry_functions.shuffle(&mut state.rng);
45        for entry in entry_functions {
46            let next_tx_time = Instant::now() + self.min_tx_interval;
47            let Some(args) = Self::choose_function_call_args(state, entry.parameters).await else {
48                debug!(
49                    "Failed to choose arguments for Move function {:?}::{:?}",
50                    entry.module, entry.function
51                );
52                continue;
53            };
54            state
55                .execute_move_transaction(entry.package, entry.module, entry.function, args)
56                .await;
57            tokio::time::sleep_until(next_tx_time).await;
58        }
59    }
60
61    async fn choose_function_call_args(
62        state: &mut SurferState,
63        params: Vec<Type>,
64    ) -> Option<Vec<CallArg>> {
65        let mut args = vec![];
66        let mut chosen_owned_objects = vec![];
67        let mut failed = false;
68        for param in params {
69            let arg = match param {
70                Type::Bool => CallArg::Pure(bcs::to_bytes(&state.rng.gen::<bool>()).unwrap()),
71                Type::U8 => CallArg::Pure(bcs::to_bytes(&state.rng.gen::<u8>()).unwrap()),
72                Type::U16 => CallArg::Pure(bcs::to_bytes(&state.rng.gen::<u16>()).unwrap()),
73                Type::U32 => CallArg::Pure(bcs::to_bytes(&state.rng.gen::<u32>()).unwrap()),
74                Type::U64 => CallArg::Pure(bcs::to_bytes(&state.rng.gen::<u64>()).unwrap()),
75                Type::U128 => CallArg::Pure(bcs::to_bytes(&state.rng.gen::<u128>()).unwrap()),
76                Type::Address => CallArg::Pure(
77                    bcs::to_bytes(&state.cluster.get_addresses().choose(&mut state.rng)).unwrap(),
78                ),
79                ty @ Type::Struct { .. } => {
80                    match Self::choose_object_call_arg(
81                        state,
82                        InputObjectPassKind::Value,
83                        ty,
84                        &mut chosen_owned_objects,
85                    )
86                    .await
87                    {
88                        Some(arg) => arg,
89                        None => {
90                            failed = true;
91                            break;
92                        }
93                    }
94                }
95                Type::Reference(ty) => {
96                    match Self::choose_object_call_arg(
97                        state,
98                        InputObjectPassKind::ByRef,
99                        *ty,
100                        &mut chosen_owned_objects,
101                    )
102                    .await
103                    {
104                        Some(arg) => arg,
105                        None => {
106                            failed = true;
107                            break;
108                        }
109                    }
110                }
111                Type::MutableReference(ty) => {
112                    match Self::choose_object_call_arg(
113                        state,
114                        InputObjectPassKind::MutRef,
115                        *ty,
116                        &mut chosen_owned_objects,
117                    )
118                    .await
119                    {
120                        Some(arg) => arg,
121                        None => {
122                            failed = true;
123                            break;
124                        }
125                    }
126                }
127                Type::U256 | Type::Signer | Type::Vector(_) | Type::TypeParameter(_) => {
128                    failed = true;
129                    break;
130                }
131            };
132            args.push(arg);
133        }
134        if failed {
135            for (struct_tag, obj_ref) in chosen_owned_objects {
136                state
137                    .owned_objects
138                    .get_mut(&struct_tag)
139                    .unwrap()
140                    .insert(obj_ref);
141            }
142            None
143        } else {
144            Some(args)
145        }
146    }
147
148    async fn choose_object_call_arg(
149        state: &mut SurferState,
150        kind: InputObjectPassKind,
151        arg_type: Type,
152        chosen_owned_objects: &mut Vec<(StructTag, ObjectRef)>,
153    ) -> Option<CallArg> {
154        let type_tag = match arg_type {
155            Type::Struct {
156                address,
157                module,
158                name,
159                type_arguments,
160            } => StructTag {
161                address,
162                module,
163                name,
164                type_params: type_arguments
165                    .into_iter()
166                    .map(|t| t.into_type_tag().unwrap())
167                    .collect(),
168            },
169            _ => {
170                return None;
171            }
172        };
173        let owned = state.matching_owned_objects_count(&type_tag);
174        let shared = state.matching_shared_objects_count(&type_tag).await;
175        let immutable = state.matching_immutable_objects_count(&type_tag).await;
176
177        let total_matching_count = match kind {
178            InputObjectPassKind::Value => owned,
179            InputObjectPassKind::MutRef => owned + shared,
180            InputObjectPassKind::ByRef => owned + shared + immutable,
181        };
182        if total_matching_count == 0 {
183            return None;
184        }
185        let mut n = state.rng.gen_range(0..total_matching_count);
186        if n < owned {
187            let obj_ref = state.choose_nth_owned_object(&type_tag, n);
188            chosen_owned_objects.push((type_tag, obj_ref));
189            return Some(CallArg::Object(ObjectArg::ImmOrOwnedObject(obj_ref)));
190        }
191        n -= owned;
192        if n < shared {
193            let (id, initial_shared_version) = state.choose_nth_shared_object(&type_tag, n).await;
194            return Some(CallArg::Object(ObjectArg::SharedObject {
195                id,
196                initial_shared_version,
197                mutable: matches!(kind, InputObjectPassKind::MutRef),
198            }));
199        }
200        n -= shared;
201        let obj_ref = state.choose_nth_immutable_object(&type_tag, n).await;
202        Some(CallArg::Object(ObjectArg::ImmOrOwnedObject(obj_ref)))
203    }
204}