Total mix stake accounting (#794)

* Feature/persistent gateway storage (#784)

* Sqlx struct stub

* Initial schema

* Initial error enum

* Managed for persisted shared keys

* Initial inbox manager

* Comments

* Using new database in clients handler

* Extending gateway storage API

* tokio::main + placeholder values

* Removed old client store

* Simplified logic of async packet processing

* Renamed table + not null restriction

* BandwidthManager

* Removed sled dependency

* Using centralised storage for bandwidth

* Dead code removal

* WIP connection_handler split and simplification

Maybe it doesn't look like it right now, but once completed it will remove bunch of redundant checks for Nones etc

* Further more explicit clients handler split

* Minor cleanup

* Temporary store for active client handles

* Fixed error types

* Error trait on iv and encrypted address

* Authentication and registration moved to the handler

* Removal of clients handler

* Further logic simplification + returned explicit bandwidth values

* Further cleanup and comments

* Updated config with relevant changes

* Basic bandwidth tracking in client

* FreshHandle doc comments + fixed stagger issue

* Removed side-effects from .map

* More doc comments

* Database migration on build

* Increased default claimed bandwidth

* Renaming

* Fixed client determining available bandwidth

* Removed dead sql table that might be used in the future

* Windows workaround

* Comment

* Return error rather than cap credential

* Feature/migrate hidden delegations (#786)

* Remove migration code

* Added function to iterate over delegation of variable type

* Add unit tests

* Refactored some naming and reused mix/gateway functionality

* Borrow bucket instead of move

* Linked with existing delegations function

* Migration of left-over delegations

* Remove unused imports

* Put a gateway test as well, next to the mix one

* Expose queries for all delegations

* Change break point

* Added client side calls to the new queries

* Fix clippy

* Added pagination and read check tests

* Fix gateway test from the last commit

* Test functions for (de)serialization of identity and owner (in)to storage keys

* Add delegation function unit test

* Feature guard import

* Changed UnpackedDelegation from type to struct

* Remove mutable parameter and put start_after in returned value

* Made all delegations into iterator for OOM safety

* Fix clippy

* Add test for delegations iterator size in memory

* Change map with if let for ease of read

* Use DENOM instead of hardcoded value

* Total mixnode stake accounting in the contract

* GetTotalMixState query

* Add total_mix_stake to validator api

* cargo fmt

* Handle errors properly

* Gateway stake accounting

* tests

* Naming :)

Co-authored-by: Jędrzej Stuczyński <jedrzej.stuczynski@gmail.com>
Co-authored-by: Bogdan-Ștefan Neacşu <bogdan@nymtech.net>
Co-authored-by: Drazen Urch <durch@users.noreply.guthub.com>

common/client-libs/validator-client/src/client.rs

@@ -182,6 +182,20 @@ impl<C> Client<C> {
         Ok(self.nymd.get_state_params().await?)
     }
 
+    pub async fn get_total_mix_stake(&self) -> Result<u128, ValidatorClientError>
+    where
+        C: CosmWasmClient + Sync,
+    {
+        Ok(self.nymd.get_total_mix_stake().await?.u128())
+    }
+
+    pub async fn get_total_gateway_stake(&self) -> Result<u128, ValidatorClientError>
+    where
+        C: CosmWasmClient + Sync,
+    {
+        Ok(self.nymd.get_total_gateway_stake().await?.u128())
+    }
+
     // basically handles paging for us
     pub async fn get_all_nymd_mixnodes(&self) -> Result<Vec<MixNodeBond>, ValidatorClientError>
     where

common/client-libs/validator-client/src/nymd/mod.rs

@@ -12,6 +12,7 @@ use crate::nymd::wallet::DirectSecp256k1HdWallet;
 use cosmrs::rpc::endpoint::broadcast;
 use cosmrs::rpc::{Error as TendermintRpcError, HttpClientUrl};
 use cosmwasm_std::Coin;
+use cosmwasm_std::{Coin, Uint128};
 use mixnet_contract::{
     Addr, Delegation, ExecuteMsg, Gateway, GatewayOwnershipResponse, IdentityKey,
     LayerDistribution, MixNode, MixOwnershipResponse, PagedAllDelegationsResponse,
@@ -212,6 +213,26 @@ impl<C> NymdClient<C> {
             .await
     }
 
+    pub async fn get_total_mix_stake(&self) -> Result<Uint128, NymdError>
+    where
+        C: CosmWasmClient + Sync,
+    {
+        let request = QueryMsg::GetTotalMixStake {};
+        self.client
+            .query_contract_smart(self.contract_address()?, &request)
+            .await
+    }
+
+    pub async fn get_total_gateway_stake(&self) -> Result<Uint128, NymdError>
+    where
+        C: CosmWasmClient + Sync,
+    {
+        let request = QueryMsg::GetTotalGatewayStake {};
+        self.client
+            .query_contract_smart(self.contract_address()?, &request)
+            .await
+    }
+
     /// Checks whether there is a bonded mixnode associated with the provided client's address
     pub async fn owns_mixnode(&self, address: &AccountId) -> Result<bool, NymdError>
     where

common/mixnet-contract/src/gateway.rs

@@ -59,6 +59,10 @@ impl GatewayBond {
     pub fn gateway(&self) -> &Gateway {
         &self.gateway
     }
+
+    pub fn total_delegation(&self) -> Coin {
+        self.total_delegation.clone()
+    }
 }
 
 impl PartialOrd for GatewayBond {

common/mixnet-contract/src/mixnode.rs

@@ -79,6 +79,10 @@ impl MixNodeBond {
     pub fn mix_node(&self) -> &MixNode {
         &self.mix_node
     }
+
+    pub fn total_delegation(&self) -> Coin {
+        self.total_delegation.clone()
+    }
 }
 
 impl PartialOrd for MixNodeBond {

common/mixnet-contract/src/msg.rs

@@ -107,6 +107,8 @@ pub enum QueryMsg {
         address: Addr,
     },
     LayerDistribution {},
+    GetTotalMixStake {},
+    GetTotalGatewayStake {},
 }
 
 #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)]

contracts/mixnet/src/contract.rs

@@ -207,13 +207,17 @@ pub fn query(deps: Deps, _env: Env, msg: QueryMsg) -> Result<QueryResponse, Cont
             gateway_identity,
             address,
         )?),
+        QueryMsg::GetTotalMixStake {} => to_binary(&queries::query_total_mix_stake(deps)),
+        QueryMsg::GetTotalGatewayStake {} => to_binary(&queries::query_total_gt_stake(deps)),
     };
 
     Ok(query_res?)
 }
