Add auth check to on_record callback

defer-stub/src/lib.rs

@@ -1,6 +1,8 @@
 use near_sdk::{
     borsh::{self, BorshDeserialize, BorshSerialize},
+    env,
     env::log_str,
+    ext_contract,
     json_types::U128,
     near_bindgen, AccountId, PanicOnDefault,
 };
@@ -17,6 +19,34 @@ impl Contract {
     }
 
     pub fn record_batch_for_hold(&mut self, amounts: Vec<(AccountId, U128)>) {
-        log_str(format!("Call record_batch_for_hold with {:?}", amounts).as_str());
+        log_str(&format!("Call record_batch_for_hold with {:?}", amounts));
     }
+
+    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

@@ -0,0 +1,66 @@
+#![cfg(test)]
+
+use integration_utils::{integration_contract::IntegrationContract, misc::ToNear};
+use near_sdk::json_types::U128;
+use serde_json::json;
+use sweat_model::FungibleTokenCoreIntegration;
+
+use crate::{
+    common::PanicFinder,
+    interface::common::ContractAccount,
+    prepare::{prepare_contract, IntegrationContext},
+};
+
+#[tokio::test]
+async fn test_call_on_record_in_callback() -> anyhow::Result<()> {
+    let mut context = prepare_contract().await?;
+
+    let alice = context.alice().await?;
+
+    let alice_balance_before_attack = context.ft_contract().ft_balance_of(alice.to_near()).await?;
+
+    let target_amount = U128(1_000_000);
+    let result = alice
+        .call(context.holding_contract().id(), "exploit_on_record")
+        .args_json(json!({
+            "ft_account_id": context.ft_contract().account(),
+            "amount": target_amount,
+        }))
+        .max_gas()
+        .transact()
+        .await?
+        .into_result()?;
+
+    assert!(result.has_panic("The operation can be only initiated by an oracle"));
+
+    let alice_balance_after_attack = context.ft_contract().ft_balance_of(alice.to_near()).await?;
+    assert_eq!(alice_balance_before_attack, alice_balance_after_attack);
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn test_call_on_record_directly() -> anyhow::Result<()> {
+    let mut context = prepare_contract().await?;
+
+    let alice = context.alice().await?;
+    let oracle = context.oracle().await?;
+
+    let intruder_id = alice.to_near();
+    let result = oracle
+        .call(context.ft_contract().contract().id(), "on_record")
+        .args_json(json!({
+            "receiver_id": intruder_id,
+            "amount": "1000000",
+            "fee_account_id": intruder_id,
+            "fee": "2000000",
+        }))
+        .max_gas()
+        .transact()
+        .await?
+        .into_result();
+
+    assert!(result.has_panic("Contract expected a result on the callback"));
+
+    Ok(())
+}

integration-tests/src/common.rs

@@ -0,0 +1,29 @@
+use near_workspaces::result::{ExecutionFailure, ExecutionResult, ExecutionSuccess};
+
+pub(crate) trait PanicFinder {
+    fn has_panic(&self, message: &str) -> bool;
+}
+
+impl PanicFinder for Result<ExecutionSuccess, ExecutionFailure> {
+    fn has_panic(&self, message: &str) -> bool {
+        match self {
+            Ok(ok) => ok.has_panic(message),
+            Err(err) => err.has_panic(message),
+        }
+    }
+}
+
+impl<T> PanicFinder for ExecutionResult<T> {
+    fn has_panic(&self, message: &str) -> bool {
+        self.outcomes()
+            .into_iter()
+            .map(|item| match item.clone().into_result() {
+                Ok(_) => None,
+                Err(err) => Some(err),
+            })
+            .any(|item| match item {
+                None => false,
+                Some(err) => format!("{err:?}").contains(message),
+            })
+    }
+}

integration-tests/src/lib.rs

@@ -5,6 +5,8 @@ use sweat_model::{FungibleTokenCoreIntegration, SweatApiIntegration, SweatDeferI
 
 use crate::prepare::{prepare_contract, IntegrationContext};
 
+mod callback_attack;
+mod common;
 mod defer;
 mod formula;
 mod interface;

sweat/src/defer.rs

@@ -1,4 +1,4 @@
-use near_sdk::{serde_json::json, Gas, Promise};
+use near_sdk::{env::panic_str, serde_json::json, Gas, Promise};
 use sweat_model::SweatDefer;
 
 use crate::*;
@@ -57,24 +57,30 @@ pub trait FungibleTokenTransferCallback {
 #[near_bindgen]
 impl FungibleTokenTransferCallback for Contract {
     fn on_record(&mut self, receiver_id: AccountId, amount: U128, fee_account_id: AccountId, fee: U128) {
-        if is_promise_success() {
-            let mut events: Vec<FtMint> = Vec::with_capacity(2);
+        if !self.oracles.contains(&env::signer_account_id()) {
+            panic_str("The operation can be only initiated by an oracle");
+        }
 
-            internal_deposit(&mut self.token, &fee_account_id, fee.0);
-            events.push(FtMint {
-                owner_id: &fee_account_id,
-                amount: &fee,
-                memo: None,
-            });
+        if !is_promise_success() {
+            panic_str("Failed to record data in holding account");
+        }
 
-            internal_deposit(&mut self.token, &receiver_id, amount.0);
-            events.push(FtMint {
-                owner_id: &receiver_id,
-                amount: &amount,
-                memo: None,
-            });
+        let mut events: Vec<FtMint> = Vec::with_capacity(2);
 
-            FtMint::emit_many(&events);
-        }
+        internal_deposit(&mut self.token, &fee_account_id, fee.0);
+        events.push(FtMint {
+            owner_id: &fee_account_id,
+            amount: &fee,
+            memo: None,
+        });
+
+        internal_deposit(&mut self.token, &receiver_id, amount.0);
+        events.push(FtMint {
+            owner_id: &receiver_id,
+            amount: &amount,
+            memo: None,
+        });
+
+        FtMint::emit_many(&events);
     }
 }