Fix ft transfer call message check

Makefile.toml

@@ -174,6 +174,7 @@ category = "Test"
 command = "${CARGO}"
 args = [
     "test",
+    "test_ft_transfer_call_user_message",
     "--features",
     "${CARGO_FEATURES_TEST}",
     "--",

eth-connector-tests/src/connector.rs

@@ -311,6 +311,93 @@ async fn test_ft_transfer_call_without_message() -> anyhow::Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn test_ft_transfer_call_user_message() {
+    let contract = TestContract::new().await.unwrap();
+    contract.call_deposit_eth_to_near().await.unwrap();
+
+    let user_acc = contract.create_sub_account("eth_recipient").await.unwrap();
+    let receiver_id = contract.contract.id();
+    let balance = contract
+        .get_eth_on_near_balance(user_acc.id())
+        .await
+        .unwrap();
+    assert_eq!(balance.0, DEPOSITED_AMOUNT - DEPOSITED_FEE);
+
+    let balance = contract
+        .get_eth_on_near_balance(contract.contract.id())
+        .await
+        .unwrap();
+    assert_eq!(balance.0, DEPOSITED_FEE);
+
+    let transfer_amount: U128 = 50.into();
+    let memo: Option<String> = None;
+    let message = "";
+    // Send to Aurora contract with wrong message should failed
+    let res = user_acc
+        .call(contract.contract.id(), "ft_transfer_call")
+        .args_json((&receiver_id, transfer_amount, &memo, message))
+        .gas(DEFAULT_GAS)
+        .deposit(ONE_YOCTO)
+        .transact()
+        .await
+        .unwrap();
+    //println!("{:#?}", res);
+    assert!(res.is_success());
+
+    let balance = contract.get_eth_on_near_balance(receiver_id).await.unwrap();
+    assert_eq!(balance.0, DEPOSITED_FEE + transfer_amount.0);
+    let balance = contract
+        .get_eth_on_near_balance(user_acc.id())
+        .await
+        .unwrap();
+    assert_eq!(
+        balance.0,
+        DEPOSITED_AMOUNT - DEPOSITED_FEE - transfer_amount.0
+    );
+
+    let res = contract
+        .contract
+        .call("set_engine_account")
+        .args_json((&receiver_id,))
+        .gas(DEFAULT_GAS)
+        .transact()
+        .await
+        .unwrap();
+    assert!(res.is_success());
+
+    let res = contract
+        .contract
+        .call("get_engine_accounts")
+        .view()
+        .await
+        .unwrap()
+        .json::<Vec<AccountId>>()
+        .unwrap();
+    assert!(res.contains(receiver_id));
+
+    let res = user_acc
+        .call(contract.contract.id(), "ft_transfer_call")
+        .args_json((&receiver_id, transfer_amount, &memo, message))
+        .gas(DEFAULT_GAS)
+        .deposit(ONE_YOCTO)
+        .transact()
+        .await
+        .unwrap();
+    assert!(res.is_failure());
+    assert!(contract.check_error_message(res, "ERR_INVALID_ON_TRANSFER_MESSAGE_FORMAT"));
+    let balance = contract.get_eth_on_near_balance(receiver_id).await.unwrap();
+    assert_eq!(balance.0, DEPOSITED_FEE + transfer_amount.0);
+    let balance = contract
+        .get_eth_on_near_balance(user_acc.id())
+        .await
+        .unwrap();
+    assert_eq!(
+        balance.0,
+        DEPOSITED_AMOUNT - DEPOSITED_FEE - transfer_amount.0
+    );
+}
+
 #[tokio::test]
 async fn test_deposit_with_0x_prefix() -> anyhow::Result<()> {
     let contract = TestContract::new().await?;

eth-connector/src/connector.rs

@@ -85,3 +85,10 @@ pub trait EngineStorageManagement {
 
     fn engine_storage_unregister(&mut self, sender_id: AccountId, force: Option<bool>) -> bool;
 }
+
+#[ext_contract(ext_known_enine_accounts)]
+pub trait KnownEnginAccountsManagement {
+    fn set_engine_account(&mut self, engine_account: AccountId);
+
+    fn get_engine_accounts(&self) -> Vec<AccountId>;
+}

eth-connector/src/lib.rs

@@ -1,7 +1,7 @@
 use crate::admin_controlled::{AdminControlled, PausedMask, PAUSE_WITHDRAW, UNPAUSE_ALL};
 use crate::connector::{
     ConnectorDeposit, ConnectorFundsFinish, ConnectorWithdraw, EngineFungibleToken,
-    EngineStorageManagement, FungibleTokeStatistic,
+    EngineStorageManagement, FungibleTokeStatistic, KnownEnginAccountsManagement,
 };
 use crate::connector_impl::{
     EthConnector, FinishDepositCallArgs, TransferCallCallArgs, WithdrawResult,
@@ -55,6 +55,7 @@ pub struct EthConnectorContract {
     metadata: LazyOption<FungibleTokenMetadata>,
     used_proofs: LookupMap<String, bool>,
     accounts_counter: u64,
+    known_engine_accounts: Vec<AccountId>,
 }
 
 #[derive(BorshSerialize, BorshStorageKey)]
@@ -95,6 +96,57 @@ impl EthConnectorContract {
             self.ft.internal_register_account(account);
         }
     }
+
+    fn internal_ft_transfer_call(
+        &mut self,
+        sender_id: AccountId,
+        receiver_id: AccountId,
+        amount: U128,
+        memo: Option<String>,
+        msg: String,
+    ) -> PromiseOrValue<U128> {
+        assert_one_yocto();
+        self.register_if_not_exists(&receiver_id);
+
+        let amount: Balance = amount.into();
+        crate::log!(
+            "Transfer call from {} to {} amount {}",
+            sender_id,
+            receiver_id,
+            amount,
+        );
+
+        // Verify message data before `ft_on_transfer` call for Engine accounts
+        // to avoid verification panics inside `ft_on_transfer`.
+        // Allowed empty message if `receiver_id != known_engin_accounts`.
+        if self.known_engine_accounts.contains(&receiver_id) {
+            let _ = FtTransferMessageData::parse_on_transfer_message(&msg).sdk_unwrap();
+        }
+
+        // Special case, we do not fail if `sender_id = receiver_id`
+        // if `predecessor_account_id` call `ft_transfer_call` as receiver itself
+        // to call `ft_on_transfer`.
+        if sender_id != receiver_id {
+            // It's panic if: `sender_id == receiver_id`
+            self.ft
+                .internal_transfer(&sender_id, &receiver_id, amount, memo);
+        }
+
+        let receiver_gas = env::prepaid_gas()
+            .0
+            .checked_sub(GAS_FOR_FT_TRANSFER_CALL.0)
+            .unwrap_or_else(|| env::panic_str("Prepaid gas overflow"));
+        // Initiating receiver's call and the callback
+        ext_ft_receiver::ext(receiver_id.clone())
+            .with_static_gas(receiver_gas.into())
+            .ft_on_transfer(sender_id.clone(), amount.into(), msg)
+            .then(
+                ext_ft_resolver::ext(env::current_account_id())
+                    .with_static_gas(GAS_FOR_RESOLVE_TRANSFER)
+                    .ft_resolve_transfer(sender_id, receiver_id, amount.into()),
+            )
+            .into()
+    }
 }
 
 #[near_bindgen]
@@ -124,6 +176,7 @@ impl EthConnectorContract {
             metadata: LazyOption::new(StorageKey::Metadata, Some(&metadata)),
             used_proofs: LookupMap::new(StorageKey::Proof),
             accounts_counter: 0,
+            known_engine_accounts: vec![],
         };
         this.register_if_not_exists(&owner_id);
         this
@@ -162,47 +215,8 @@ impl FungibleTokenCore for EthConnectorContract {
         memo: Option<String>,
         msg: String,
     ) -> PromiseOrValue<U128> {
-        self.register_if_not_exists(&receiver_id);
-
         let sender_id = env::predecessor_account_id();
-        let amount: Balance = amount.into();
-        crate::log!(
-            "Transfer call from {} to {} amount {}",
-            sender_id,
-            receiver_id,
-            amount,
-        );
-
-        // Verify message data before `ft_on_transfer` call to avoid verification panics
-        // It's allowed empty message if `receiver_id =! current_account_id`
-        if sender_id == receiver_id {
-            let message_data = FtTransferMessageData::parse_on_transfer_message(&msg).sdk_unwrap();
-            // Check is transfer amount > fee
-            if message_data.fee.as_u128() >= amount {
-                panic_err("insufficient balance for fee");
-            }
-        }
-
-        // Special case for Aurora transfer itself - we shouldn't transfer
-        if sender_id != receiver_id {
-            self.ft
-                .internal_transfer(&sender_id, &receiver_id, amount, memo);
-        }
-
-        let receiver_gas = env::prepaid_gas()
-            .0
-            .checked_sub(GAS_FOR_FT_TRANSFER_CALL.0)
-            .unwrap_or_else(|| env::panic_str("Prepaid gas overflow"));
-        // Initiating receiver's call and the callback
-        ext_ft_receiver::ext(receiver_id.clone())
-            .with_static_gas(receiver_gas.into())
-            .ft_on_transfer(sender_id.clone(), amount.into(), msg)
-            .then(
-                ext_ft_resolver::ext(env::current_account_id())
-                    .with_static_gas(GAS_FOR_RESOLVE_TRANSFER)
-                    .ft_resolve_transfer(sender_id, receiver_id, amount.into()),
-            )
-            .into()
+        self.internal_ft_transfer_call(sender_id, receiver_id, amount, memo, msg)
     }
 
     fn ft_total_supply(&self) -> U128 {
@@ -247,46 +261,19 @@ impl EngineFungibleToken for EthConnectorContract {
         msg: String,
     ) -> PromiseOrValue<U128> {
         self.assert_access_right().sdk_unwrap();
-        self.register_if_not_exists(&receiver_id);
-
-        let amount: Balance = amount.into();
-        crate::log!(
-            "Transfer call from {} to {} amount {}",
-            sender_id,
-            receiver_id,
-            amount,
-        );
-
-        // Verify message data before `ft_on_transfer` call to avoid verification panics
-        // It's allowed empty message if `receiver_id =! current_account_id`
-        if sender_id == receiver_id {
-            let message_data = FtTransferMessageData::parse_on_transfer_message(&msg).sdk_unwrap();
-            // Check is transfer amount > fee
-            if message_data.fee.as_u128() >= amount {
-                panic_err("insufficient balance for fee");
-            }
-        }
+        self.internal_ft_transfer_call(sender_id, receiver_id, amount, memo, msg)
+    }
+}
 
-        // Special case for Aurora transfer itself - we shouldn't transfer
-        if sender_id != receiver_id {
-            self.ft
-                .internal_transfer(&sender_id, &receiver_id, amount, memo);
-        }
+/// Management for a known Engine accounts
+#[near_bindgen]
+impl KnownEnginAccountsManagement for EthConnectorContract {
+    fn set_engine_account(&mut self, engine_account: AccountId) {
+        self.known_engine_accounts.push(engine_account);
+    }
 
-        let receiver_gas = env::prepaid_gas()
-            .0
-            .checked_sub(GAS_FOR_FT_TRANSFER_CALL.0)
-            .unwrap_or_else(|| env::panic_str("Prepaid gas overflow"));
-        // Initiating receiver's call and the callback
-        ext_ft_receiver::ext(receiver_id.clone())
-            .with_static_gas(receiver_gas.into())
-            .ft_on_transfer(sender_id.clone(), amount.into(), msg)
-            .then(
-                ext_ft_resolver::ext(env::current_account_id())
-                    .with_static_gas(GAS_FOR_RESOLVE_TRANSFER)
-                    .ft_resolve_transfer(sender_id, receiver_id, amount.into()),
-            )
-            .into()
+    fn get_engine_accounts(&self) -> Vec<AccountId> {
+        self.known_engine_accounts.clone()
     }
 }