From 5a46f26e861a2debd3a228eca785acbf31a9fc1b Mon Sep 17 00:00:00 2001
From: rigelrozanski <rigel.rozanski@gmail.com>
Date: Wed, 23 May 2018 17:06:54 -0400
Subject: [PATCH] cleanup handleMsgUnbond and revoke logic

---
 x/stake/handler.go | 15 ++++-----------
 x/stake/keeper.go  | 35 ++++++++++++++++++++++++++---------
 2 files changed, 30 insertions(+), 20 deletions(-)

diff --git a/x/stake/handler.go b/x/stake/handler.go
index 017c4d5ad1..53653557cc 100644
--- a/x/stake/handler.go
+++ b/x/stake/handler.go
@@ -226,29 +226,22 @@ func handleMsgUnbond(ctx sdk.Context, msg MsgUnbond, k Keeper) sdk.Result {
 	// Add the coins
 	pool := k.GetPool(ctx)
 	validator, pool, returnAmount := validator.removeDelShares(pool, delShares)
+	k.setPool(ctx, pool)
 	returnCoins := sdk.Coins{{k.GetParams(ctx).BondDenom, returnAmount}}
 	k.coinKeeper.AddCoins(ctx, bond.DelegatorAddr, returnCoins)
 
 	/////////////////////////////////////
-
 	// revoke validator if necessary
 	if revokeCandidacy {
-
-		// change the share types to unbonded if they were not already
-		if validator.Status() == sdk.Bonded {
-			validator, pool = validator.UpdateStatus(pool, sdk.Unbonded)
-		}
-
 		validator.Revoked = true
 	}
 
-	// deduct shares from the validator
+	validator = k.updateValidator(ctx, validator)
+
 	if validator.DelegatorShares.IsZero() {
 		k.removeValidator(ctx, validator.Owner)
-	} else {
-		k.updateValidator(ctx, validator)
 	}
-	k.setPool(ctx, pool)
+
 	tags := sdk.NewTags("action", []byte("unbond"), "delegator", msg.DelegatorAddr.Bytes(), "validator", msg.ValidatorAddr.Bytes())
 	return sdk.Result{
 		Tags: tags,
diff --git a/x/stake/keeper.go b/x/stake/keeper.go
index 1e8697e1df..ada19df9f0 100644
--- a/x/stake/keeper.go
+++ b/x/stake/keeper.go
@@ -202,6 +202,11 @@ func (k Keeper) updateValidator(ctx sdk.Context, validator Validator) Validator
 	// retreive the old validator record
 	oldValidator, oldFound := k.GetValidator(ctx, ownerAddr)
 
+	if validator.Revoked && oldValidator.Status() == sdk.Bonded {
+		validator, pool = validator.UpdateStatus(pool, sdk.Unbonded)
+		k.setPool(ctx, pool)
+	}
+
 	powerIncreasing := false
 	if oldFound && oldValidator.PoolShares.Bonded().LT(validator.PoolShares.Bonded()) {
 		powerIncreasing = true
@@ -227,7 +232,7 @@ func (k Keeper) updateValidator(ctx sdk.Context, validator Validator) Validator
 
 	// efficiency case:
 	// if already bonded and power increasing only need to update tendermint
-	if powerIncreasing && oldValidator.Status() == sdk.Bonded {
+	if powerIncreasing && !validator.Revoked && oldValidator.Status() == sdk.Bonded {
 		bz := k.cdc.MustMarshalBinary(validator.abciValidator(k.cdc))
 		store.Set(GetTendermintUpdatesKey(ownerAddr), bz)
 		return validator
@@ -274,13 +279,13 @@ func (k Keeper) updateBondedValidators(ctx sdk.Context, store sdk.KVStore,
 	// add the actual validator power sorted store
 	maxValidators := k.GetParams(ctx).MaxValidators
 	iterator := store.ReverseSubspaceIterator(ValidatorsByPowerKey) // largest to smallest
-	i := 0
+	bondedValidatorsCount := 0
 	var validator Validator
-	for ; ; i++ {
-		if !iterator.Valid() || i > int(maxValidators-1) {
+	for {
+		if !iterator.Valid() || bondedValidatorsCount > int(maxValidators-1) {
 
 			// TODO benchmark if we should read the current power and not write if it's the same
-			if i-1 == int(maxValidators-1) {
+			if bondedValidatorsCount == int(maxValidators) { // is cliff validator
 				k.setCliffValidator(ctx, validator, k.GetPool(ctx))
 			}
 			iterator.Close()
@@ -313,6 +318,12 @@ func (k Keeper) updateBondedValidators(ctx sdk.Context, store sdk.KVStore,
 			}
 		}
 
+		if validator.Revoked && validator.Status() == sdk.Bonded {
+			panic(fmt.Sprintf("revoked validator cannot be bonded, address: %v\n", ownerAddr))
+		} else {
+			bondedValidatorsCount++
+		}
+
 		iterator.Next()
 	}
 
@@ -342,12 +353,12 @@ func (k Keeper) updateBondedValidatorsFull(ctx sdk.Context, store sdk.KVStore) {
 	// add the actual validator power sorted store
 	maxValidators := k.GetParams(ctx).MaxValidators
 	iterator = store.ReverseSubspaceIterator(ValidatorsByPowerKey) // largest to smallest
-	i := 0
+	bondedValidatorsCount := 0
 	var validator Validator
-	for ; ; i++ {
-		if !iterator.Valid() || i > int(maxValidators-1) {
+	for {
+		if !iterator.Valid() || bondedValidatorsCount > int(maxValidators-1) {
 
-			if i-1 == int(maxValidators-1) {
+			if bondedValidatorsCount == int(maxValidators) { // is cliff validator
 				k.setCliffValidator(ctx, validator, k.GetPool(ctx))
 			}
 			iterator.Close()
@@ -376,6 +387,12 @@ func (k Keeper) updateBondedValidatorsFull(ctx sdk.Context, store sdk.KVStore) {
 			validator = k.bondValidator(ctx, store, validator)
 		}
 
+		if validator.Revoked && validator.Status() == sdk.Bonded {
+			panic(fmt.Sprintf("revoked validator cannot be bonded, address: %v\n", ownerAddr))
+		} else {
+			bondedValidatorsCount++
+		}
+
 		iterator.Next()
 	}