fixed exploit test

Cargo.toml

@@ -4,6 +4,7 @@ members = [
     "sweat",
     "sweat-integration",
     "integration-tests",
+    "exploit-stub"
 ]
 
 resolver = "2"

exploit-stub/Cargo.toml

@@ -0,0 +1,12 @@
+[package]
+name = "exploit-stub"
+version = "0.1.0"
+authors = []
+edition = "2021"
+publish = false
+
+[lib]
+crate-type = ["cdylib"]
+
+[dependencies]
+near-sdk = { workspace = true }

exploit-stub/src/lib.rs

@@ -0,0 +1,55 @@
+use near_sdk::{
+    borsh::{self, BorshDeserialize, BorshSerialize},
+    env,
+    env::log_str,
+    ext_contract,
+    json_types::U128,
+    near_bindgen, AccountId, PanicOnDefault,
+};
+
+#[near_bindgen]
+#[derive(BorshSerialize, BorshDeserialize, PanicOnDefault)]
+struct Contract {}
+
+#[near_bindgen]
+#[allow(dead_code)]
+impl Contract {
+    #[init]
+    pub fn new() -> Self {
+        Self {}
+    }
+
+    #[allow(clippy::unused_self)]
+    pub fn record_batch_for_hold(&mut self, amounts: Vec<(AccountId, U128)>) {
+        log_str(&format!("Call record_batch_for_hold with {amounts:?}"));
+    }
+
+    #[allow(clippy::unused_self)]
+    pub fn exploit_on_record(&mut self, ft_account_id: AccountId, amount: U128) {
+        log_str(&format!(
+            "Try to call on_record in callback, ft account = {ft_account_id}"
+        ));
+
+        let intruder_id = env::predecessor_account_id();
+        ext_self::ext(env::current_account_id())
+            .some_function()
+            .then(ext_token::ext(ft_account_id).on_record(intruder_id.clone(), amount, intruder_id, U128(0)));
+    }
+}
+
+#[ext_contract(ext_self)]
+pub trait Callback {
+    fn some_function(&mut self);
+}
+
+#[near_bindgen]
+impl Callback for Contract {
+    fn some_function(&mut self) {
+        log_str("Call some_function in stub contract");
+    }
+}
+
+#[ext_contract(ext_token)]
+pub trait FungibleTokenTransferCallback {
+    fn on_record(&mut self, receiver_id: AccountId, amount: U128, fee_account_id: AccountId, fee: U128);
+}

integration-tests/src/callback_attack.rs

@@ -21,7 +21,7 @@ async fn test_call_on_record_in_callback() -> anyhow::Result<()> {
 
     let target_amount = U128(1_000_000);
     let result = alice
-        .call(context.claim_contract().id(), "exploit_on_record")
+        .call(context.stub_contract().id(), "exploit_on_record")
         .args_json(json!({
             "ft_account_id": ft_contract_id,
             "amount": target_amount,

integration-tests/src/prepare.rs

@@ -7,6 +7,7 @@ use sweat_integration::{SweatFt, FT_CONTRACT};
 use sweat_model::{StorageManagementIntegration, SweatApiIntegration};
 
 const CLAIM_CONTRACT: &str = "sweat_claim";
+const HOLDING_STUB_CONTRACT: &str = "exploit_stub";
 
 pub type Context = integration_utils::context::Context<near_workspaces::network::Sandbox>;
 
@@ -16,9 +17,10 @@ pub trait IntegrationContext {
     async fn alice(&mut self) -> Result<Account>;
     async fn bob(&mut self) -> Result<Account>;
     async fn long_account_name(&mut self) -> Result<Account>;
-    fn ft_contract(&self) -> SweatFt;
 
+    fn ft_contract(&self) -> SweatFt;
     fn claim_contract(&self) -> &Contract;
+    fn stub_contract(&self) -> &Contract;
 }
 
 #[async_trait]
@@ -46,10 +48,18 @@ impl IntegrationContext for Context {
     fn claim_contract(&self) -> &Contract {
         &self.contracts[CLAIM_CONTRACT]
     }
+
+    fn stub_contract(&self) -> &Contract {
+        &self.contracts[HOLDING_STUB_CONTRACT]
+    }
 }
 
 pub async fn prepare_contract() -> Result<Context> {
-    let mut context = Context::new(&[FT_CONTRACT, CLAIM_CONTRACT], "build-integration".into()).await?;
+    let mut context = Context::new(
+        &[FT_CONTRACT, CLAIM_CONTRACT, HOLDING_STUB_CONTRACT],
+        "build-integration".into(),
+    )
+    .await?;
     let oracle = context.oracle().await?;
     let alice = context.alice().await?;
     let long = context.long_account_name().await?;
@@ -81,7 +91,7 @@ pub async fn prepare_contract() -> Result<Context> {
 
     context.ft_contract().add_oracle(&oracle.to_near()).call().await?;
 
-    let holding_contract_init_result = context
+    let claim_contract_result = context
         .claim_contract()
         .call("init")
         .args_json(json!({ "token_account_id": token_account_id }))
@@ -90,7 +100,17 @@ pub async fn prepare_contract() -> Result<Context> {
         .await?
         .into_result()?;
 
-    println!("Initialized holding contract: {:?}", holding_contract_init_result);
+    println!("Initialized claim contract: {:?}", claim_contract_result);
+
+    let exploit_stup_contract_result = context
+        .stub_contract()
+        .call("new")
+        .max_gas()
+        .transact()
+        .await?
+        .into_result()?;
+
+    println!("Initialized exploit stub contract: {:?}", exploit_stup_contract_result);
 
     context
         .claim_contract()

scripts/build-stub.sh

@@ -4,6 +4,6 @@ set -eox pipefail
 echo ">> Building stub contract"
 
 rustup target add wasm32-unknown-unknown
-cargo build -p defer-stub --target wasm32-unknown-unknown --profile=contract
+cargo build -p exploit-stub --target wasm32-unknown-unknown --profile=contract
 
-cp ./target/wasm32-unknown-unknown/contract/defer_stub.wasm res/defer_stub.wasm
+cp ./target/wasm32-unknown-unknown/contract/exploit_stub.wasm res/exploit_stub.wasm