From 89e21733e95e6b3e1f32f52734d75f6bd256ffc4 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

(cherry picked from commit b0a2115a2600c3e382e3d0dae47c842858add710)
---
 .../getters_participation_test.go             |  4 +--
 .../stateutil/unrealized_justification.go     | 28 +++++++++++--------
 .../unrealized_justification_test.go          |  1 -
 3 files changed, 18 insertions(+), 15 deletions(-)

diff --git a/beacon-chain/state/state-native/getters_participation_test.go b/beacon-chain/state/state-native/getters_participation_test.go
index 65c24cd7db1..c30fdaa80b7 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 d1a427c05ac..e20f76a7bce 100644
--- a/beacon-chain/state/stateutil/unrealized_justification.go
+++ b/beacon-chain/state/stateutil/unrealized_justification.go
@@ -24,23 +24,27 @@ func UnrealizedCheckpointBalances(cp, pp []byte, validators []*ethpb.Validator,
 
 	var err error
 	for i, v := range validators {
-		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 ccd064d4cc7..b5baa25d42e 100644
--- a/beacon-chain/state/stateutil/unrealized_justification_test.go
+++ b/beacon-chain/state/stateutil/unrealized_justification_test.go
@@ -78,7 +78,6 @@ func TestState_UnrealizedCheckpointBalances(t *testing.T) {
 		validators[1].Slashed = true
 		active, previous, current, err := UnrealizedCheckpointBalances(cp, pp, 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)