feat: add macro to inject contract version validation and setting the…

… new version during migration (#746)

contracts/router/src/contract.rs

@@ -1,11 +1,10 @@
-use axelar_wasm_std::{address, killswitch, permission_control, FnExt};
+use axelar_wasm_std::{address, killswitch, migrate_from_version, permission_control, FnExt};
 #[cfg(not(feature = "library"))]
 use cosmwasm_std::entry_point;
 use cosmwasm_std::{
     to_json_binary, Addr, Binary, Deps, DepsMut, Env, MessageInfo, Response, Storage,
 };
 use router_api::error::Error;
-use semver::{Version, VersionReq};
 
 use crate::contract::migrations::v1_1_1;
 use crate::events::RouterInstantiated;
@@ -21,21 +20,13 @@ pub const CONTRACT_NAME: &str = env!("CARGO_PKG_NAME");
 pub const CONTRACT_VERSION: &str = env!("CARGO_PKG_VERSION");
 
 #[cfg_attr(not(feature = "library"), entry_point)]
+#[migrate_from_version("1.1")]
 pub fn migrate(
     deps: DepsMut,
     _env: Env,
     msg: MigrateMsg,
 ) -> Result<Response, axelar_wasm_std::error::ContractError> {
-    let old_version = Version::parse(&cw2::get_contract_version(deps.storage)?.version)?;
-    let version_requirement = VersionReq::parse(">= 1.1.0, < 1.2.0")?;
-    assert!(version_requirement.matches(&old_version));
-
     v1_1_1::migrate(deps.storage, msg.chains_to_remove)?;
-
-    // this needs to be the last thing to do during migration,
-    // because previous migration steps should check the old version
-    cw2::set_contract_version(deps.storage, CONTRACT_NAME, CONTRACT_VERSION)?;
-
     Ok(Response::default())
 }
 

packages/axelar-wasm-std-derive/Cargo.toml

@@ -16,11 +16,14 @@ itertools = { workspace = true }
 proc-macro2 = { workspace = true }
 quote = { workspace = true }
 report = { workspace = true }
+semver = { workspace = true }
 syn = { workspace = true }
 thiserror = { workspace = true }
 
 [dev-dependencies]
-axelar-wasm-std = { workspace = true }
+assert_ok = { workspace = true }
+axelar-wasm-std = { workspace = true, features = ["derive"] }
+cw2 = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
 

packages/axelar-wasm-std-derive/src/lib.rs

@@ -5,6 +5,7 @@ use itertools::Itertools;
 use proc_macro::TokenStream;
 use proc_macro2::{Ident, Span, TokenStream as TokenStream2};
 use quote::quote;
+use syn::spanned::Spanned;
 use syn::{DeriveInput, FieldsNamed, Generics, ItemEnum, Variant};
 
 #[proc_macro_derive(IntoContractError)]
@@ -282,3 +283,165 @@ fn match_unit_variant(event_enum: &Ident, variant_name: &Ident) -> TokenStream2
         #event_enum::#variant_name => cosmwasm_std::Event::new(#event_name)
     }
 }
