From 9919f7d75cdd9d7d53d3b89feb6333b584b5d42d Mon Sep 17 00:00:00 2001
From: javiersuweijie <javier.su.weijie@gmail.com>
Date: Thu, 8 Feb 2024 13:30:07 +0800
Subject: [PATCH] fix: issue with querying rewards for asset not staked

---
 contracts/alliance-hub/src/contract.rs  | 12 ++++++++++--
 contracts/alliance-hub/src/query.rs     |  4 ++--
 packages/alliance-protocol/src/error.rs |  3 +++
 3 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/contracts/alliance-hub/src/contract.rs b/contracts/alliance-hub/src/contract.rs
index d0d31bb..08d4bf6 100644
--- a/contracts/alliance-hub/src/contract.rs
+++ b/contracts/alliance-hub/src/contract.rs
@@ -39,7 +39,7 @@ const CREATE_REPLY_ID: u64 = 1;
 const CLAIM_REWARD_ERROR_REPLY_ID: u64 = 2;
 
 #[cfg_attr(not(feature = "library"), entry_point)]
-pub fn migrate(_deps: DepsMut, _env: Env, _msg: MigrateMsg) -> StdResult<Response> {
+pub fn migrate(deps: DepsMut, _env: Env, _msg: MigrateMsg) -> StdResult<Response> {
     Ok(Response::default())
 }
 
@@ -238,7 +238,7 @@ fn unstake(deps: DepsMut, info: MessageInfo, asset: Asset) -> Result<Response, C
                     }
                     Ok(balance - asset.amount)
                 }
-                None => Err(ContractError::InsufficientBalance {}),
+                None => Err(ContractError::AssetNotStaked {}),
             }
         },
     )?;
@@ -308,6 +308,14 @@ fn _claim_reward(
     asset: AssetInfo,
 ) -> Result<Uint128, ContractError> {
     let asset_key = AssetInfoKey::from(&asset);
+    // If the user did not stake the asset, do nothing and return 0
+    if BALANCES
+        .load(storage, (user.clone(), asset_key.clone()))
+        .is_err()
+    {
+        return Ok(Uint128::zero());
+    }
+
     let user_reward_rate = USER_ASSET_REWARD_RATE.load(storage, (user.clone(), asset_key.clone()));
     let asset_reward_rate = ASSET_REWARD_RATE.load(storage, asset_key.clone())?;
 
diff --git a/contracts/alliance-hub/src/query.rs b/contracts/alliance-hub/src/query.rs
index 63da2f8..9795cec 100644
--- a/contracts/alliance-hub/src/query.rs
+++ b/contracts/alliance-hub/src/query.rs
@@ -79,7 +79,7 @@ fn get_pending_rewards(deps: Deps, asset_query: AssetQuery) -> StdResult<Binary>
     let user_reward_rate = USER_ASSET_REWARD_RATE.load(deps.storage, key.clone())?;
     let asset_reward_rate =
         ASSET_REWARD_RATE.load(deps.storage, AssetInfoKey::from(asset_query.asset.clone()))?;
-    let user_balance = BALANCES.load(deps.storage, key.clone())?;
+    let user_balance = BALANCES.load(deps.storage, key.clone()).unwrap_or(Uint128::zero());
     let unclaimed_rewards = UNCLAIMED_REWARDS
         .load(deps.storage, key)
         .unwrap_or(Uint128::zero());
@@ -131,7 +131,7 @@ fn get_all_pending_rewards(deps: Deps, query: AllPendingRewardsQuery) -> StdResu
             let user_balance = BALANCES.load(
                 deps.storage,
                 (addr.clone(), AssetInfoKey::from(asset.clone())),
-            )?;
+            ).unwrap_or(Uint128::zero());
             let unclaimed_rewards = UNCLAIMED_REWARDS
                 .load(
                     deps.storage,
diff --git a/packages/alliance-protocol/src/error.rs b/packages/alliance-protocol/src/error.rs
index 244b935..7b3f795 100644
--- a/packages/alliance-protocol/src/error.rs
+++ b/packages/alliance-protocol/src/error.rs
@@ -39,4 +39,7 @@ pub enum ContractError {
 
     #[error("Invalid total distribution: {0}")]
     InvalidTotalDistribution(Decimal),
+
+    #[error("Asset not staked")]
+    AssetNotStaked {},
 }