Add IsWithinWeakSubjectivityPeriod helper method (#8706)

* Add IsWithinWeakSubjectivityPeriod helper method

* Nishant's suggestion on cache reset

Co-authored-by: prylabs-bulldozer[bot] <58059840+prylabs-bulldozer[bot]@users.noreply.github.com>

beacon-chain/core/helpers/weak_subjectivity.go

@@ -1,10 +1,12 @@
 package helpers
 
 import (
+	"bytes"
 	"errors"
 	"fmt"
 
 	types "github.com/prysmaticlabs/eth2-types"
+	eth "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
 	iface "github.com/prysmaticlabs/prysm/beacon-chain/state/interface"
 	"github.com/prysmaticlabs/prysm/shared/mathutil"
 	"github.com/prysmaticlabs/prysm/shared/params"
@@ -95,3 +97,47 @@ func ComputeWeakSubjectivityPeriod(st iface.ReadOnlyBeaconState) (types.Epoch, e
 
 	return types.Epoch(wsp), nil
 }
+
+// IsWithinWeakSubjectivityPeriod verifies if a given weak subjectivity checkpoint is not stale i.e.
+// the current node is so far beyond, that a given state and checkpoint are not for the latest weak
+// subjectivity point. Provided checkpoint still can be used to double-check that node's block root
+// at a given epoch matches that of the checkpoint.
+//
+// Reference implementation:
+// https://github.com/ethereum/eth2.0-specs/blob/master/specs/phase0/weak-subjectivity.md#checking-for-stale-weak-subjectivity-checkpoint
+
+// def is_within_weak_subjectivity_period(store: Store, ws_state: BeaconState, ws_checkpoint: Checkpoint) -> bool:
+//    # Clients may choose to validate the input state against the input Weak Subjectivity Checkpoint
+//    assert ws_state.latest_block_header.state_root == ws_checkpoint.root
+//    assert compute_epoch_at_slot(ws_state.slot) == ws_checkpoint.epoch
+//
+//    ws_period = compute_weak_subjectivity_period(ws_state)
+//    ws_state_epoch = compute_epoch_at_slot(ws_state.slot)
+//    current_epoch = compute_epoch_at_slot(get_current_slot(store))
+//    return current_epoch <= ws_state_epoch + ws_period
+func IsWithinWeakSubjectivityPeriod(
+	currentEpoch types.Epoch, wsState iface.ReadOnlyBeaconState, wsCheckpoint *eth.WeakSubjectivityCheckpoint) (bool, error) {
+	// Make sure that incoming objects are not nil.
+	if wsState == nil || wsState.LatestBlockHeader() == nil || wsCheckpoint == nil {
+		return false, errors.New("invalid weak subjectivity state or checkpoint")
+	}
+
+	// Assert that state and checkpoint have the same root and epoch.
+	if !bytes.Equal(wsState.LatestBlockHeader().StateRoot, wsCheckpoint.StateRoot) {
+		return false, fmt.Errorf("state (%#x) and checkpoint (%#x) roots do not match",
+			wsState.LatestBlockHeader().StateRoot, wsCheckpoint.StateRoot)
+	}
+	if SlotToEpoch(wsState.Slot()) != wsCheckpoint.Epoch {
+		return false, fmt.Errorf("state (%v) and checkpoint (%v) epochs do not match",
+			SlotToEpoch(wsState.Slot()), wsCheckpoint.Epoch)
+	}
+
+	// Compare given epoch to state epoch + weak subjectivity period.
+	wsPeriod, err := ComputeWeakSubjectivityPeriod(wsState)
+	if err != nil {
+		return false, fmt.Errorf("cannot compute weak subjectivity period: %w", err)
+	}
+	wsStateEpoch := SlotToEpoch(wsState.Slot())
+
+	return currentEpoch <= wsStateEpoch+wsPeriod, nil
+}

beacon-chain/core/helpers/weak_subjectivity_test.go

@@ -1,4 +1,4 @@
-package helpers
+package helpers_test
 
 import (
 	"fmt"
@@ -6,32 +6,16 @@ import (
 
 	types "github.com/prysmaticlabs/eth2-types"
 	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
-	"github.com/prysmaticlabs/prysm/beacon-chain/cache"
+	"github.com/prysmaticlabs/prysm/beacon-chain/core/helpers"
 	iface "github.com/prysmaticlabs/prysm/beacon-chain/state/interface"
-	"github.com/prysmaticlabs/prysm/beacon-chain/state/stateV0"
-	pb "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
+	"github.com/prysmaticlabs/prysm/shared/bytesutil"
 	"github.com/prysmaticlabs/prysm/shared/params"
+	"github.com/prysmaticlabs/prysm/shared/testutil"
 	"github.com/prysmaticlabs/prysm/shared/testutil/assert"
 	"github.com/prysmaticlabs/prysm/shared/testutil/require"
 )
 
 func TestWeakSubjectivity_ComputeWeakSubjectivityPeriod(t *testing.T) {
-	genState := func(valCount uint64, avgBalance uint64) iface.ReadOnlyBeaconState {
-		registry := make([]*ethpb.Validator, valCount)
-		for i := uint64(0); i < valCount; i++ {
-			registry[i] = &ethpb.Validator{
-				EffectiveBalance: avgBalance * 1e9,
-				ExitEpoch:        params.BeaconConfig().FarFutureEpoch,
-			}
-		}
-		beaconState, err := stateV0.InitializeFromProto(&pb.BeaconState{
-			Validators: registry,
-			Slot:       200,
-		})
-		require.NoError(t, err)
-
-		return beaconState
-	}
 	tests := []struct {
 		valCount   uint64
 		avgBalance uint64
@@ -62,10 +46,182 @@ func TestWeakSubjectivity_ComputeWeakSubjectivityPeriod(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(fmt.Sprintf("valCount: %d, avgBalance: %d", tt.valCount, tt.avgBalance), func(t *testing.T) {
 			// Reset committee cache - as we need to recalculate active validator set for each test.
-			committeeCache = cache.NewCommitteesCache()
-			got, err := ComputeWeakSubjectivityPeriod(genState(tt.valCount, tt.avgBalance))
+			helpers.ClearCache()
+			got, err := helpers.ComputeWeakSubjectivityPeriod(genState(t, tt.valCount, tt.avgBalance))
 			require.NoError(t, err)
 			assert.Equal(t, tt.want, got, "valCount: %v, avgBalance: %v", tt.valCount, tt.avgBalance)
 		})
 	}
 }
+func TestWeakSubjectivity_IsWithinWeakSubjectivityPeriod(t *testing.T) {
+	tests := []struct {
+		name            string
+		epoch           types.Epoch
+		genWsState      func() iface.ReadOnlyBeaconState
+		genWsCheckpoint func() *ethpb.WeakSubjectivityCheckpoint
+		want            bool
+		wantedErr       string
+	}{
+		{
+			name: "nil weak subjectivity state",
+			genWsState: func() iface.ReadOnlyBeaconState {
+				return nil
+			},
+			genWsCheckpoint: func() *ethpb.WeakSubjectivityCheckpoint {
+				return &ethpb.WeakSubjectivityCheckpoint{
+					BlockRoot: make([]byte, 32),
+					StateRoot: make([]byte, 32),
+					Epoch:     42,
+				}
+			},
+			wantedErr: "invalid weak subjectivity state or checkpoint",
+		},
+		{
+			name: "nil weak subjectivity checkpoint",
+			genWsState: func() iface.ReadOnlyBeaconState {
+				return genState(t, 128, 32)
+			},
+			genWsCheckpoint: func() *ethpb.WeakSubjectivityCheckpoint {
+				return nil
+			},
+			wantedErr: "invalid weak subjectivity state or checkpoint",
+		},
+		{
+			name: "state and checkpoint roots do not match",
+			genWsState: func() iface.ReadOnlyBeaconState {
+				beaconState := genState(t, 128, 32)
+				require.NoError(t, beaconState.SetSlot(42*params.BeaconConfig().SlotsPerEpoch))
+				err := beaconState.SetLatestBlockHeader(&ethpb.BeaconBlockHeader{
+					Slot:      42 * params.BeaconConfig().SlotsPerEpoch,
+					StateRoot: bytesutil.PadTo([]byte("stateroot1"), 32),
+				})
+				require.NoError(t, err)
+				return beaconState
+			},
+			genWsCheckpoint: func() *ethpb.WeakSubjectivityCheckpoint {
+				return &ethpb.WeakSubjectivityCheckpoint{
+					StateRoot: bytesutil.PadTo([]byte("stateroot2"), 32),
+					Epoch:     42,
+				}
+			},
+			wantedErr: fmt.Sprintf("state (%#x) and checkpoint (%#x) roots do not match",
+				bytesutil.PadTo([]byte("stateroot1"), 32), bytesutil.PadTo([]byte("stateroot2"), 32)),
+		},
+		{
+			name: "state and checkpoint epochs do not match",
+			genWsState: func() iface.ReadOnlyBeaconState {
+				beaconState := genState(t, 128, 32)
+				require.NoError(t, beaconState.SetSlot(42*params.BeaconConfig().SlotsPerEpoch))
+				err := beaconState.SetLatestBlockHeader(&ethpb.BeaconBlockHeader{
+					Slot:      42 * params.BeaconConfig().SlotsPerEpoch,
+					StateRoot: bytesutil.PadTo([]byte("stateroot"), 32),
+				})
+				require.NoError(t, err)
+				return beaconState
+			},
+			genWsCheckpoint: func() *ethpb.WeakSubjectivityCheckpoint {
+				return &ethpb.WeakSubjectivityCheckpoint{
+					StateRoot: bytesutil.PadTo([]byte("stateroot"), 32),
+					Epoch:     43,
+				}
+			},
+			wantedErr: "state (42) and checkpoint (43) epochs do not match",
+		},
+		{
+			name: "no active validators",
+			genWsState: func() iface.ReadOnlyBeaconState {
+				beaconState := genState(t, 0, 32)
+				require.NoError(t, beaconState.SetSlot(42*params.BeaconConfig().SlotsPerEpoch))
+				err := beaconState.SetLatestBlockHeader(&ethpb.BeaconBlockHeader{
+					Slot:      42 * params.BeaconConfig().SlotsPerEpoch,
+					StateRoot: bytesutil.PadTo([]byte("stateroot"), 32),
+				})
+				require.NoError(t, err)
+				return beaconState
+			},
+			genWsCheckpoint: func() *ethpb.WeakSubjectivityCheckpoint {
+				return &ethpb.WeakSubjectivityCheckpoint{
+					StateRoot: bytesutil.PadTo([]byte("stateroot"), 32),
+					Epoch:     42,
+				}
+			},
+			wantedErr: "cannot compute weak subjectivity period: no active validators found",
+		},
+		{
+			name:  "outside weak subjectivity period",
+			epoch: 300,
+			genWsState: func() iface.ReadOnlyBeaconState {
+				beaconState := genState(t, 128, 32)
+				require.NoError(t, beaconState.SetSlot(42*params.BeaconConfig().SlotsPerEpoch))
+				err := beaconState.SetLatestBlockHeader(&ethpb.BeaconBlockHeader{
+					Slot:      42 * params.BeaconConfig().SlotsPerEpoch,
+					StateRoot: bytesutil.PadTo([]byte("stateroot"), 32),
+				})
+				require.NoError(t, err)
+				return beaconState
+			},
+			genWsCheckpoint: func() *ethpb.WeakSubjectivityCheckpoint {
+				return &ethpb.WeakSubjectivityCheckpoint{
+					StateRoot: bytesutil.PadTo([]byte("stateroot"), 32),
+					Epoch:     42,
+				}
+			},
+			want: false,
+		},
+		{
+			name:  "within weak subjectivity period",
+			epoch: 299,
+			genWsState: func() iface.ReadOnlyBeaconState {
+				beaconState := genState(t, 128, 32)
+				require.NoError(t, beaconState.SetSlot(42*params.BeaconConfig().SlotsPerEpoch))
+				err := beaconState.SetLatestBlockHeader(&ethpb.BeaconBlockHeader{
+					Slot:      42 * params.BeaconConfig().SlotsPerEpoch,
+					StateRoot: bytesutil.PadTo([]byte("stateroot"), 32),
+				})
+				require.NoError(t, err)
+				return beaconState
+			},
+			genWsCheckpoint: func() *ethpb.WeakSubjectivityCheckpoint {
+				return &ethpb.WeakSubjectivityCheckpoint{
+					StateRoot: bytesutil.PadTo([]byte("stateroot"), 32),
+					Epoch:     42,
+				}
+			},
+			want: true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := helpers.IsWithinWeakSubjectivityPeriod(tt.epoch, tt.genWsState(), tt.genWsCheckpoint())
+			if tt.wantedErr != "" {
+				assert.Equal(t, false, got)
+				assert.ErrorContains(t, tt.wantedErr, err)
+				return
+			}
+			require.NoError(t, err)
+			assert.Equal(t, tt.want, got)
+		})
+	}
+}
+
+func genState(t *testing.T, valCount uint64, avgBalance uint64) iface.BeaconState {
+	beaconState, err := testutil.NewBeaconState()
+	require.NoError(t, err)
+
+	validators := make([]*ethpb.Validator, valCount)
+	balances := make([]uint64, len(validators))
+	for i := uint64(0); i < valCount; i++ {
+		validators[i] = &ethpb.Validator{
+			PublicKey:             make([]byte, params.BeaconConfig().BLSPubkeyLength),
+			WithdrawalCredentials: make([]byte, 32),
+			EffectiveBalance:      avgBalance * 1e9,
+			ExitEpoch:             params.BeaconConfig().FarFutureEpoch,
+		}
+		balances[i] = validators[i].EffectiveBalance
+	}
+
+	require.NoError(t, beaconState.SetValidators(validators))
+	require.NoError(t, beaconState.SetBalances(balances))
+
+	return beaconState
+}