From e8dc90a6231a616cb03dfef96e24d83ee1dd8c1b Mon Sep 17 00:00:00 2001
From: Potuz <potuz@prysmaticlabs.com>
Date: Wed, 6 Mar 2024 17:46:16 -0300
Subject: [PATCH] Fix UJ (#13688)

* Fix UJ

* gate slashed

* don't filter slashed for active balance

* don't overflow

* fix tests

* fix tests
---
 .../getters_participation_test.go             |  4 +--
 .../stateutil/unrealized_justification.go     | 28 +++++++++++--------
 .../unrealized_justification_test.go          |  2 --
 3 files changed, 18 insertions(+), 16 deletions(-)

diff --git a/beacon-chain/state/state-native/getters_participation_test.go b/beacon-chain/state/state-native/getters_participation_test.go
index 65c24cd7db16..c30fdaa80b76 100644
--- a/beacon-chain/state/state-native/getters_participation_test.go
+++ b/beacon-chain/state/state-native/getters_participation_test.go
@@ -19,7 +19,7 @@ func TestState_UnrealizedCheckpointBalances(t *testing.T) {
 		balances[i] = params.BeaconConfig().MaxEffectiveBalance
 	}
 	base := &ethpb.BeaconStateAltair{
-		Slot:        2,
+		Slot:        66,
 		RandaoMixes: make([][]byte, params.BeaconConfig().EpochsPerHistoricalVector),
 
 		Validators:                 validators,
@@ -57,7 +57,7 @@ func TestState_UnrealizedCheckpointBalances(t *testing.T) {
 	require.NoError(t, err)
 	active, previous, current, err = state.UnrealizedCheckpointBalances()
 	require.NoError(t, err)
-	require.Equal(t, allActive-params.BeaconConfig().MaxEffectiveBalance, active)
+	require.Equal(t, allActive, active)
 	require.Equal(t, uint64(0), current)
 	require.Equal(t, params.BeaconConfig().MaxEffectiveBalance, previous)
 
diff --git a/beacon-chain/state/stateutil/unrealized_justification.go b/beacon-chain/state/stateutil/unrealized_justification.go
index d6dc40a24277..81d9b8aa2e15 100644
--- a/beacon-chain/state/stateutil/unrealized_justification.go
+++ b/beacon-chain/state/stateutil/unrealized_justification.go
@@ -27,23 +27,27 @@ func UnrealizedCheckpointBalances(cp, pp []byte, validators ValReader, currentEp
 		if err != nil {
 			return 0, 0, 0, err
 		}
-		active := v.ActivationEpoch <= currentEpoch && currentEpoch < v.ExitEpoch
-		if active && !v.Slashed {
+		activeCurrent := v.ActivationEpoch <= currentEpoch && currentEpoch < v.ExitEpoch
+		if activeCurrent {
 			activeBalance, err = math.Add64(activeBalance, v.EffectiveBalance)
 			if err != nil {
 				return 0, 0, 0, err
 			}
-			if ((cp[i] >> targetIdx) & 1) == 1 {
-				currentTarget, err = math.Add64(currentTarget, v.EffectiveBalance)
-				if err != nil {
-					return 0, 0, 0, err
-				}
+		}
+		if v.Slashed {
+			continue
+		}
+		if activeCurrent && ((cp[i]>>targetIdx)&1) == 1 {
+			currentTarget, err = math.Add64(currentTarget, v.EffectiveBalance)
+			if err != nil {
+				return 0, 0, 0, err
 			}
-			if ((pp[i] >> targetIdx) & 1) == 1 {
-				prevTarget, err = math.Add64(prevTarget, v.EffectiveBalance)
-				if err != nil {
-					return 0, 0, 0, err
-				}
+		}
+		activePrevious := v.ActivationEpoch+1 <= currentEpoch && currentEpoch <= v.ExitEpoch
+		if activePrevious && ((pp[i]>>targetIdx)&1) == 1 {
+			prevTarget, err = math.Add64(prevTarget, v.EffectiveBalance)
+			if err != nil {
+				return 0, 0, 0, err
 			}
 		}
 	}
diff --git a/beacon-chain/state/stateutil/unrealized_justification_test.go b/beacon-chain/state/stateutil/unrealized_justification_test.go
index 1ba9b37f22c4..f6a3500c51a6 100644
--- a/beacon-chain/state/stateutil/unrealized_justification_test.go
+++ b/beacon-chain/state/stateutil/unrealized_justification_test.go
@@ -79,7 +79,6 @@ func TestState_UnrealizedCheckpointBalances(t *testing.T) {
 		validators[1].Slashed = true
 		active, previous, current, err := UnrealizedCheckpointBalances(cp, pp, NewValSliceReader(validators), 1)
 		require.NoError(tt, err)
-		expectedActive -= params.BeaconConfig().MaxEffectiveBalance
 		require.Equal(tt, expectedActive, active)
 		require.Equal(tt, params.BeaconConfig().MaxEffectiveBalance-params.BeaconConfig().MinDepositAmount, current)
 		require.Equal(tt, 2*params.BeaconConfig().MaxEffectiveBalance, previous)
@@ -169,7 +168,6 @@ func TestState_MVSlice_UnrealizedCheckpointBalances(t *testing.T) {
 		validators[1].Slashed = true
 		active, previous, current, err := UnrealizedCheckpointBalances(cp, pp, NewValMultiValueSliceReader(mv, &testObject{id: 0}), 1)
 		require.NoError(tt, err)
-		expectedActive -= params.BeaconConfig().MaxEffectiveBalance
 		require.Equal(tt, expectedActive, active)
 		require.Equal(tt, params.BeaconConfig().MaxEffectiveBalance-params.BeaconConfig().MinDepositAmount, current)
 		require.Equal(tt, 2*params.BeaconConfig().MaxEffectiveBalance, previous)