added migration version and contract checks (#869)

contracts/distribution/dao-rewards-distributor/Cargo.toml

@@ -29,6 +29,7 @@ cw-utils = { workspace = true }
 dao-hooks = { workspace = true }
 dao-interface = { workspace = true }
 dao-voting = { workspace = true }
+semver = { workspace = true }
 thiserror = { workspace = true }
 
 [dev-dependencies]

contracts/distribution/dao-rewards-distributor/schema/dao-rewards-distributor.json

@@ -847,8 +847,8 @@
   "migrate": {
     "$schema": "http://json-schema.org/draft-07/schema#",
     "title": "MigrateMsg",
-    "type": "string",
-    "enum": []
+    "type": "object",
+    "additionalProperties": false
   },
   "sudo": null,
   "responses": {

contracts/distribution/dao-rewards-distributor/src/contract.rs

@@ -9,6 +9,7 @@ use cw20::{Cw20ReceiveMsg, Denom};
 use cw_storage_plus::Bound;
 use cw_utils::{must_pay, nonpayable, Duration, Expiration};
 use dao_interface::voting::InfoResponse;
+use semver::Version;
 
 use std::ops::Add;
 
@@ -27,8 +28,8 @@ use crate::rewards::{
 use crate::state::{DistributionState, EmissionRate, Epoch, COUNT, DISTRIBUTIONS, USER_REWARDS};
 use crate::ContractError;
 
-const CONTRACT_NAME: &str = env!("CARGO_PKG_NAME");
-const CONTRACT_VERSION: &str = env!("CARGO_PKG_VERSION");
+pub(crate) const CONTRACT_NAME: &str = env!("CARGO_PKG_NAME");
+pub(crate) const CONTRACT_VERSION: &str = env!("CARGO_PKG_VERSION");
 
 pub const DEFAULT_LIMIT: u32 = 10;
 pub const MAX_LIMIT: u32 = 50;
@@ -657,6 +658,27 @@ fn query_distributions(
 
 #[cfg_attr(not(feature = "library"), entry_point)]
 pub fn migrate(deps: DepsMut, _env: Env, _msg: MigrateMsg) -> Result<Response, ContractError> {
+    let contract_version = get_contract_version(deps.storage)?;
+
+    if contract_version.contract != CONTRACT_NAME {
+        return Err(ContractError::MigrationErrorIncorrectContract {
+            expected: CONTRACT_NAME.to_string(),
+            actual: contract_version.contract,
+        });
+    }
+
+    let new_version: Version = CONTRACT_VERSION.parse()?;
+    let current_version: Version = contract_version.version.parse()?;
+
+    // only allow upgrades
+    if new_version <= current_version {
+        return Err(ContractError::MigrationErrorInvalidVersion {
+            new: new_version.to_string(),
+            current: current_version.to_string(),
+        });
+    }
+
     set_contract_version(deps.storage, CONTRACT_NAME, CONTRACT_VERSION)?;
+
     Ok(Response::default())
 }

contracts/distribution/dao-rewards-distributor/src/error.rs

@@ -22,6 +22,9 @@ pub enum ContractError {
     #[error(transparent)]
     Payment(#[from] PaymentError),
 
+    #[error("semver parsing error: {0}")]
+    SemVer(String),
+
     #[error("Invalid CW20")]
     InvalidCw20 {},
 
@@ -54,4 +57,16 @@ pub enum ContractError {
 
     #[error("Cannot update emission rate because this distribution has accumulated the maximum rewards. Start a new distribution with the new emission rate instead. (Overflow: {err})")]
     DistributionHistoryTooLarge { err: String },
+
+    #[error("Invalid version migration. {new} is not newer than {current}.")]
+    MigrationErrorInvalidVersion { new: String, current: String },
+
+    #[error("Expected to migrate from contract {expected}. Got {actual}.")]
+    MigrationErrorIncorrectContract { expected: String, actual: String },
+}
+
+impl From<semver::Error> for ContractError {
+    fn from(err: semver::Error) -> Self {
+        Self::SemVer(err.to_string())
+    }
 }

contracts/distribution/dao-rewards-distributor/src/msg.rs

@@ -133,4 +133,4 @@ pub struct DistributionPendingRewards {
 }
 
 #[cw_serde]
-pub enum MigrateMsg {}
+pub struct MigrateMsg {}

contracts/distribution/dao-rewards-distributor/src/testing/suite.rs

@@ -109,6 +109,7 @@ impl SuiteBuilder {
             owner: Some(owner.clone()),
             staking_addr: Addr::unchecked(""),
             voting_power_addr: Addr::unchecked(""),
+            reward_code_id: 0,
             distribution_contract: Addr::unchecked(""),
             cw20_addr: Addr::unchecked(""),
             reward_denom: DENOM.to_string(),
@@ -229,12 +230,12 @@ impl SuiteBuilder {
         };
 
         // initialize the rewards distributor
-        let reward_code_id = suite_built.app.borrow_mut().store_code(contract_rewards());
+        suite_built.reward_code_id = suite_built.app.borrow_mut().store_code(contract_rewards());
         let reward_addr = suite_built
             .app
             .borrow_mut()
             .instantiate_contract(
-                reward_code_id,
+                suite_built.reward_code_id,
                 owner.clone(),
                 &InstantiateMsg {
                     owner: Some(owner.clone().into_string()),
@@ -327,6 +328,7 @@ pub struct Suite {
     pub voting_power_addr: Addr,
     pub reward_denom: String,
 
+    pub reward_code_id: u64,
     pub distribution_contract: Addr,
 
     // cw20 type fields

contracts/distribution/dao-rewards-distributor/src/testing/tests.rs

@@ -1,5 +1,6 @@
 use std::borrow::BorrowMut;
 
+use cosmwasm_std::testing::{mock_dependencies, mock_env};
 use cosmwasm_std::{coin, coins, to_json_binary, Addr, Timestamp};
 use cosmwasm_std::{Uint128, Uint256};
 use cw2::ContractVersion;
@@ -9,7 +10,8 @@ use cw_multi_test::Executor;
 use cw_utils::Duration;
 use dao_interface::voting::InfoResponse;
 
-use crate::msg::{CreateMsg, FundMsg};
+use crate::contract::{CONTRACT_NAME, CONTRACT_VERSION};
+use crate::msg::{CreateMsg, FundMsg, MigrateMsg};
 use crate::state::{EmissionRate, Epoch};
 use crate::testing::native_setup::setup_native_token_test;
 use crate::ContractError;
@@ -2481,3 +2483,48 @@ fn test_large_stake_before_claim() {
     suite.claim_rewards(ADDR2, 1);
     suite.claim_rewards(ADDR3, 1);
 }
+
+#[test]
+fn test_migrate() {
+    let mut deps = mock_dependencies();
+
+    cw2::set_contract_version(&mut deps.storage, "test", "0.0.1").unwrap();
+
+    // wrong contract name errors
+    let err: ContractError =
+        crate::contract::migrate(deps.as_mut(), mock_env(), MigrateMsg {}).unwrap_err();
+    assert_eq!(
+        err,
+        ContractError::MigrationErrorIncorrectContract {
+            expected: CONTRACT_NAME.to_string(),
+            actual: "test".to_string(),
+        }
+    );
+
+    // migration succeeds from past version of same contract
+    cw2::set_contract_version(&mut deps.storage, CONTRACT_NAME, "0.0.1").unwrap();
+    crate::contract::migrate(deps.as_mut(), mock_env(), MigrateMsg {}).unwrap();
+
+    // same-version migration errors
+    let err: ContractError =
+        crate::contract::migrate(deps.as_mut(), mock_env(), MigrateMsg {}).unwrap_err();
+    assert_eq!(
+        err,
+        ContractError::MigrationErrorInvalidVersion {
+            new: CONTRACT_VERSION.to_string(),
+            current: CONTRACT_VERSION.to_string(),
+        }
+    );
+
+    // future version errors
+    cw2::set_contract_version(&mut deps.storage, CONTRACT_NAME, "9.9.9").unwrap();
+    let err: ContractError =
+        crate::contract::migrate(deps.as_mut(), mock_env(), MigrateMsg {}).unwrap_err();
+    assert_eq!(
+        err,
+        ContractError::MigrationErrorInvalidVersion {
+            new: CONTRACT_VERSION.to_string(),
+            current: "9.9.9".to_string(),
+        }
+    );
+}