From e6d29552914b011b7598370277ce159034e72a8b Mon Sep 17 00:00:00 2001
From: Jiri Vestfal <45360788+MissingNO57@users.noreply.github.com>
Date: Fri, 25 Oct 2024 11:43:12 +0200
Subject: [PATCH] refactor: cudos limits config (#411)

---
 app/upgrade_cudos.go              | 12 ++++++++++--
 app/upgrade_cudos_distribution.go | 16 +++++++++++++---
 app/upgrade_network_config.go     | 17 ++++++++++++++---
 3 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/app/upgrade_cudos.go b/app/upgrade_cudos.go
index 7c968524..add353a3 100644
--- a/app/upgrade_cudos.go
+++ b/app/upgrade_cudos.go
@@ -1060,7 +1060,11 @@ func withdrawGenesisStakingDelegations(logger log.Logger, genesisData *GenesisDa
 	// Handle remaining bonded pool balance
 	bondedPool := genesisData.Accounts.MustGet(genesisData.BondedPoolAddress)
 
-	// TODO: Write to manifest?
+	maxToleratedRemainingStakingBalance := unwrapOrDefault(
+		cudosCfg.Config.MaxToleratedRemainingStakingBalance,
+		DefaultMaxToleratedRemainingStakingBalance,
+	)
+
 	err := checkTolerance(bondedPool.Balance, maxToleratedRemainingStakingBalance)
 	if err != nil {
 		return fmt.Errorf("remaining bonded pool balance %s is too high", bondedPool.Balance.String())
@@ -1078,7 +1082,6 @@ func withdrawGenesisStakingDelegations(logger log.Logger, genesisData *GenesisDa
 	// Handle remaining not-bonded pool balance
 	notBondedPool := genesisData.Accounts.MustGet(genesisData.NotBondedPoolAddress)
 
-	// TODO: Write to manifest?
 	err = checkTolerance(notBondedPool.Balance, maxToleratedRemainingStakingBalance)
 	if err != nil {
 		return fmt.Errorf("remaining not-bonded pool balance %s is too high", notBondedPool.Balance.String())
@@ -2262,6 +2265,11 @@ func MigrateGenesisAccounts(genesisData *GenesisData, ctx sdk.Context, app *App,
 	remainingMintBalance := app.BankKeeper.GetAllBalances(ctx, mintModuleAddr)
 	remainingMintBalance = remainingMintBalance.Sub(initialMintBalance)
 
+	maxToleratedRemainingMintBalance := unwrapOrDefault(
+		cudosCfg.Config.MaxToleratedRemainingMintBalance,
+		DefaultMaxToleratedRemainingMintBalance,
+	)
+
 	err = checkTolerance(remainingMintBalance, maxToleratedRemainingMintBalance)
 	if err != nil {
 		return err
diff --git a/app/upgrade_cudos_distribution.go b/app/upgrade_cudos_distribution.go
index e7d52c72..04879d31 100644
--- a/app/upgrade_cudos_distribution.go
+++ b/app/upgrade_cudos_distribution.go
@@ -360,7 +360,13 @@ func checkTolerance(coins sdk.Coins, maxToleratedDiff sdk.Int) error {
 	return nil
 }
 
-func verifyOutstandingBalances(genesisData *GenesisData) error {
+func verifyOutstandingBalances(genesisData *GenesisData, cudosCfg *CudosMergeConfig) error {
+
+	maxToleratedRemainingDistributionBalance := unwrapOrDefault(
+		cudosCfg.Config.MaxToleratedRemainingDistributionBalance,
+		DefaultMaxToleratedRemainingDistributionBalance,
+	)
+
 	for i := range genesisData.DistributionInfo.OutstandingRewards.Iterate() {
 		validatorAddr, validatorOutstandingReward := i.Key, i.Value
 
@@ -393,7 +399,7 @@ func withdrawGenesisDistributionRewards(logger log.Logger, genesisData *GenesisD
 	}
 
 	// Check that remaining balance is equal to AccumulatedCommissions
-	err := verifyOutstandingBalances(genesisData)
+	err := verifyOutstandingBalances(genesisData, cudosCfg)
 	if err != nil {
 		return err
 	}
@@ -416,7 +422,11 @@ func withdrawGenesisDistributionRewards(logger log.Logger, genesisData *GenesisD
 		logger.Info("cudos merge: remaining dist balance", "amount", remainingBalance.String())
 	}
 
-	// TODO: Write to manifest?
+	maxToleratedRemainingDistributionBalance := unwrapOrDefault(
+		cudosCfg.Config.MaxToleratedRemainingDistributionBalance,
+		DefaultMaxToleratedRemainingDistributionBalance,
+	)
+
 	err = checkTolerance(remainingBalance, maxToleratedRemainingDistributionBalance)
 	if err != nil {
 		return fmt.Errorf("remaining distribution balance %s is too high", remainingBalance.String())
diff --git a/app/upgrade_network_config.go b/app/upgrade_network_config.go
index 20911fdb..b8cbd3dd 100644
--- a/app/upgrade_network_config.go
+++ b/app/upgrade_network_config.go
@@ -9,11 +9,18 @@ import (
 )
 
 var (
-	maxToleratedRemainingDistributionBalance, _ = sdk.NewIntFromString("1000000000000000000")
-	maxToleratedRemainingStakingBalance, _      = sdk.NewIntFromString("100000000")
-	maxToleratedRemainingMintBalance, _         = sdk.NewIntFromString("100000000")
+	DefaultMaxToleratedRemainingDistributionBalance, _ = sdk.NewIntFromString("1000000000000000000")
+	DefaultMaxToleratedRemainingStakingBalance, _      = sdk.NewIntFromString("100000000")
+	DefaultMaxToleratedRemainingMintBalance, _         = sdk.NewIntFromString("100000000")
 )
 
+func unwrapOrDefault[T any](ptr *T, defaultValue T) T {
+	if ptr != nil {
+		return *ptr
+	}
+	return defaultValue
+}
+
 func newInt(val string) sdk.Int {
 	res, ok := sdk.NewIntFromString(val)
 	if !ok {
@@ -361,6 +368,10 @@ type CudosMergeConfigJSON struct {
 	ValidatorsMap []Pair[string, string] `json:"validators_map,omitempty"`
 
 	BackupValidators []string `json:"backup_validators,omitempty"`
+
+	MaxToleratedRemainingDistributionBalance *sdk.Int `json:"max_remaining_distribution_module_balance,omitempty"`
+	MaxToleratedRemainingStakingBalance      *sdk.Int `json:"max_remaining_staking_module_balance,omitempty"`
+	MaxToleratedRemainingMintBalance         *sdk.Int `json:"max_remaining_mint_module_balance,omitempty"`
 }
 
 type CudosMergeConfig struct {