-
 #[entry_point]
 pub fn migrate(_deps: DepsMut, _env: Env, _msg: MigrateMsg) -> Result<Response, ContractError> {
+    todo!("Calculate initial total mix and gateway stake after initial deployment");
+
+    #[allow(unreachable_code)]
     Ok(Default::default())
 }
 

contracts/mixnet/src/error.rs

@@ -89,4 +89,7 @@ pub enum ContractError {
         identity: IdentityKey,
         address: Addr,
     },
+
+    #[error("Overflow error!")]
+    Overflow(#[from] cosmwasm_std::OverflowError),
 }

contracts/mixnet/src/queries.rs

@@ -7,10 +7,10 @@ use crate::storage::{
     all_gateway_delegations_read, all_mix_delegations_read, gateway_delegations_read,
     gateways_owners_read, gateways_read, mix_delegations_read, mixnodes_owners_read, mixnodes_read,
     read_layer_distribution, read_state_params, reverse_gateway_delegations_read,
-    reverse_mix_delegations_read,
+    reverse_mix_delegations_read, total_gateway_stake_value, total_mix_stake_value,
 };
 use config::defaults::DENOM;
-use cosmwasm_std::{coin, Addr, Deps, Order, StdResult};
+use cosmwasm_std::{coin, Addr, Deps, Order, StdResult, Uint128};
 use mixnet_contract::{
     Delegation, GatewayBond, GatewayOwnershipResponse, IdentityKey, LayerDistribution, MixNodeBond,
     MixOwnershipResponse, PagedAllDelegationsResponse, PagedGatewayDelegationsResponse,
@@ -93,6 +93,14 @@ pub(crate) fn query_layer_distribution(deps: Deps) -> LayerDistribution {
     read_layer_distribution(deps.storage)
 }
 
+pub(crate) fn query_total_mix_stake(deps: Deps) -> Uint128 {
+    total_mix_stake_value(deps.storage)
+}
+
+pub(crate) fn query_total_gt_stake(deps: Deps) -> Uint128 {
+    total_gateway_stake_value(deps.storage)
+}
+
 /// Adds a 0 byte to terminate the `start_after` value given. This allows CosmWasm
 /// to get the succeeding key as the start of the next page.
 // S works for both `String` and `Addr` and that's what we wanted

contracts/mixnet/src/storage.rs

@@ -1,9 +1,9 @@
 // Copyright 2021 - Nym Technologies SA <contact@nymtech.net>
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::queries;
 use crate::state::State;
 use crate::transactions::MINIMUM_BLOCK_AGE_FOR_REWARDING;
+use crate::{error::ContractError, queries};
 use cosmwasm_std::{Decimal, Order, StdResult, Storage, Uint128};
 use cosmwasm_storage::{
     bucket, bucket_read, singleton, singleton_read, Bucket, ReadonlyBucket, ReadonlySingleton,
@@ -24,6 +24,9 @@ use serde::Serialize;
 // singletons
 const CONFIG_KEY: &[u8] = b"config";
 const LAYER_DISTRIBUTION_KEY: &[u8] = b"layers";
+// Keeps total amount of stake towards mixnodes and gateways. Removing a bond removes all its delegations from the total, the reverse is true for adding a bond.
+const TOTAL_MIX_STAKE_KEY: &[u8] = b"total_mn";
+const TOTAL_GATEWAY_STAKE_KEY: &[u8] = b"total_gt";
 
 // buckets
 const PREFIX_MIXNODES: &[u8] = b"mn";
@@ -48,6 +51,72 @@ pub fn config_read(storage: &dyn Storage) -> ReadonlySingleton<State> {
     singleton_read(storage, CONFIG_KEY)
 }
 
+fn total_mix_stake(storage: &dyn Storage) -> ReadonlySingleton<Uint128> {
+    singleton_read(storage, TOTAL_MIX_STAKE_KEY)
+}
+
+pub fn mut_total_mix_stake(storage: &mut dyn Storage) -> Singleton<Uint128> {
+    singleton(storage, TOTAL_MIX_STAKE_KEY)
+}
+
+fn total_gateway_stake(storage: &dyn Storage) -> ReadonlySingleton<Uint128> {
+    singleton_read(storage, TOTAL_GATEWAY_STAKE_KEY)
+}
+
+pub fn mut_total_gateway_stake(storage: &mut dyn Storage) -> Singleton<Uint128> {
+    singleton(storage, TOTAL_GATEWAY_STAKE_KEY)
+}
+
+pub fn total_mix_stake_value(storage: &dyn Storage) -> Uint128 {
+    match total_mix_stake(storage).load() {
+        Ok(value) => value,
+        Err(_e) => Uint128(0),
+    }
+}
+
+pub fn total_gateway_stake_value(storage: &dyn Storage) -> Uint128 {
+    match total_gateway_stake(storage).load() {
+        Ok(value) => value,
+        Err(_e) => Uint128(0),
+    }
+}
+
+pub fn incr_total_mix_stake(
+    amount: Uint128,
+    storage: &mut dyn Storage,
+) -> Result<Uint128, ContractError> {
+    let stake = total_mix_stake_value(storage).checked_add(amount)?;
+    mut_total_mix_stake(storage).save(&stake)?;
+    Ok(stake)
+}
+
+pub fn decr_total_mix_stake(
+    amount: Uint128,
+    storage: &mut dyn Storage,
+) -> Result<Uint128, ContractError> {
+    let stake = total_mix_stake_value(storage).checked_sub(amount)?;
+    mut_total_mix_stake(storage).save(&stake)?;
+    Ok(stake)
+}
+
+pub fn incr_total_gateway_stake(
+    amount: Uint128,
+    storage: &mut dyn Storage,
+) -> Result<Uint128, ContractError> {
+    let stake = total_gateway_stake_value(storage).checked_add(amount)?;
+    mut_total_gateway_stake(storage).save(&stake)?;
+    Ok(stake)
+}
+
+pub fn decr_total_gateway_stake(
+    amount: Uint128,
+    storage: &mut dyn Storage,
+) -> Result<Uint128, ContractError> {
+    let stake = total_gateway_stake_value(storage).checked_sub(amount)?;
+    mut_total_gateway_stake(storage).save(&stake)?;
+    Ok(stake)
+}
+
 pub(crate) fn read_state_params(storage: &dyn Storage) -> StateParams {
     // note: In any other case, I wouldn't have attempted to unwrap this result, but in here
     // if we fail to load the stored state we would already be in the undefined behaviour land,