fix: fast transfers state migration

near/omni-bridge/src/lib.rs

@@ -37,6 +37,7 @@ use storage::{
 
 mod errors;
 mod storage;
+mod migrate;
 
 #[cfg(test)]
 mod tests;

near/omni-bridge/src/migrate.rs

@@ -0,0 +1,50 @@
+use borsh::{BorshDeserialize, BorshSerialize};
+use near_contract_standards::storage_management::StorageBalance;
+use near_sdk::{collections::{LookupMap, LookupSet}, env, near, AccountId, PanicOnDefault};
+use omni_types::{ChainKind, Nonce, OmniAddress, TransferId};
+use crate::{storage::{Decimals, TransferMessageStorage}, Contract, ContractExt, StorageKey};
+
+#[derive(BorshDeserialize, BorshSerialize, PanicOnDefault)]
+pub struct OldState {
+    pub prover_account: AccountId,
+    pub factories: LookupMap<ChainKind, OmniAddress>,
+    pub pending_transfers: LookupMap<TransferId, TransferMessageStorage>,
+    pub finalised_transfers: LookupSet<TransferId>,
+    pub token_id_to_address: LookupMap<(ChainKind, AccountId), OmniAddress>,
+    pub token_address_to_id: LookupMap<OmniAddress, AccountId>,
+    pub token_decimals: LookupMap<OmniAddress, Decimals>,
+    pub deployed_tokens: LookupSet<AccountId>,
+    pub token_deployer_accounts: LookupMap<ChainKind, AccountId>,
+    pub mpc_signer: AccountId,
+    pub current_origin_nonce: Nonce,
+    pub destination_nonces: LookupMap<ChainKind, Nonce>,
+    pub accounts_balances: LookupMap<AccountId, StorageBalance>,
+    pub wnear_account_id: AccountId,
+}
+
+#[near]
+impl Contract {
+    #[private]
+    #[init(ignore_state)]
+    pub fn migrate() -> Self {
+        let old_state: OldState = env::state_read().expect("failed");
+
+        Self {
+            prover_account: old_state.prover_account,
+            factories: old_state.factories,
+            pending_transfers: old_state.pending_transfers,
+            finalised_transfers: old_state.finalised_transfers,
+            fast_transfers: LookupMap::new(StorageKey::FastTransfers),
+            token_id_to_address: old_state.token_id_to_address,
+            token_address_to_id: old_state.token_address_to_id,
+            token_decimals: old_state.token_decimals,
+            deployed_tokens: old_state.deployed_tokens,
+            token_deployer_accounts: old_state.token_deployer_accounts,
+            mpc_signer: old_state.mpc_signer,
+            current_origin_nonce: old_state.current_origin_nonce,
+            destination_nonces: old_state.destination_nonces,
+            accounts_balances: old_state.accounts_balances,
+            wnear_account_id: old_state.wnear_account_id,
+        }
+    }
+}

near/omni-bridge/src/storage.rs

@@ -1,7 +1,7 @@
 use near_contract_standards::storage_management::{StorageBalance, StorageBalanceBounds};
 use near_sdk::{assert_one_yocto, borsh, near};
 use near_sdk::{env, near_bindgen, AccountId, NearToken};
-use omni_types::{FastTransferStatus, TransferId};
+use omni_types::{FastTransferStatus, TransferId, TransferMessageV0};
 
 use crate::{
     require, ChainKind, Contract, ContractExt, Fee, OmniAddress, Promise, SdkExpect,
@@ -11,6 +11,13 @@ use crate::{
 pub const BRIDGE_TOKEN_INIT_BALANCE: NearToken = NearToken::from_near(3);
 pub const NEP141_DEPOSIT: NearToken = NearToken::from_yoctonear(1_250_000_000_000_000_000_000);
 
+#[near(serializers=[borsh, json])]
+#[derive(Debug, Clone)]
+pub struct TransferMessageStorageValueV0 {
+    pub message: TransferMessageV0,
+    pub owner: AccountId,
+}
+
 #[near(serializers=[borsh, json])]
 #[derive(Debug, Clone)]
 pub struct TransferMessageStorageValue {
@@ -22,21 +29,36 @@ pub struct TransferMessageStorageValue {
 #[near(serializers=[borsh, json])]
 #[derive(Debug, Clone)]
 pub enum TransferMessageStorage {
-    V0(TransferMessageStorageValue),
+    V0(TransferMessageStorageValueV0),
+    V1(TransferMessageStorageValue),
 }
 
 impl TransferMessageStorage {
     pub fn into_main(self) -> TransferMessageStorageValue {
         match self {
-            TransferMessageStorage::V0(m) => m,
+            TransferMessageStorage::V0(m) => TransferMessageStorageValue {
+                message: TransferMessage {
+                    origin_nonce: m.message.origin_nonce,
+                    token: m.message.token,
+                    amount: m.message.amount,
+                    recipient: m.message.recipient,
+                    fee: m.message.fee,
+                    sender: m.message.sender,
+                    msg: m.message.msg,
+                    destination_nonce: m.message.destination_nonce,
+                    origin_transfer_id: None,
+                },
+                owner: m.owner,
+            },
+            TransferMessageStorage::V1(m) => m,
         }
     }
 
     pub fn encode_borsh(
         message: TransferMessage,
         owner: AccountId,
     ) -> Result<Vec<u8>, std::io::Error> {
-        borsh::to_vec(&TransferMessageStorage::V0(TransferMessageStorageValue {
+        borsh::to_vec(&TransferMessageStorage::V1(TransferMessageStorageValue {
             message,
             owner,
         }))
@@ -174,7 +196,7 @@ impl Contract {
             .sdk_expect("ERR_CAST");
 
         let value_len: u64 =
-            borsh::to_vec(&TransferMessageStorage::V0(TransferMessageStorageValue {
+            borsh::to_vec(&TransferMessageStorage::V1(TransferMessageStorageValue {
                 message: TransferMessage {
                     origin_nonce: 0,
                     token: OmniAddress::Near(max_account_id.clone()),

near/omni-tests/src/init_transfer.rs

@@ -22,6 +22,8 @@ mod tests {
     const EXPECTED_RELAYER_GAS_COST: NearToken =
         NearToken::from_yoctonear(1_500_000_000_000_000_000_000);
 
+    const PREV_LOCKER_WASM_FILEPATH: &str = "src/data/omni_bridge-0_2_6.wasm";
+
     struct TestEnv {
         worker: near_workspaces::Worker<near_workspaces::network::Sandbox>,
         token_contract: near_workspaces::Contract,
@@ -177,6 +179,53 @@ mod tests {
         }
     }
 
+    async fn init_transfer_legacy(
+        env: &TestEnv,
+        transfer_amount: u128,
+        init_transfer_msg: InitTransferMsg,
+    ) -> anyhow::Result<TransferMessage> {
+        let storage_deposit_amount = get_balance_required_for_account(
+            &env.locker_contract,
+            &env.sender_account,
+            &init_transfer_msg,
+            None,
+        )
+        .await?;
+
+        // Storage deposit
+        env.sender_account
+            .call(env.locker_contract.id(), "storage_deposit")
+            .args_json(json!({
+                "account_id": env.sender_account.id(),
+            }))
+            .deposit(storage_deposit_amount)
+            .max_gas()
+            .transact()
+            .await?
+            .into_result()?;
+
+        // Initiate the transfer
+        let transfer_result = env
+            .sender_account
+            .call(env.token_contract.id(), "ft_transfer_call")
+            .args_json(json!({
+                "receiver_id": env.locker_contract.id(),
+                "amount": U128(transfer_amount),
+                "memo": None::<String>,
+                "msg": serde_json::to_string(&init_transfer_msg)?,
+            }))
+            .deposit(NearToken::from_yoctonear(1))
+            .max_gas()
+            .transact()
+            .await?
+            .into_result()?;
+
+        // Ensure the transfer event is emitted
+        let transfer_message = get_transfer_message_from_event(&transfer_result)?;
+
+        Ok(transfer_message)
+    }
+
     async fn init_transfer_flow_on_near(
         env: &TestEnv,
         transfer_amount: u128,
@@ -860,4 +909,71 @@ mod tests {
         .await
         .unwrap();
     }
+
+    #[rstest]
+    #[tokio::test]
+    async fn test_migrate(
+        mock_token_wasm: Vec<u8>,
+        mock_prover_wasm: Vec<u8>,
+        locker_wasm: Vec<u8>,
+    ) -> anyhow::Result<()> {
+        let sender_balance_token = 1_000_000;
+        let transfer_amount = 5000;
+        let init_transfer_msg = InitTransferMsg {
+            native_token_fee: U128(0),
+            fee: U128(0),
+            recipient: eth_eoa_address(),
+        };
+
+        let prev_locker_wasm = std::fs::read(PREV_LOCKER_WASM_FILEPATH).unwrap();
+        let env = TestEnv::new(
+            sender_balance_token,
+            mock_token_wasm,
+            mock_prover_wasm,
+            prev_locker_wasm,
+        )
+        .await?;
+
+        let transfer_message = init_transfer_legacy(
+            &env,
+            transfer_amount,
+            init_transfer_msg.clone(),
+        )
+        .await?;
+
+        let res = env.locker_contract
+            .as_account()
+            .deploy(&locker_wasm)
+            .await
+            .unwrap();
+        assert!(res.is_success(), "Failed to upgrade locker");
+
+        let res = env.locker_contract.call("migrate").args_json(json!({}))
+            .max_gas()
+            .transact()
+            .await?;
+        assert!(res.is_success(), "Migration didn't succeed");
+
+        let transfer = env.locker_contract.call("get_transfer_message").args_json(json!({
+            "transfer_id": TransferId {
+                origin_chain: ChainKind::Near,
+                origin_nonce: transfer_message.origin_nonce,
+            },
+        }))
+        .max_gas()
+        .transact()
+        .await?;
+
+        let migrated_transfer = transfer.json::<TransferMessage>()?;
+        assert_eq!(migrated_transfer.origin_transfer_id, None);
+        assert_eq!(migrated_transfer.origin_nonce, transfer_message.origin_nonce);
+        assert_eq!(migrated_transfer.recipient, transfer_message.recipient);
+        assert_eq!(migrated_transfer.token, transfer_message.token);
+        assert_eq!(migrated_transfer.amount, transfer_message.amount);
+        assert_eq!(migrated_transfer.fee, transfer_message.fee);
+        assert_eq!(migrated_transfer.sender, transfer_message.sender);
+        assert_eq!(migrated_transfer.destination_nonce, transfer_message.destination_nonce);
+
+        Ok(())
+    }
 }

near/omni-types/src/lib.rs

@@ -438,6 +438,19 @@ pub struct TransferId {
     pub origin_nonce: Nonce,
 }
 
+#[near(serializers=[borsh, json])]
+#[derive(Debug, Clone)]
+pub struct TransferMessageV0 {
+    pub origin_nonce: Nonce,
+    pub token: OmniAddress,
+    pub amount: U128,
+    pub recipient: OmniAddress,
+    pub fee: Fee,
+    pub sender: OmniAddress,
+    pub msg: String,
+    pub destination_nonce: Nonce
+}
+
 #[near(serializers=[borsh, json])]
 #[derive(Debug, Clone)]
 pub struct TransferMessage {