+
+/// Attribute macro for handling contract version migrations. Must be applied to the `migrate` contract entry point.
+/// Checks if migrating from the current version is supported and sets the new version. The base version must be a valid semver without patch, pre, or build.
+///
+/// # Example
+/// ```
+/// use cosmwasm_std::{ DepsMut, Env, Response, Empty};
+/// use axelar_wasm_std_derive::migrate_from_version;
+///
+/// #[migrate_from_version("1.1")]
+/// pub fn migrate(
+///     deps: DepsMut,
+///     _env: Env,
+///     _msg: Empty,
+/// ) -> Result<Response, axelar_wasm_std::error::ContractError> {
+///     // migration logic
+///     Ok(Response::default())
+/// }
+/// ```
+///
+/// ```compile_fail
+/// # use cosmwasm_std::{ DepsMut, Env, Response, Empty};
+/// # use axelar_wasm_std_derive::migrate_from_version;
+///
+/// # #[migrate_from_version("1.1")] // compilation error because the macro is not applied to a function `migrate`
+/// # pub fn execute(
+/// #     deps: DepsMut,
+/// #     _env: Env,
+/// #     _msg: Empty,
+/// # ) -> Result<Response, axelar_wasm_std::error::ContractError> {
+/// #     Ok(Response::default())
+/// # }
+/// ```
+///
+/// ```compile_fail
+/// # use cosmwasm_std::{ Deps, Env, Response, Empty};
+/// # use axelar_wasm_std_derive::migrate_from_version;
+///
+/// # #[migrate_from_version("1.1")] // compilation error because it cannot parse a `DepsMut` parameter
+/// # pub fn migrate(
+/// #     deps: Deps,
+/// #     _env: Env,
+/// #     _msg: Empty,
+/// # ) -> Result<Response, axelar_wasm_std::error::ContractError> {
+/// #     Ok(Response::default())
+/// # }
+/// ```
+///
+/// ```compile_fail
+/// # use cosmwasm_std::{ DepsMut, Env, Response, Empty};
+/// # use axelar_wasm_std_derive::migrate_from_version;
+///
+/// # #[migrate_from_version("~1.1.0")] // compilation error because the base version is not formatted correctly
+/// # pub fn migrate(
+/// #     deps: DepsMut,
+/// #     _env: Env,
+/// #     _msg: Empty,
+/// # ) -> Result<Response, axelar_wasm_std::error::ContractError> {
+/// #     Ok(Response::default())
+/// # }
+/// ```
+///
+#[proc_macro_attribute]
+pub fn migrate_from_version(input: TokenStream, item: TokenStream) -> TokenStream {
+    let base_version_req = syn::parse_macro_input!(input as syn::LitStr);
+    let annotated_fn = syn::parse_macro_input!(item as syn::ItemFn);
+
+    try_migrate_from_version(base_version_req, annotated_fn)
+        .unwrap_or_else(syn::Error::into_compile_error)
+        .into()
+}
+
+fn try_migrate_from_version(
+    base_version: syn::LitStr,
+    annotated_fn: syn::ItemFn,
+) -> syn::Result<TokenStream2> {
+    let fn_name = &annotated_fn.sig.ident;
+    let fn_inputs = &annotated_fn.sig.inputs;
+    let fn_output = &annotated_fn.sig.output;
+    let fn_block = &annotated_fn.block;
+
+    let base_semver_req = base_semver_req(&base_version)?;
+    let deps = validate_migrate_signature(&annotated_fn.sig)?;
+
+    let gen = quote! {
+        pub fn #fn_name(#fn_inputs) #fn_output {
+            let pkg_name = env!("CARGO_PKG_NAME");
+            let pkg_version = env!("CARGO_PKG_VERSION");
+
+            let contract_version = cw2::get_contract_version(#deps.storage)?;
+            assert_eq!(contract_version.contract, pkg_name, "contract name mismatch: actual {}, expected {}", contract_version.contract, pkg_name);
+
+            let curr_version = semver::Version::parse(&contract_version.version)?;
+            let version_requirement = semver::VersionReq::parse(#base_semver_req)?;
+            assert!(version_requirement.matches(&curr_version), "base version {} does not match {} version requirement", curr_version, #base_semver_req);
+
+            cw2::set_contract_version(#deps.storage, pkg_name, pkg_version)?;
+
+            #fn_block
+        }
+    };
+
+    Ok(gen)
+}
+
+fn base_semver_req(base_version: &syn::LitStr) -> syn::Result<String> {
+    let base_semver = semver::Version::parse(&format!("{}.0", base_version.value()))
+        .map_err(|_| syn::Error::new(base_version.span(), "base version format must be semver without patch, pre, or build. Example: '1.2'"))
+        .and_then(|version| {
+            if version.patch == 0 && version.pre.is_empty() && version.build.is_empty() {
+                Ok(version)
+            } else {
+                Err(syn::Error::new(base_version.span(), "base version format must be semver without patch, pre, or build. Example: '1.2'"))
+            }
+        })?;
+
+    Ok(format!("~{}.{}.0", base_semver.major, base_semver.minor))
+}
+
+fn validate_migrate_signature(sig: &syn::Signature) -> syn::Result<syn::Ident> {
+    if sig.ident != "migrate"
+        || sig.inputs.len() != 3
+        || !matches!(sig.output, syn::ReturnType::Type(_, _))
+    {
+        return Err(syn::Error::new(
+            sig.ident.span(),
+            "invalid function signature for 'migrate' entry point",
+        ));
+    }
+
+    validate_migrate_param(&sig.inputs[1], "Env")?;
+    validate_migrate_param(&sig.inputs[0], "DepsMut")
+}
+
+fn validate_migrate_param(param: &syn::FnArg, expected_type: &str) -> syn::Result<syn::Ident> {
+    let (ty, pat) = match param {
+        syn::FnArg::Typed(syn::PatType { ty, pat, .. }) => (ty, pat),
+        _ => {
+            return Err(syn::Error::new(
+                param.span(),
+                format!(
+                    "parameter for 'migrate' entry point expected to be of type {}",
+                    expected_type
+                ),
+            ));
+        }
+    };
+    match (&**ty, &**pat) {
+        (syn::Type::Path(syn::TypePath { path, .. }), syn::Pat::Ident(pat_ident))
+            if path.is_ident(expected_type) =>
+        {
+            Ok(pat_ident.ident.clone())
+        }
+        _ => Err(syn::Error::new(
+            ty.span(),
+            format!(
+                "parameter for 'migrate' entry point expected to be of type {}",
+                expected_type
+            ),
+        )),
+    }
+}

packages/axelar-wasm-std-derive/tests/derive.rs

@@ -1,5 +1,8 @@
+use assert_ok::assert_ok;
 use axelar_wasm_std::error::ContractError;
-use axelar_wasm_std::IntoContractError;
+use axelar_wasm_std::{migrate_from_version, IntoContractError};
+use cosmwasm_std::testing::{mock_dependencies, mock_env};
+use cosmwasm_std::{DepsMut, Empty, Env, Response};
 use thiserror::Error;
 
 #[derive(Error, Debug, IntoContractError)]
@@ -12,3 +15,55 @@ enum TestError {
 fn can_convert_error() {
     _ = ContractError::from(TestError::Something);
 }
+
+#[migrate_from_version("999.1")]
+pub fn migrate(
+    deps: DepsMut,
+    _env: Env,
+    _msg: Empty,
+) -> Result<Response, axelar_wasm_std::error::ContractError> {
+    // migration logic
+    deps.storage.set(b"key", b"migrated value");
+    Ok(Response::default())
+}
+
+#[test]
+fn should_handle_version_migration() {
+    let mut deps = mock_dependencies();
+
+    let base_contract = env!("CARGO_PKG_NAME");
+    let base_version = "999.1.1";
+    cw2::set_contract_version(deps.as_mut().storage, base_contract, base_version).unwrap();
+    deps.as_mut().storage.set(b"key", b"original value");
+
+    migrate(deps.as_mut(), mock_env(), Empty {}).unwrap();
+
+    let contract_version = assert_ok!(cw2::get_contract_version(deps.as_ref().storage));
+    assert_eq!(contract_version.contract, base_contract);
+    assert_eq!(contract_version.version, env!("CARGO_PKG_VERSION"));
+
+    let migrated_value = deps.as_ref().storage.get(b"key").unwrap();
+    assert_eq!(migrated_value, b"migrated value")
+}
+
+#[test]
+#[should_panic(expected = "base version 999.2.1 does not match ~999.1.0 version requirement")]
+fn should_fail_version_migration_if_not_supported() {
+    let mut deps = mock_dependencies();
+    let base_contract = env!("CARGO_PKG_NAME");
+    let base_version = "999.2.1";
+    cw2::set_contract_version(deps.as_mut().storage, base_contract, base_version).unwrap();
+
+    migrate(deps.as_mut(), mock_env(), Empty {}).unwrap();
+}
+
+#[test]
+#[should_panic(expected = "contract name mismatch: actual wrong-base-contract, expected ")]
+fn should_fail_version_migration_using_wrong_contract() {
+    let mut deps = mock_dependencies();
+    let base_contract = "wrong-base-contract";
+    let base_version = "999.1.1";
+    cw2::set_contract_version(deps.as_mut().storage, base_contract, base_version).unwrap();
+
+    migrate(deps.as_mut(), mock_env(), Empty {}).unwrap();
+}