Aurora-to-eth Fast Bridge: combine unlock with withdraw

aurora/contracts/src/AuroraErc20FastBridge.sol

@@ -27,7 +27,7 @@ contract AuroraErc20FastBridge is Initializable, UUPSUpgradeable, AccessControlU
     uint64 constant BASE_NEAR_GAS = 10_000_000_000_000;
     uint64 constant WITHDRAW_NEAR_GAS = 50_000_000_000_000;
     uint64 constant INIT_TRANSFER_NEAR_GAS = 100_000_000_000_000;
-    uint64 constant UNLOCK_NEAR_GAS = 150_000_000_000_000;
+    uint64 constant UNLOCK_NEAR_GAS = 160_000_000_000_000;
 
     uint128 constant NEAR_STORAGE_DEPOSIT = 12_500_000_000_000_000_000_000;
 
@@ -252,7 +252,7 @@ contract AuroraErc20FastBridge is Initializable, UUPSUpgradeable, AccessControlU
 
         PromiseCreateArgs memory callUnlock = near.call(
             bridgeAddressOnNear,
-            "unlock",
+            "unlock_and_withdraw",
             args,
             NO_DEPOSIT,
             UNLOCK_NEAR_GAS
@@ -282,34 +282,6 @@ contract AuroraErc20FastBridge is Initializable, UUPSUpgradeable, AccessControlU
         );
     }
 
