From 89bd35040cdf8f1b002589ba2e2b308e417a7106 Mon Sep 17 00:00:00 2001
From: Noah Saso <noahsaso@gmail.com>
Date: Thu, 1 Aug 2024 23:03:40 -0400
Subject: [PATCH] migrate version only if newer pattern, and prevent migrating
 from other contracts

---
 Cargo.lock                                     |  1 +
 .../dao-rewards-distributor/Cargo.toml         |  1 +
 .../dao-rewards-distributor/src/contract.rs    | 18 +++++++++++++++++-
 .../dao-rewards-distributor/src/error.rs       | 12 ++++++++++++
 4 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/Cargo.lock b/Cargo.lock
index f821573c9..b1e7d67d1 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2222,6 +2222,7 @@ dependencies = [
  "dao-voting-cw4 2.5.0",
  "dao-voting-cw721-staked",
  "dao-voting-token-staked",
+ "semver",
  "thiserror",
 ]
 
diff --git a/contracts/distribution/dao-rewards-distributor/Cargo.toml b/contracts/distribution/dao-rewards-distributor/Cargo.toml
index 0df6ef1d5..4d38cbc3b 100644
--- a/contracts/distribution/dao-rewards-distributor/Cargo.toml
+++ b/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]
diff --git a/contracts/distribution/dao-rewards-distributor/src/contract.rs b/contracts/distribution/dao-rewards-distributor/src/contract.rs
index 3a6e3c6e0..db4b1e9ab 100644
--- a/contracts/distribution/dao-rewards-distributor/src/contract.rs
+++ b/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;
 
@@ -592,6 +593,21 @@ fn query_distributions(
 
 #[cfg_attr(not(feature = "library"), entry_point)]
 pub fn migrate(deps: DepsMut, _env: Env, _msg: MigrateMsg) -> Result<Response, ContractError> {
-    set_contract_version(deps.storage, CONTRACT_NAME, CONTRACT_VERSION)?;
+    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 version: Version = CONTRACT_VERSION.parse()?;
+    let storage_version: Version = contract_version.version.parse()?;
+
+    if storage_version < version {
+        set_contract_version(deps.storage, CONTRACT_NAME, CONTRACT_VERSION)?;
+    }
+
     Ok(Response::default())
 }
diff --git a/contracts/distribution/dao-rewards-distributor/src/error.rs b/contracts/distribution/dao-rewards-distributor/src/error.rs
index 94d7bbb85..18101e127 100644
--- a/contracts/distribution/dao-rewards-distributor/src/error.rs
+++ b/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,13 @@ 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("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())
+    }
 }