Sui: Support One-time Witness (#451)

* add swap

* fmt

src/move/corpus_initializer.rs

@@ -19,6 +19,7 @@ use move_vm_types::{
 };
 use revm_primitives::HashSet;
 use sui_types::base_types::{TX_CONTEXT_MODULE_NAME, TX_CONTEXT_STRUCT_NAME};
+use tracing::{debug, info};
 
 use crate::{
     generic_vm::vm_executor::GenericVM,
@@ -64,6 +65,39 @@ pub fn is_tx_context(struct_tag: &StructTag) -> bool {
         struct_tag.name == TX_CONTEXT_STRUCT_NAME.into()
 }
 
+pub fn create_tx_context(caller: AccountAddress, ty: Type) -> Value {
+    match ty {
+        Type::MutableReference(ty) | Type::Reference(ty) => {
+            if let Type::Struct(_struct_tag) = *ty {
+                // struct TxContext has drop {
+                //     /// The address of the user that signed the current transaction
+                //     sender: address,
+                //     /// Hash of the current transaction
+                //     tx_hash: vector<u8>,
+                //     /// The current epoch number
+                //     epoch: u64,
+                //     /// Timestamp that the epoch started at
+                //     epoch_timestamp_ms: u64,
+                //     /// Counter recording the number of fresh id's created while executing
+                //     /// this transaction. Always 0 at the start of a transaction
+                //     ids_created: u64
+                // }
+                let inner = Container::Struct(Rc::new(RefCell::new(vec![
+                    ValueImpl::Address(caller),
+                    ValueImpl::Container(Container::VecU8(Rc::new(RefCell::new(vec![6; 32])))),
+                    ValueImpl::U64(123213),
+                    ValueImpl::U64(2130127412),
+                    ValueImpl::U64(0),
+                ])));
+
+                return Value(ValueImpl::ContainerRef(ContainerRef::Local(inner)));
+            }
+        }
+        _ => unreachable!("tx context type mismatch"),
+    }
+    unreachable!()
+}
+
 impl<'a, SC, ISC> MoveCorpusInitializer<'a, SC, ISC>
 where
     SC: Scheduler<State = MoveFuzzState>,
@@ -230,7 +264,8 @@ where
                 continue;
             }
 
-            for (_, func) in funcs {
+            for (name, func) in funcs {
+                debug!("fuzzing: {:?}::{:?}", module_id, name);
                 let input = self.build_input(&module_id, func.clone());
                 match input {
                     Some(input) => {
@@ -347,34 +382,8 @@ where
         }
     }
 
-    fn gen_tx_context(&mut self, ty: Type) -> Value {
-        if let Type::MutableReference(ty) = ty {
-            if let Type::Struct(_struct_tag) = *ty {
-                // struct TxContext has drop {
-                //     /// The address of the user that signed the current transaction
-                //     sender: address,
-                //     /// Hash of the current transaction
-                //     tx_hash: vector<u8>,
-                //     /// The current epoch number
-                //     epoch: u64,
-                //     /// Timestamp that the epoch started at
-                //     epoch_timestamp_ms: u64,
-                //     /// Counter recording the number of fresh id's created while executing
-                //     /// this transaction. Always 0 at the start of a transaction
-                //     ids_created: u64
-                // }
-                let inner = Container::Struct(Rc::new(RefCell::new(vec![
-                    ValueImpl::Address(self.state.get_rand_caller()),
-                    ValueImpl::Container(Container::VecU8(Rc::new(RefCell::new(vec![6; 32])))),
-                    ValueImpl::U64(123213),
-                    ValueImpl::U64(2130127412),
-                    ValueImpl::U64(0),
-                ])));
-
-                return Value(ValueImpl::ContainerRef(ContainerRef::Local(inner)));
-            }
-        }
-        unreachable!()
+    pub fn gen_tx_context(&mut self, ty: Type) -> Value {
+        create_tx_context(self.state.get_rand_caller(), ty)
     }
 
     fn build_input(&mut self, module_id: &ModuleId, function: Arc<Function>) -> Option<MoveFunctionInput> {

src/move/movevm.rs

@@ -37,6 +37,7 @@ use sui_types::{
     object::{Object, Owner},
     storage::ChildObjectResolver,
 };
+use tracing::debug;
 
 use super::types::MoveFuzzState;
 use crate::{
@@ -46,8 +47,8 @@ use crate::{
     },
     input::VMInputT,
     r#move::{
-        corpus_initializer::is_tx_context,
-        input::{ConciseMoveInput, MoveFunctionInput, MoveFunctionInputT},
+        corpus_initializer::{create_tx_context, is_tx_context, MoveCorpusInitializer},
+        input::{ConciseMoveInput, FunctionDefaultable, MoveFunctionInput, MoveFunctionInputT},
         types::{MoveAddress, MoveOutput},
         vm_state::{Gate, GatedValue, MoveVMState},
     },
@@ -464,8 +465,6 @@ where
             state.metadata_map_mut().insert(TypeTagInfoMeta::new());
         }
 
-        let meta = state.metadata_map_mut().get_mut::<TypeTagInfoMeta>().unwrap();
-
         let func_off = self.loader.module_cache.read().functions.len();
         let _module_name = module.name().to_owned();
         let deployed_module_idx = module.self_id();
@@ -479,18 +478,55 @@ where
                 &module,
             )
             .expect("internal deploy error");
+
         for f in &self.loader.module_cache.read().functions[func_off..] {
             // debug!("deployed function: {:?}@{}({:?}) returns {:?}", deployed_module_idx,
             // f.name.as_str(), f.parameter_types, f.return_types());
             self.functions
                 .entry(deployed_module_idx.clone())
                 .or_default()
                 .insert(f.name.to_owned(), f.clone());
-
+            let meta = state.metadata_map_mut().get_mut::<TypeTagInfoMeta>().unwrap();
             for ty in &f.parameter_types {
                 meta.register_type_tag(ty.clone(), &self.loader);
             }
         }
+
+        let init_func = self.loader.module_cache.read().functions[func_off..]
+            .iter()
+            .filter(|f| f.name.as_str() == "init")
+            .map(|f| f.clone())
+            .next();
+        if let Some(init_func) = init_func {
+            let otw = &init_func.parameters;
+            debug!(
+                "init function found {:?} {_module_name} {:?}",
+                otw, init_func.parameter_types
+            );
+            let mut args = vec![];
+            if otw.len() == 1 {
+                args.push(create_tx_context(
+                    state.get_rand_caller(),
+                    init_func.parameter_types[0].clone(),
+                ));
+            }
+
+            let move_input = MoveFunctionInput {
+                module: deployed_module_idx.clone(),
+                function: init_func.name.clone(),
+                function_info: Arc::new(FunctionDefaultable::new(init_func.clone())),
+                args: vec![],
+                ty_args: vec![],
+                caller: AccountAddress::ZERO,
+                vm_state: StagedVMState::new_with_state(MoveVMState::default()),
+                vm_state_idx: 0,
+                _deps: Default::default(),
+                _resolved: true,
+            };
+            let res = self.execute(&move_input.as_any().downcast_ref::<I>().unwrap(), state);
+            println!("init function found {:?}", res);
+        }
+
         Some(*deployed_module_idx.address())
     }
 
@@ -596,7 +632,7 @@ where
             // debug!("{:?}", ret);
 
             if ret.is_err() {
-                // debug!("reverted {:?}", ret);
+                debug!("reverted {:?}", ret);
                 reverted = true;
                 break;
             }

src/move/scheduler.rs

@@ -164,7 +164,12 @@ where
 
     fn next(&mut self, state: &mut Self::State) -> Result<CorpusId, Error> {
         let mut next_idx = self.inner.next(state)?;
+        let mut retries = 0;
         loop {
+            retries += 1;
+            if retries > 10000 {
+                panic!("All functions depend on structs that are not available to be forged by the fuzzer. ");
+            }
             let tc = state.corpus().get(next_idx).expect("Missing testcase");
             let input = tc.borrow().input().clone().expect("Missing input");
             let meta = state

tests/move/swap/Move.toml

@@ -0,0 +1,10 @@
+[package]
+name = "swap"
+version = "0.0.1"
+
+
+[dependencies]
+Sui = { git = "https://github.com/MystenLabs/sui.git", subdir = "crates/sui-framework/packages/sui-framework", rev = "mainnet" }
+
+[addresses]
+swap = "0x0"

tests/move/swap/sources/Swap.move

@@ -0,0 +1,98 @@
+module swap::vault{
+    use sui::coin::{Self, Coin};
+    use sui::tx_context::{Self, TxContext};
+    use sui::balance::{Self, Balance};
+    use sui::object::{Self, ID, UID};
+    use sui::transfer;
+    use sui::event;
+    use swap::ctfa::{Self, MintA};
+    use swap::ctfb::{Self, MintB};
+
+    struct Vault<phantom A, phantom B> has key {
+        id: UID,
+        coin_a: Balance<A>,
+        coin_b: Balance<B>,
+        flashed: bool
+    }
+
+    struct Flag has copy, drop {
+        win: bool,
+        sender: address
+    }
+
+    struct Receipt {
+        id: ID,
+        a_to_b: bool,
+        repay_amount: u64
+    }
+
+    public entry fun initialize<A,B>(capa: MintA<A>, capb: MintB<B>,ctx: &mut TxContext) {
+        let vault = Vault<A, B> {
+            id: object::new(ctx),
+            coin_a: coin::into_balance(ctfa::mint_for_vault(capa, ctx)),
+            coin_b: coin::into_balance(ctfb::mint_for_vault(capb, ctx)),
+            flashed: false
+        };
+        transfer::share_object(vault);
+    }
+
+    public fun flash<A,B>(vault: &mut Vault<A,B>, amount: u64, a_to_b: bool, ctx: &mut TxContext): (Coin<A>, Coin<B>, Receipt) {
+        assert!(!vault.flashed, 1);
+        let (coin_a, coin_b) = if (a_to_b) {
+        (coin::zero<A>(ctx), coin::from_balance(balance::split(&mut vault.coin_b, amount ), ctx))
+        }
+        else {
+        (coin::from_balance(balance::split(&mut vault.coin_a, amount ), ctx), coin::zero<B>(ctx))
+        };
+
+        let receipt = Receipt {
+            id: object::id(vault),
+            a_to_b,
+            repay_amount: amount
+        };
+        vault.flashed = true;
+
+        (coin_a, coin_b, receipt)
+
+    }
+
+    public fun repay_flash<A,B>(vault: &mut Vault<A,B>, coina: Coin<A>, coinb: Coin<B>, receipt: Receipt) {
+        let Receipt {
+            id: _,
+            a_to_b: a2b,
+            repay_amount: amount
+        } = receipt;
+        if (a2b) {
+            assert!(coin::value(&coinb) >= amount, 0);
+        } else {
+            assert!(coin::value(&coina) >= amount, 1);
+        };
+        balance::join(&mut vault.coin_a, coin::into_balance(coina));
+        balance::join(&mut vault.coin_b, coin::into_balance(coinb));
+        vault.flashed = false;
+    }
+
+    public fun swap_a_to_b<A,B>(vault: &mut Vault<A,B>, coina:Coin<A>, ctx: &mut TxContext): Coin<B> {
+            let amount_out_B = coin::value(&coina) * balance::value(&vault.coin_b) / balance::value(&vault.coin_a);
+            coin::put<A>(&mut vault.coin_a, coina);
+            coin::take(&mut vault.coin_b, amount_out_B, ctx)
+    }
+
+    public fun swap_b_to_a<A,B>(vault: &mut Vault<A,B>, coinb:Coin<B>, ctx: &mut TxContext): Coin<A> {
+            let amount_out_A = coin::value(&coinb) * balance::value(&vault.coin_a) / balance::value(&vault.coin_b);
+            coin::put<B>(&mut vault.coin_b, coinb);
+            coin::take(&mut vault.coin_a, amount_out_A, ctx)
+    }
+
+    public fun get_flag<A,B>(vault: &Vault<A,B>, ctx: &TxContext) {
+        assert!(
+            balance::value(&vault.coin_a) == 0 && balance::value(&vault.coin_b) == 0, 123
+        );
+        event::emit(
+            Flag {
+                win: true,
+                sender: tx_context::sender(ctx)
+            }
+        );
+    }
+}

tests/move/swap/sources/TokenA.move

@@ -0,0 +1,40 @@
+module swap::ctfa {
+    use sui::coin::{Self, Coin, TreasuryCap};
+    use sui::tx_context::{Self, TxContext};
+    use sui::transfer;
+    use sui::object::{Self, UID};
+    use std::option;
+
+    friend swap::vault;
+
+    struct CTFA has drop {}
+
+    struct MintA<phantom CTFA> has key, store{
+        id: UID,
+        cap: TreasuryCap<CTFA>
+    }
+
+    fun init(witness: CTFA, ctx: &mut TxContext){
+        // Get a treasury cap for the coin and give it to the transaction sender
+        let (treasury_cap, metadata) = coin::create_currency<CTFA>(witness, 1, b"CTF", b"CTF", b"Token for move ctf", option::none(), ctx);
+        let mint = MintA<CTFA> {
+            id: object::new(ctx),
+            cap:treasury_cap
+        };
+        transfer::share_object(mint);
+        transfer::public_freeze_object(metadata);
+    }
+
+    public(friend) fun mint_for_vault<CTFA>(mint: MintA<CTFA>, ctx: &mut TxContext): Coin<CTFA> {
+        let coinb = coin::mint<CTFA>(&mut mint.cap, 100, ctx);
+        coin::mint_and_transfer(&mut mint.cap, 10, tx_context::sender(ctx), ctx);
+        let MintA<CTFA> {
+            id: ida,
+            cap: capa
+        } = mint;
+        object::delete(ida);
+        transfer::public_freeze_object(capa);
+        coinb
+    }
+
+}

tests/move/swap/sources/TokenB.move

@@ -0,0 +1,40 @@
+module swap::ctfb {
+    use sui::coin::{Self, Coin, TreasuryCap};
+    use sui::tx_context::{Self, TxContext};
+    use sui::transfer;
+    use sui::object::{Self, UID};
+    use std::option;
+
+    friend swap::vault;
+
+    struct CTFB has drop {}
+
+    struct MintB<phantom CTFB> has key, store {
+        id: UID,
+        cap: TreasuryCap<CTFB>
+    }
+
+    fun init(witness: CTFB, ctx: &mut TxContext) {
+        // Get a treasury cap for the coin and give it to the transaction sender
+        let (treasury_cap, metadata) = coin::create_currency<CTFB>(witness, 1, b"CTF", b"CTF", b"Token for move ctf", option::none(), ctx);
+        let mint = MintB<CTFB> {
+            id: object::new(ctx),
+            cap:treasury_cap
+        };
+        transfer::share_object(mint);
+        transfer::public_freeze_object(metadata);
+    }
+
+    public(friend) fun mint_for_vault<CTFB>(mint: MintB<CTFB>, ctx: &mut TxContext): Coin<CTFB> {
+        let coinb = coin::mint<CTFB>(&mut mint.cap, 100, ctx);
+        coin::mint_and_transfer(&mut mint.cap, 10, tx_context::sender(ctx), ctx);
+        let MintB<CTFB> {
+            id: idb,
+            cap: capb
+        } = mint;
+        object::delete(idb);
+        transfer::public_freeze_object(capb);
+        coinb
+    }
+
+}