-    function withdrawFromNear(string calldata tokenId, uint128 amount) external whenNotPaused {
-        require(near.wNEAR.balanceOf(address(this)) >= ONE_YOCTO, "Not enough wNEAR balance");
-        bytes memory args = bytes(
-            string.concat('{"token_id": "', tokenId, '", "amount": "', Strings.toString(amount), '"}')
-        );
-        PromiseCreateArgs memory callWithdraw = _callWithoutTransferWNear(
-            near,
-            bridgeAddressOnNear,
-            "withdraw",
-            args,
-            ONE_YOCTO,
-            WITHDRAW_NEAR_GAS
-        );
-        bytes memory callbackArg = abi.encodeWithSelector(this.withdrawFromNearCallback.selector, tokenId, amount);
-        PromiseCreateArgs memory callback = near.auroraCall(address(this), callbackArg, NO_DEPOSIT, BASE_NEAR_GAS);
-
-        callWithdraw.then(callback).transact();
-    }
-
-    function withdrawFromNearCallback(string calldata tokenId, uint128 amount) external onlyRole(CALLBACK_ROLE) {
-        require(
-            AuroraSdk.promiseResult(0).status == PromiseResultStatus.Successful,
-            "ERROR: The `Withdraw From Near` XCC is fail"
-        );
-
-        emit WithdrawFromNear(tokenId, amount);
-    }
-
     function withdraw(string calldata token) external whenNotPaused {
         require(near.wNEAR.balanceOf(address(this)) >= ONE_YOCTO, "Not enough wNEAR balance");
         uint128 signerBalance = balance[token][msg.sender];

aurora/integration-tests/src/lib.rs

@@ -224,20 +224,6 @@ mod tests {
             self.call_aurora_contract(contract_args).await;
         }
 
-        pub async fn withdraw_from_near(&self) {
-            let contract_args = self
-                .aurora_fast_bridge_contract
-                .create_call_method_bytes_with_args(
-                    "withdrawFromNear",
-                    &[
-                        ethabi::Token::String(self.mock_token.id().to_string()),
-                        ethabi::Token::Uint(U256::from(TRANSFER_TOKENS_AMOUNT)),
-                    ],
-                );
-
-            self.call_aurora_contract(contract_args).await;
-        }
-
         pub async fn call_aurora_contract(&self, contract_args: Vec<u8>) {
             call_aurora_contract(
                 self.aurora_fast_bridge_contract.address,
@@ -346,7 +332,6 @@ mod tests {
         infra
             .assert_user_balance_in_fast_bridge_on_aurora(None, TRANSFER_TOKENS_AMOUNT)
             .await;
-        infra.withdraw_from_near().await;
         infra
             .assert_user_balance_in_fast_bridge_on_aurora(None, TRANSFER_TOKENS_AMOUNT)
             .await;

near/contracts/bridge/src/lib.rs

@@ -96,6 +96,20 @@ trait FastBridgeInterface {
         #[serializer(borsh)] sender_id: AccountId,
         #[serializer(borsh)] update_balance: Option<UpdateBalance>,
     ) -> PromiseOrValue<U128>;
+
+    fn unlock_and_withdraw_callback(
+        &mut self,
+        #[callback]
+        #[serializer(borsh)]
+        transfer_data: TransferMessage,
+        #[serializer(borsh)] sender_id: AccountId,
+    ) -> Promise;
+
+    fn unlock_and_withdraw_return_transfer_msg(
+        &self,
+        #[callback] withdraw_amount: U128,
+        transfer_data: TransferMessage,
+    ) -> TransferMessage;
 }
 
 #[derive(
@@ -375,6 +389,48 @@ impl FastBridge {
         U128::from(0)
     }
 
+    #[pause(except(roles(Role::UnrestrictedUnlock)))]
+    pub fn unlock_and_withdraw(
+        &self,
+        nonce: U128,
+        proof: near_sdk::json_types::Base64VecU8,
+    ) -> Promise {
+        self.unlock(nonce, proof).then(
+            ext_self::ext(current_account_id())
+                .with_static_gas(utils::tera_gas(25))
+                .with_attached_deposit(utils::NO_DEPOSIT)
+                .unlock_and_withdraw_callback(env::predecessor_account_id()),
+        )
+    }
+
+    #[private]
+    pub fn unlock_and_withdraw_callback(
+        &mut self,
+        #[callback]
+        #[serializer(borsh)]
+        transfer_data: TransferMessage,
+        #[serializer(borsh)] sender_id: AccountId,
+    ) -> Promise {
+        self.withdraw_internal(transfer_data.transfer.token_near.clone(), None, sender_id)
+            .then(
+                ext_self::ext(current_account_id())
+                    .with_static_gas(utils::tera_gas(1))
+                    .with_attached_deposit(utils::NO_DEPOSIT)
+                    .unlock_and_withdraw_return_transfer_msg(transfer_data),
+            )
+    }
+
+    #[private]
+    #[result_serializer(borsh)]
+    pub fn unlock_and_withdraw_return_transfer_msg(
+        &self,
+        #[callback] withdraw_amount: U128,
+        transfer_data: TransferMessage,
+    ) -> TransferMessage {
+        require!(withdraw_amount.0 > 0, "Withdraw failed");
+        transfer_data
+    }
+
     /// Unlocks the transfer with the given `nonce`, using the provided `proof` of the non-existence
     /// of the transfer on Ethereum. The unlock could be possible only if the transfer on Ethereum
     /// didn't happen and its validity time is already expired.
@@ -421,7 +477,7 @@ impl FastBridge {
             )
             .then(
                 ext_self::ext(current_account_id())
-                    .with_static_gas(utils::tera_gas(50))
+                    .with_static_gas(utils::tera_gas(10))
                     .with_attached_deposit(utils::NO_DEPOSIT)
                     .unlock_callback(nonce, env::predecessor_account_id()),
             )
@@ -793,8 +849,17 @@ impl FastBridge {
     /// * The caller does not have any balance.
     #[payable]
     #[pause(except(roles(Role::UnrestrictedWithdraw)))]
-    pub fn withdraw(&mut self, token_id: AccountId, amount: Option<U128>) -> PromiseOrValue<U128> {
+    pub fn withdraw(&mut self, token_id: AccountId, amount: Option<U128>) -> Promise {
         let recipient_id = env::predecessor_account_id();
+        self.withdraw_internal(token_id, amount, recipient_id)
+    }
+
+    fn withdraw_internal(
+        &mut self,
+        token_id: AccountId,
+        amount: Option<U128>,
+        recipient_id: AccountId,
+    ) -> Promise {
         let user_balance = self.get_user_balance(&recipient_id, &token_id);
         let amount = amount.unwrap_or(user_balance);
 
@@ -820,7 +885,6 @@ impl FastBridge {
                     .with_attached_deposit(utils::NO_DEPOSIT)
                     .withdraw_callback(token_id, amount, recipient_id),
             )
-            .into()
     }
 
     /// This function finalizes the execution flow of the `withdraw()` function. This private function is called after