fix: pre-cw2 migrate validations (#117)

## Description

<!-- Brief description of changes -->

## Checklist

- [x] I have updated the
[docs/SPEC.md](https://github.com/babylonlabs-io/rollup-bsn-contracts/blob/main/docs/SPEC.md)
file if this change affects the specification
- [x] I have updated the schema by running `cargo gen-schema`

Author: gusin13

CHANGELOG.md

@@ -39,7 +39,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 
 ### Improvements
 
-* [#114](https://github.com/babylonlabs-io/rollup-bsn-contracts/pull/114) feat: migration setup
+* [#114](https://github.com/babylonlabs-io/rollup-bsn-contracts/pull/114) feat:
+  migration setup
+* [#117](https://github.com/babylonlabs-io/rollup-bsn-contracts/pull/117) fix:
+  pre-cw2 migrate validations
 
 
 ## v1.0.0-rc.0

contracts/finality/schema/finality.json

@@ -647,17 +647,8 @@
   "migrate": {
     "$schema": "http://json-schema.org/draft-07/schema#",
     "title": "MigrateMsg",
-    "description": "Migration message for contract upgrades. This can be extended in the future to include migration-specific parameters.",
+    "description": "Migration message for contract upgrades. Empty for non-state-breaking migrations. Can be extended in the future to include migration-specific parameters for state transformations.",
     "type": "object",
-    "properties": {
-      "version": {
-        "description": "Optional version string for tracking migration",
-        "type": [
-          "string",
-          "null"
-        ]
-      }
-    },
     "additionalProperties": false
   },
   "sudo": null,

contracts/finality/schema/raw/migrate.json

@@ -1,16 +1,7 @@
 {
   "$schema": "http://json-schema.org/draft-07/schema#",
   "title": "MigrateMsg",
-  "description": "Migration message for contract upgrades. This can be extended in the future to include migration-specific parameters.",
+  "description": "Migration message for contract upgrades. Empty for non-state-breaking migrations. Can be extended in the future to include migration-specific parameters for state transformations.",
   "type": "object",
-  "properties": {
-    "version": {
-      "description": "Optional version string for tracking migration",
-      "type": [
-        "string",
-        "null"
-      ]
-    }
-  },
   "additionalProperties": false
 }

contracts/finality/src/contract.rs

@@ -24,6 +24,8 @@ use crate::state::public_randomness::{
 // Contract metadata for version tracking
 const CONTRACT_NAME: &str = "rollup-bsn/finality";
 const CONTRACT_VERSION: &str = env!("CARGO_PKG_VERSION");
+/// Placeholder version for contracts deployed before cw2 version tracking was implemented
+const LEGACY_VERSION: &str = "pre-cw2";
 
 pub fn instantiate(
     mut deps: DepsMut<BabylonQuery>,
@@ -84,25 +86,42 @@ pub fn instantiate(
 }
 
 /// Handle contract migration.
+///
 /// This function is called when the contract is migrated to a new version.
+/// It supports both:
+/// - Migrating from contracts with existing cw2 version info
+/// - Migrating from legacy contracts deployed without cw2 version tracking
+///
 /// For non-state-breaking migrations, this updates the contract version and logs the migration.
+/// State-breaking migrations would require additional logic in this function.
 pub fn migrate(
     deps: DepsMut<BabylonQuery>,
     _env: Env,
     _msg: MigrateMsg,
 ) -> Result<Response<BabylonMsg>, ContractError> {
-    // Get the current version stored in the contract
-    let prev_version = cw2::get_contract_version(deps.storage)?;
-
-    // Validate that this is the expected contract
-    if prev_version.contract != CONTRACT_NAME {
-        return Err(ContractError::InvalidContractName {
-            expected: CONTRACT_NAME.to_string(),
-            actual: prev_version.contract,
-        });
-    }
+    // Attempt to get the current version stored in the contract
+    let prev_version = match cw2::get_contract_version(deps.storage) {
+        Ok(version) => {
+            // Contract has version info - validate it's the expected contract
+            if version.contract != CONTRACT_NAME {
+                return Err(ContractError::InvalidContractName {
+                    expected: CONTRACT_NAME.to_string(),
+                    actual: version.contract,
+                });
+            }
+            version
+        }
+        Err(_) => {
+            // No version info exists - this is a legacy contract
+            // Create a placeholder version for logging purposes
+            cw2::ContractVersion {
+                contract: CONTRACT_NAME.to_string(),
+                version: LEGACY_VERSION.to_string(),
+            }
+        }
+    };
 
-    // Update to the new version
+    // Update to the new version (this creates version info if it didn't exist)
     cw2::set_contract_version(deps.storage, CONTRACT_NAME, CONTRACT_VERSION)?;
 
     Ok(Response::new()
@@ -1638,6 +1657,7 @@ pub(crate) mod tests {
 
     #[test]
     fn test_migrate_basic() {
+        // Test 1: Migration from versioned contract
         let mut deps = mock_deps_babylon();
 
         // Simulate an older version being deployed initially (realistic migration scenario)
@@ -1655,12 +1675,36 @@ pub(crate) mod tests {
         assert_eq!(res.attributes[1].key, "from_version");
         assert_eq!(res.attributes[1].value, old_version); // Should be "0.9.0"
         assert_eq!(res.attributes[2].key, "to_version");
-        assert_eq!(res.attributes[2].value, CONTRACT_VERSION); // Should be "1.0.0-rc.0"
+        assert_eq!(res.attributes[2].value, CONTRACT_VERSION);
 
         // Verify the version was actually updated in storage
         let stored_version = cw2::get_contract_version(deps.as_ref().storage).unwrap();
         assert_eq!(stored_version.contract, CONTRACT_NAME);
         assert_eq!(stored_version.version, CONTRACT_VERSION);
+
+        // Test 2: Migration from legacy contract (no version info)
+        let mut legacy_deps = mock_deps_babylon();
+
+        // Don't set any version info - simulates a legacy contract deployed without cw2
+        // This is the real-world scenario for testnet contracts
+
+        // Test migration from legacy contract
+        let migrate_msg = MigrateMsg {};
+        let res = migrate(legacy_deps.as_mut(), mock_env(), migrate_msg).unwrap();
+
+        // Check that migration response has correct attributes for legacy migration
+        assert_eq!(res.attributes.len(), 3);
+        assert_eq!(res.attributes[0].key, "action");
+        assert_eq!(res.attributes[0].value, "migrate");
+        assert_eq!(res.attributes[1].key, "from_version");
+        assert_eq!(res.attributes[1].value, LEGACY_VERSION); // Should be "pre-cw2"
+        assert_eq!(res.attributes[2].key, "to_version");
+        assert_eq!(res.attributes[2].value, CONTRACT_VERSION);
+
+        // Verify the version was set in storage (first time)
+        let stored_version = cw2::get_contract_version(legacy_deps.as_ref().storage).unwrap();
+        assert_eq!(stored_version.contract, CONTRACT_NAME);
+        assert_eq!(stored_version.version, CONTRACT_VERSION);
     }
 
     #[test]

docs/SPEC.md

@@ -477,24 +477,38 @@ in the contract's state and managed by the migration function.
 ```rust
 /// Handle contract migration.
 /// This function is called when the contract is migrated to a new version.
+/// It supports both:
+/// - Migrating from contracts with existing cw2 version info
+/// - Migrating from legacy contracts deployed without cw2 version tracking
 /// For non-state-breaking migrations, this updates the contract version and logs the migration.
 pub fn migrate(
     deps: DepsMut<BabylonQuery>,
     _env: Env,
     _msg: MigrateMsg,
 ) -> Result<Response<BabylonMsg>, ContractError> {
-    // Get the current version stored in the contract
-    let prev_version = cw2::get_contract_version(deps.storage)?;
-
-    // Validate that this is the expected contract
-    if prev_version.contract != CONTRACT_NAME {
-        return Err(ContractError::InvalidContractName {
-            expected: CONTRACT_NAME.to_string(),
-            actual: prev_version.contract,
-        });
-    }
-
-    // Update to the new version
+    // Attempt to get the current version stored in the contract
+    let prev_version = match cw2::get_contract_version(deps.storage) {
+        Ok(version) => {
+            // Contract has version info - validate it's the expected contract
+            if version.contract != CONTRACT_NAME {
+                return Err(ContractError::InvalidContractName {
+                    expected: CONTRACT_NAME.to_string(),
+                    actual: version.contract,
+                });
+            }
+            version
+        }
+        Err(_) => {
+            // No version info exists - this is a legacy contract
+            // Create a placeholder version for logging purposes
+            cw2::ContractVersion {
+                contract: CONTRACT_NAME.to_string(),
+                version: LEGACY_VERSION.to_string(),
+            }
+        }
+    };
+
+    // Update to the new version (this creates version info if it didn't exist)
     cw2::set_contract_version(deps.storage, CONTRACT_NAME, CONTRACT_VERSION)?;
 
     Ok(Response::new()
@@ -507,6 +521,8 @@ pub fn migrate(
 **Migration Features:**
 - **Automatic Version Tracking**: Uses the `cw2` library for contract version
   management
+- **Legacy Contract Support**: Gracefully handles contracts deployed before
+  `cw2` version tracking was implemented
 - **Contract Validation**: Ensures the contract name matches before migration
 - **Version Updates**: Automatically updates stored version information
 - **Detailed Logging**: Returns attributes showing the version transition
@@ -530,6 +546,8 @@ the expected name (`"rollup-bsn/finality"`).
 **Version Management:**
 - **Contract Name**: `"rollup-bsn/finality"` (constant)
 - **Contract Version**: Uses `CARGO_PKG_VERSION` from the package manifest
+- **Legacy Version**: `"pre-cw2"` (constant used for contracts deployed before
+  version tracking)
 - **Storage**: Version information is stored using the `cw2` library
 
 **Security Considerations:**