securesecrets · FloppyDisck · Oct 23, 2021 · Oct 11, 2021 · Oct 12, 2021 · Oct 17, 2021
diff --git a/contracts/governance/src/contract.rs b/contracts/governance/src/contract.rs
@@ -12,6 +12,7 @@ use crate::{
 };
 use shade_protocol::asset::Contract;
 use crate::state::{admin_commands_list_w, supported_contracts_list_w};
+use crate::handle::try_disable_staker;
 
 pub fn init<S: Storage, A: Api, Q: Querier>(
     deps: &mut Extern<S, A, Q>,
@@ -24,6 +25,7 @@ pub fn init<S: Storage, A: Api, Q: Querier>(
             None => { env.message.sender.clone() }
             Some(admin) => { admin }
         },
+        staker: msg.staker,
         proposal_deadline: msg.proposal_deadline,
         minimum_votes: msg.quorum
     };
@@ -56,9 +58,11 @@ pub fn handle<S: Storage, A: Api, Q: Querier>(
 
         /// Self interactions
         // Config
-        HandleMsg::UpdateConfig { admin, proposal_deadline,
+        HandleMsg::UpdateConfig { admin, staker, proposal_deadline,
             minimum_votes } =>
-            handle::try_update_config(deps, &env, admin, proposal_deadline, minimum_votes),
+            handle::try_update_config(deps, &env, admin, staker, proposal_deadline, minimum_votes),
+
+        HandleMsg::DisableStaker {} => try_disable_staker(deps, &env),
 
         // Supported contract
         HandleMsg::AddSupportedContract { name, contract
@@ -81,8 +85,8 @@ pub fn handle<S: Storage, A: Api, Q: Querier>(
         } => handle::try_update_admin_command(deps, &env, name, proposal),
 
         /// User interaction
-        HandleMsg::MakeVote { proposal_id, option
-        } => handle::try_vote(deps, &env, proposal_id, option),
+        HandleMsg::MakeVote { voter, proposal_id, votes
+        } => handle::try_vote(deps, &env, voter, proposal_id, votes),
 
         HandleMsg::TriggerProposal { proposal_id
         } => handle::try_trigger_proposal(deps, &env, proposal_id),
@@ -108,6 +112,9 @@ pub fn query<S: Storage, A: Api, Q: Querier>(
         QueryMsg::GetTotalProposals {} => to_binary(
             &query::total_proposals(deps)?),
 
+        QueryMsg::GetProposalVotes { proposal_id } => to_binary(
+            &query::proposal_votes(deps, proposal_id)?),
+
         QueryMsg::GetSupportedContracts { } => to_binary(&query::supported_contracts(deps)?),
 
         QueryMsg::GetSupportedContract { name } => to_binary(

diff --git a/contracts/governance/src/handle.rs b/contracts/governance/src/handle.rs
@@ -1,13 +1,14 @@
-use cosmwasm_std::{to_binary, Api, Binary, Env, Extern, HandleResponse, Querier, StdError, StdResult, Storage, CosmosMsg, HumanAddr, Uint128, WasmMsg};
-use crate::state::{supported_contract_r, config_r, total_proposals_w, proposal_w, config_w, supported_contract_w, proposal_r, admin_commands_r, admin_commands_w, admin_commands_list_w, supported_contracts_list_w};
+use cosmwasm_std::{to_binary, Api, Binary, Env, Extern, HandleResponse, Querier, StdError, StdResult, Storage, CosmosMsg, HumanAddr, Uint128, WasmMsg, Empty};
+use crate::state::{supported_contract_r, config_r, total_proposals_w, proposal_w, config_w, supported_contract_w, proposal_r, admin_commands_r, admin_commands_w, admin_commands_list_w, supported_contracts_list_w, total_proposals_r, total_proposal_votes_w, proposal_votes_w, total_proposal_votes_r, proposal_votes_r};
 use shade_protocol::{
     governance::{Proposal, ProposalStatus, HandleAnswer, GOVERNANCE_SELF, HandleMsg},
     generic_response::ResponseStatus,
     asset::Contract,
 };
-use shade_protocol::governance::ProposalStatus::Accepted;
+use shade_protocol::governance::ProposalStatus::{Accepted, Expired, Rejected};
 use shade_protocol::generic_response::ResponseStatus::{Success, Failure};
-use shade_protocol::governance::{AdminCommand, ADMIN_COMMAND_VARIABLE, Vote};
+use shade_protocol::governance::{AdminCommand, ADMIN_COMMAND_VARIABLE, Vote, UserVote, VoteTally};
+use secret_toolkit::utils::HandleCallback;
 
 pub fn create_proposal<S: Storage, A: Api, Q: Querier>(
     deps: &mut Extern<S, A, Q>,
@@ -45,6 +46,14 @@ pub fn create_proposal<S: Storage, A: Api, Q: Querier>(
     // Store the proposal
     proposal_w(&mut deps.storage).save(proposal_id.to_string().as_bytes(), &proposal)?;
 
+    // Create proposal votes
+    total_proposal_votes_w(&mut deps.storage).save(
+        proposal_id.to_string().as_bytes(), &VoteTally{
+        yes: Uint128(0),
+        no: Uint128(0),
+        abstain: Uint128(0)
+    })?;
+
     Ok(proposal_id)
 }
 
@@ -64,8 +73,37 @@ pub fn try_trigger_proposal<S: Storage, A: Api, Q: Querier>(
             backtrace: None })
     }
 
-    // TODO: missing part where proposal votes are counted and vote_status is updated
-    proposal.vote_status = Accepted;
+    // Change proposal behavior according to stake availability
+    let config = config_r(&deps.storage).load()?;
+    proposal.vote_status = match config.staker {
+        Some(staker) => {
+
+            let total_votes = total_proposal_votes_r(&deps.storage).load(
+                proposal_id.to_string().as_bytes())?;
+
+            // Check if proposal can be run
+            if proposal.due_date > env.block.time {
+                Err(StdError::Unauthorized { backtrace: None })
+            }
+            else if total_votes.yes + total_votes.no + total_votes.abstain < config.minimum_votes {
+                Ok(Expired)
+            }
+            else if total_votes.yes > total_votes.no {
+                Ok(Accepted)
+            }
+            else {
+                Ok(Rejected)
+            }
+        }
+        None => {
+            // Check if user is an admin in order to trigger the proposal
+            if config.admin == env.message.sender {
+                Ok(Accepted)
+            } else {
+                Err(StdError::Unauthorized { backtrace: None })
+            }
+        }
+    }?;
 
     let mut messages: Vec<CosmosMsg> = vec![];
 
@@ -120,19 +158,61 @@ pub fn try_execute_msg(
 }
 
 pub fn try_vote<S: Storage, A: Api, Q: Querier>(
-    _deps: &mut Extern<S, A, Q>,
-    _env: &Env,
-    _proposal_id: Uint128,
-    _option: Vote,
+    deps: &mut Extern<S, A, Q>,
+    env: &Env,
+    voter: HumanAddr,
+    proposal_id: Uint128,
+    votes: VoteTally,
 ) -> StdResult<HandleResponse> {
 
-    // TODO: reach consensus on voting
-    // TODO: SNIP20 could have a voting command and it will automatically vote on governance
+    // Check that sender is staking contract and staking is enabled
+    let config = config_r(&deps.storage).load()?;
+    if config.staker.is_none() || config.staker.unwrap().address != env.message.sender {
+        return Err(StdError::Unauthorized { backtrace: None })
+    }
+
+    // Check that proposal exists
+    if proposal_id > total_proposals_r(&deps.storage).load()? {
+        return Err(StdError::NotFound { kind: "Proposal".to_string(), backtrace: None })
+    }
+
+    // Check that proposal is still votable
+    let proposal = proposal_r(&deps.storage).load(proposal_id.to_string().as_bytes())?;
+
+    if proposal.vote_status != ProposalStatus::InProgress || proposal.due_date <= env.block.time {
+        return Err(StdError::Unauthorized { backtrace: None })
+    }
+
+    // Get proposal voting state
+    let mut proposal_voting_state = total_proposal_votes_r(&deps.storage).load(
+        proposal_id.to_string().as_bytes())?;
+
+    // Check if user has already voted
+    match proposal_votes_r(&deps.storage, proposal_id).may_load(
+        voter.to_string().as_bytes())? {
+        None => {}
+        Some(old_votes) => {
+            // Remove those votes from state
+            proposal_voting_state.yes = (proposal_voting_state.yes - old_votes.yes)?;
+            proposal_voting_state.no = (proposal_voting_state.no - old_votes.no)?;
+            proposal_voting_state.abstain = (proposal_voting_state.abstain - old_votes.abstain)?;
+        }
+    }
+
+    // Update state
+    proposal_voting_state.yes += votes.yes;
+    proposal_voting_state.no += votes.no;
+    proposal_voting_state.abstain += votes.abstain;
+
+    // Save staker info
+    total_proposal_votes_w(&mut deps.storage).save(proposal_id.to_string().as_bytes(),
+                                                   &proposal_voting_state)?;
+    proposal_votes_w(&mut deps.storage, proposal_id).save(voter.to_string().as_bytes(), &votes)?;
 
     Ok(HandleResponse {
         messages: vec![],
         log: vec![],
-        data: Some( to_binary( &HandleAnswer::Vote {
+        data: Some( to_binary( &HandleAnswer::MakeVote {
             status: Success,
         })?),
     })
@@ -241,6 +321,7 @@ pub fn try_update_config<S: Storage, A: Api, Q: Querier>(
     deps: &mut Extern<S, A, Q>,
     env: &Env,
     admin: Option<HumanAddr>,
+    staker: Option<Contract>,
     proposal_deadline: Option<u64>,
     minimum_votes: Option<Uint128>,
 ) -> StdResult<HandleResponse> {
@@ -254,6 +335,9 @@ pub fn try_update_config<S: Storage, A: Api, Q: Querier>(
         if let Some(admin) = admin {
             state.admin = admin;
         }
+        if staker.is_some() {
+            state.staker = staker;
+        }
         if let Some(proposal_deadline) = proposal_deadline {
             state.proposal_deadline = proposal_deadline;
         }
@@ -273,6 +357,25 @@ pub fn try_update_config<S: Storage, A: Api, Q: Querier>(
     })
 }
 
+pub fn try_disable_staker<S: Storage, A: Api, Q: Querier>(
+    deps: &mut Extern<S, A, Q>,
+    env: &Env,
+) -> StdResult<HandleResponse> {
+
+    config_w(&mut deps.storage).update(|mut state| {
+        state.staker = None;
+        Ok(state)
+    })?;
+
+    Ok(HandleResponse {
+        messages: vec![],
+        log: vec![],
+        data: Some( to_binary( &HandleAnswer::DisableStaker {
+            status: ResponseStatus::Success
+        })?),
+    })
+}
+
 pub fn try_add_supported_contract<S: Storage, A: Api, Q: Querier>(
     deps: &mut Extern<S, A, Q>,
     env: &Env,
@@ -307,7 +410,7 @@ pub fn try_add_supported_contract<S: Storage, A: Api, Q: Querier>(
     Ok(HandleResponse {
         messages: vec![],
         log: vec![],
-        data: Some( to_binary( &HandleAnswer::UpdateConfig {
+        data: Some( to_binary( &HandleAnswer::AddSupportedContract {
             status: ResponseStatus::Success
         })?),
     })
@@ -341,7 +444,7 @@ pub fn try_remove_supported_contract<S: Storage, A: Api, Q: Querier>(
     Ok(HandleResponse {
         messages: vec![],
         log: vec![],
-        data: Some( to_binary( &HandleAnswer::UpdateConfig {
+        data: Some( to_binary( &HandleAnswer::RemoveSupportedContract {
             status: ResponseStatus::Success
         })?),
     })
@@ -367,7 +470,7 @@ pub fn try_update_supported_contract<S: Storage, A: Api, Q: Querier>(
     Ok(HandleResponse {
         messages: vec![],
         log: vec![],
-        data: Some( to_binary( &HandleAnswer::UpdateConfig {
+        data: Some( to_binary( &HandleAnswer::UpdateSupportedContract {
             status: ResponseStatus::Success
         })?),
     })
@@ -405,7 +508,7 @@ pub fn try_add_admin_command<S: Storage, A: Api, Q: Querier>(
     Ok(HandleResponse {
         messages: vec![],
         log: vec![],
-        data: Some( to_binary( &HandleAnswer::UpdateConfig {
+        data: Some( to_binary( &HandleAnswer::AddAdminCommand {
             status: ResponseStatus::Success
         })?),
     })
@@ -434,7 +537,7 @@ pub fn try_remove_admin_command<S: Storage, A: Api, Q: Querier>(
     Ok(HandleResponse {
         messages: vec![],
         log: vec![],
-        data: Some( to_binary( &HandleAnswer::UpdateConfig {
+        data: Some( to_binary( &HandleAnswer::RemoveAdminCommand {
             status: ResponseStatus::Success
         })?),
     })
@@ -463,7 +566,7 @@ pub fn try_update_admin_command<S: Storage, A: Api, Q: Querier>(
     Ok(HandleResponse {
         messages: vec![],
         log: vec![],
-        data: Some( to_binary( &HandleAnswer::UpdateConfig {
+        data: Some( to_binary( &HandleAnswer::UpdateAdminCommand {
             status: ResponseStatus::Success
         })?),
     })

diff --git a/contracts/governance/src/query.rs b/contracts/governance/src/query.rs
@@ -1,6 +1,6 @@
 use cosmwasm_std::{Api, Extern, Querier, StdError, StdResult, Storage, Uint128};
 use shade_protocol::governance::{QueryAnswer, Proposal};
-use crate::state::{total_proposals_r, proposal_r, supported_contracts_list_r, admin_commands_list_r, supported_contract_r, admin_commands_r};
+use crate::state::{total_proposals_r, proposal_r, supported_contracts_list_r, admin_commands_list_r, supported_contract_r, admin_commands_r, total_proposal_votes_r};
 
 pub fn proposals<S: Storage, A: Api, Q: Querier>(
     deps: &Extern<S, A, Q>,
@@ -41,6 +41,14 @@ pub fn total_proposals<S: Storage, A: Api, Q: Querier>(
     })
 }
 
+pub fn proposal_votes<S: Storage, A: Api, Q: Querier>(
+    deps: &Extern<S, A, Q>, proposal_id: Uint128) -> StdResult<QueryAnswer> {
+
+    Ok(QueryAnswer::ProposalVotes {
+        status: total_proposal_votes_r(&deps.storage).load(proposal_id.to_string().as_bytes())?
+    })
+}
+
 pub fn supported_contracts<S: Storage, A: Api, Q: Querier>(
     deps: &Extern<S, A, Q>) -> StdResult<QueryAnswer> {
 

diff --git a/contracts/governance/src/state.rs b/contracts/governance/src/state.rs
@@ -4,12 +4,14 @@ use shade_protocol::{
     governance::{Config, Proposal},
     asset::Contract,
 };
-use shade_protocol::governance::AdminCommand;
+use shade_protocol::governance::{AdminCommand, UserVote, VoteTally};
 
 pub static CONFIG_KEY: &[u8] = b"config";
 pub static CONTRACT_KEY: &[u8] = b"supported_contracts";
 pub static CONTRACT_LIST_KEY: &[u8] = b"supported_contracts_list";
 pub static PROPOSAL_KEY: &[u8] = b"proposals";
+pub static PROPOSAL_VOTES_KEY: &str = "proposal_votes";
+pub static TOTAL_PROPOSAL_VOTES_KEY: &[u8] = b"total_proposal_votes";
 pub static TOTAL_PROPOSAL_KEY: &[u8] = b"total_proposals";
 pub static ADMIN_COMMANDS_KEY: &[u8] = b"admin_commands";
 pub static ADMIN_COMMANDS_LIST_KEY: &[u8] = b"admin_commands_list";
@@ -39,6 +41,23 @@ pub fn proposal_w<S: Storage>(storage: &mut S) -> Bucket<S, Proposal> {
     bucket(PROPOSAL_KEY, storage)
 }
 
+// Adds the proposal ID to have better range
+pub fn proposal_votes_r<S: Storage>(storage: & S, proposal: Uint128) -> ReadonlyBucket<S, VoteTally> {
+    bucket_read((proposal.to_string() + PROPOSAL_VOTES_KEY).as_bytes(), storage)
+}
+
+pub fn proposal_votes_w<S: Storage>(storage: &mut S, proposal: Uint128) -> Bucket<S, VoteTally> {
+    bucket((proposal.to_string() + PROPOSAL_VOTES_KEY).as_bytes(), storage)
+}
+
+pub fn total_proposal_votes_r<S: Storage>(storage: & S) -> ReadonlyBucket<S, VoteTally> {
+    bucket_read(TOTAL_PROPOSAL_VOTES_KEY, storage)
+}
+
+pub fn total_proposal_votes_w<S: Storage>(storage: &mut S) -> Bucket<S, VoteTally> {
+    bucket(TOTAL_PROPOSAL_VOTES_KEY, storage)
+}
+
 pub fn supported_contract_r<S: Storage>(storage: & S) -> ReadonlyBucket<S, Contract> {
     bucket_read(CONTRACT_KEY, storage)
 }

diff --git a/contracts/staking/src/contract.rs b/contracts/staking/src/contract.rs
@@ -13,10 +13,12 @@ use crate::{
 use secret_toolkit::snip20::register_receive_msg;
 use binary_heap_plus::{BinaryHeap, MinComparator};
 use shade_protocol::{staking::Unbonding, snip20};
-use crate::{handle::{try_update_unbond_time, try_stake, try_unbond, try_get_staker, try_get_stakers, try_trigger_unbounds},
+use crate::{handle::{try_update_config, try_stake, try_unbond, try_get_staker, try_get_stakers, try_trigger_unbounds},
             state::{unbonding_w}};
 use secret_toolkit::utils::HandleCallback;
 use crate::state::total_staked_w;
+use crate::handle::try_vote;
+use shade_protocol::asset::Contract;
 
 pub fn init<S: Storage, A: Api, Q: Querier>(
     deps: &mut Extern<S, A, Q>,
@@ -26,7 +28,7 @@ pub fn init<S: Storage, A: Api, Q: Querier>(
 
     let state = Config {
         admin: match msg.admin {
-            None => { env.message.sender.clone() }
+            None => { Contract { address: env.message.sender.clone(), code_hash: "".to_string() } }
             Some(admin) => { admin }
         },
         unbond_time: msg.unbond_time,
@@ -61,12 +63,14 @@ pub fn handle<S: Storage, A: Api, Q: Querier>(
     msg: HandleMsg,
 ) -> StdResult<HandleResponse> {
     match msg {
-        HandleMsg::UpdateUnbondTime { unbond_time
-        } => try_update_unbond_time(deps, &env, unbond_time),
+        HandleMsg::UpdateConfig { admin, unbond_time
+        } => try_update_config(deps, &env, admin, unbond_time),
         HandleMsg::Receive { sender, from, amount
         } => try_stake(deps, &env, sender, from, amount),
         HandleMsg::Unbond { amount
         } => try_unbond(deps, &env, amount),
+        HandleMsg::Vote { proposal_id, votes
+        } => try_vote(deps, &env, proposal_id, votes),
         HandleMsg::GetStaker { account
         } => try_get_staker(deps, &env, account),
         HandleMsg::GetStakers { accounts