From ea1043999ffbc66577dfd026366bb5ca39ea4808 Mon Sep 17 00:00:00 2001
From: Rongjian Lan <rongjian.lan@gmail.com>
Date: Thu, 23 Apr 2020 15:54:56 -0700
Subject: [PATCH] [audit] Fix and make exact 7 epoch for undelegation (#2874)

* Fix and make exact 7 epoch for undelegation

* Make slash explicit on reference
---
 staking/slash/double-sign.go     |  4 +-
 staking/types/delegation.go      |  4 +-
 staking/types/delegation_test.go | 64 ++++++++++++++++++++++++++++++--
 3 files changed, 65 insertions(+), 7 deletions(-)

diff --git a/staking/slash/double-sign.go b/staking/slash/double-sign.go
index c28977d16..a32ce26fc 100644
--- a/staking/slash/double-sign.go
+++ b/staking/slash/double-sign.go
@@ -352,7 +352,8 @@ func delegatorSlashApply(
 				}
 
 				// NOTE Assume did as much as could above, now check the undelegations
-				for _, undelegate := range delegationNow.Undelegations {
+				for i := range delegationNow.Undelegations {
+					undelegate := delegationNow.Undelegations[i]
 					// the epoch matters, only those undelegation
 					// such that epoch>= doubleSignEpoch should be slashable
 					if undelegate.Epoch.Cmp(doubleSignEpoch) >= 0 {
@@ -371,7 +372,6 @@ func delegatorSlashApply(
 						}
 
 						if nowAmt.Cmp(common.Big0) == 0 {
-							// TODO(audit): need to remove the undelegate
 							utils.Logger().Info().
 								RawJSON("delegation-snapshot", []byte(delegationSnapshot.String())).
 								RawJSON("delegation-current", []byte(delegationNow.String())).
diff --git a/staking/types/delegation.go b/staking/types/delegation.go
index cc59e45ab..efd942fc6 100644
--- a/staking/types/delegation.go
+++ b/staking/types/delegation.go
@@ -181,8 +181,8 @@ func (d *Delegation) RemoveUnlockedUndelegations(
 	totalWithdraw := big.NewInt(0)
 	count := 0
 	for j := range d.Undelegations {
-		if big.NewInt(0).Sub(curEpoch, d.Undelegations[j].Epoch).Int64() > LockPeriodInEpoch ||
-			big.NewInt(0).Sub(curEpoch, lastEpochInCommittee).Int64() > LockPeriodInEpoch {
+		if big.NewInt(0).Sub(curEpoch, d.Undelegations[j].Epoch).Int64() >= LockPeriodInEpoch ||
+			big.NewInt(0).Sub(curEpoch, lastEpochInCommittee).Int64() >= LockPeriodInEpoch {
 			// need to wait at least 7 epochs to withdraw; or the validator has been out of committee for 7 epochs
 			totalWithdraw.Add(totalWithdraw, d.Undelegations[j].Amount)
 			count++
diff --git a/staking/types/delegation_test.go b/staking/types/delegation_test.go
index 6e00f0c90..aacc509f5 100644
--- a/staking/types/delegation_test.go
+++ b/staking/types/delegation_test.go
@@ -67,16 +67,74 @@ func TestDeleteEntry(t *testing.T) {
 	}
 }
 
-func TestRemoveUnlockUndelegations(t *testing.T) {
-	lastEpochInCommitte := big.NewInt(16)
+func TestUnlockedLastEpochInCommittee(t *testing.T) {
+	lastEpochInCommittee := big.NewInt(17)
 	curEpoch := big.NewInt(24)
 
 	epoch4 := big.NewInt(21)
 	amount4 := big.NewInt(4000)
 	delegation.Undelegate(epoch4, amount4)
 
-	result := delegation.RemoveUnlockedUndelegations(curEpoch, lastEpochInCommitte)
+	result := delegation.RemoveUnlockedUndelegations(curEpoch, lastEpochInCommittee)
 	if result.Cmp(big.NewInt(8000)) != 0 {
 		t.Errorf("removing an unlocked undelegation fails")
 	}
 }
+
+func TestUnlockedLastEpochInCommitteeFail(t *testing.T) {
+	delegation := NewDelegation(delegatorAddr, delegationAmt)
+	lastEpochInCommittee := big.NewInt(18)
+	curEpoch := big.NewInt(24)
+
+	epoch4 := big.NewInt(21)
+	amount4 := big.NewInt(4000)
+	delegation.Undelegate(epoch4, amount4)
+
+	result := delegation.RemoveUnlockedUndelegations(curEpoch, lastEpochInCommittee)
+	if result.Cmp(big.NewInt(0)) != 0 {
+		t.Errorf("premature delegation shouldn't be unlocked")
+	}
+}
+
+func TestUnlockedFullPeriod(t *testing.T) {
+	lastEpochInCommittee := big.NewInt(34)
+	curEpoch := big.NewInt(34)
+
+	epoch5 := big.NewInt(27)
+	amount5 := big.NewInt(4000)
+	delegation.Undelegate(epoch5, amount5)
+
+	result := delegation.RemoveUnlockedUndelegations(curEpoch, lastEpochInCommittee)
+	if result.Cmp(big.NewInt(4000)) != 0 {
+		t.Errorf("removing an unlocked undelegation fails")
+	}
+}
+
+func TestUnlockedFullPeriodFail(t *testing.T) {
+	delegation := NewDelegation(delegatorAddr, delegationAmt)
+	lastEpochInCommittee := big.NewInt(34)
+	curEpoch := big.NewInt(34)
+
+	epoch5 := big.NewInt(28)
+	amount5 := big.NewInt(4000)
+	delegation.Undelegate(epoch5, amount5)
+
+	result := delegation.RemoveUnlockedUndelegations(curEpoch, lastEpochInCommittee)
+	if result.Cmp(big.NewInt(0)) != 0 {
+		t.Errorf("premature delegation shouldn't be unlocked")
+	}
+}
+
+func TestUnlockedPremature(t *testing.T) {
+	lastEpochInCommittee := big.NewInt(44)
+	curEpoch := big.NewInt(44)
+
+	epoch6 := big.NewInt(42)
+	amount6 := big.NewInt(4000)
+	delegation.Undelegate(epoch6, amount6)
+
+	result := delegation.RemoveUnlockedUndelegations(curEpoch, lastEpochInCommittee)
+	if result.Cmp(big.NewInt(0)) != 0 {
+		t.Errorf("premature delegation shouldn't be unlocked")
+	}
+}