Fixed HasAnyValidators to Fallback to Looking Into the BeaconState (#2453)

beacon-chain/db/validator.go

@@ -7,6 +7,7 @@ import (
 	"fmt"
 
 	"github.com/boltdb/bolt"
+	pb "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
 	"github.com/prysmaticlabs/prysm/shared/hashutil"
 )
 
@@ -101,19 +102,9 @@ func (db *BeaconDB) HasValidator(pubKey []byte) bool {
 	return exists
 }
 
-// HasAllValidators returns true if all validators in a list of public keys
-// are in the bucket.
-func (db *BeaconDB) HasAllValidators(pubKeys [][]byte) bool {
-	return db.hasValidators(pubKeys, true /* requireAll */)
-}
-
 // HasAnyValidators returns true if any validator in a list of public keys
 // are in the bucket.
-func (db *BeaconDB) HasAnyValidators(pubKeys [][]byte) bool {
-	return db.hasValidators(pubKeys, false /* requireAll */)
-}
-
-func (db *BeaconDB) hasValidators(pubKeys [][]byte, requireAll bool) bool {
+func (db *BeaconDB) HasAnyValidators(state *pb.BeaconState, pubKeys [][]byte) (bool, error) {
 	exists := false
 	// #nosec G104, similar to HasBlock, HasAttestation... etc
 	db.view(func(tx *bolt.Tx) error {
@@ -121,14 +112,23 @@ func (db *BeaconDB) hasValidators(pubKeys [][]byte, requireAll bool) bool {
 		for _, pk := range pubKeys {
 			h := hashutil.Hash(pk)
 			exists = bkt.Get(h[:]) != nil
-			if !exists && requireAll {
-				break
-			} else if exists && !requireAll {
-				break
-			}
+			break
 		}
 		return nil
 	})
 
-	return exists
+	if !exists {
+		for _, pubKey := range pubKeys {
+			for i := 0; i < len(state.ValidatorRegistry); i++ {
+				v := state.ValidatorRegistry[i]
+				if bytes.Equal(v.Pubkey, pubKey) {
+					if err := db.SaveValidatorIndex(pubKey, i); err != nil {
+						return false, err
+					}
+					exists = true
+				}
+			}
+		}
+	}
+	return exists, nil
 }

beacon-chain/db/validator_test.go

@@ -96,43 +96,6 @@ func TestHasValidator(t *testing.T) {
 	}
 }
 
-func TestHasAllValidators(t *testing.T) {
-	db := setupDB(t)
-	defer teardownDB(t, db)
-
-	knownPubKeys := [][]byte{
-		[]byte("pk1"),
-		[]byte("pk2"),
-	}
-	unknownPubKeys := [][]byte{
-		[]byte("pk3"),
-		[]byte("pk4"),
-	}
-
-	// Populate the db with some public key
-	if err := db.db.Update(func(tx *bolt.Tx) error {
-		bkt := tx.Bucket(validatorBucket)
-		for _, pk := range knownPubKeys {
-			h := hashutil.Hash(pk)
-			if err := bkt.Put(h[:], []byte("data")); err != nil {
-				return err
-			}
-		}
-
-		return nil
-	}); err != nil {
-		t.Fatal(err)
-	}
-
-	if !db.HasAllValidators(knownPubKeys) {
-		t.Error("Database did not have expected validators")
-	}
-
-	if db.HasAllValidators(append(knownPubKeys, unknownPubKeys...)) {
-		t.Error("Database returned true when there are pubkeys that did not exist")
-	}
-}
-
 func TestHasAnyValidator(t *testing.T) {
 	db := setupDB(t)
 	defer teardownDB(t, db)
@@ -161,11 +124,23 @@ func TestHasAnyValidator(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	if !db.HasAnyValidators(append(knownPubKeys, unknownPubKeys...)) {
+	beaconState := &pb.BeaconState{
+		ValidatorRegistry: []*pb.Validator{},
+	}
+
+	has, err := db.HasAnyValidators(beaconState, append(knownPubKeys, unknownPubKeys...))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !has {
 		t.Error("Database did not have expected validators")
 	}
 
-	if db.HasAnyValidators(unknownPubKeys) {
+	has, err = db.HasAnyValidators(beaconState, unknownPubKeys)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if has {
 		t.Error("Database returned true when there are only pubkeys that did not exist")
 	}
 }

beacon-chain/rpc/validator_server.go

@@ -31,9 +31,10 @@ type ValidatorServer struct {
 // beacon state, if not, then it creates a stream which listens for canonical states which contain
 // the validator with the public key as an active validator record.
 func (vs *ValidatorServer) WaitForActivation(req *pb.ValidatorActivationRequest, stream pb.ValidatorService_WaitForActivationServer) error {
+	beaconState, err := vs.beaconDB.HeadState(stream.Context())
 
 	reply := func() error {
-		beaconState, err := vs.beaconDB.HeadState(stream.Context())
+		beaconState, err = vs.beaconDB.HeadState(stream.Context())
 		if err != nil {
 			return fmt.Errorf("could not retrieve beacon state: %v", err)
 		}
@@ -45,13 +46,21 @@ func (vs *ValidatorServer) WaitForActivation(req *pb.ValidatorActivationRequest,
 
 	}
 
-	if vs.beaconDB.HasAnyValidators(req.PublicKeys) {
+	hasAny, err := vs.beaconDB.HasAnyValidators(beaconState, req.PublicKeys)
+	if err != nil {
+		return err
+	}
+	if hasAny {
 		return reply()
 	}
 	for {
 		select {
 		case <-time.After(3 * time.Second):
-			if !vs.beaconDB.HasAnyValidators(req.PublicKeys) {
+			hasAny, err := vs.beaconDB.HasAnyValidators(beaconState, req.PublicKeys)
+			if err != nil {
+				return err
+			}
+			if !hasAny {
 				continue
 			}
 			return reply()

beacon-chain/rpc/validator_server_test.go

@@ -574,7 +574,8 @@ func TestWaitForActivation_ContextClosed(t *testing.T) {
 	ctx := context.Background()
 
 	beaconState := &pbp2p.BeaconState{
-		Slot: params.BeaconConfig().GenesisSlot,
+		Slot:              params.BeaconConfig().GenesisSlot,
+		ValidatorRegistry: []*pbp2p.Validator{},
 	}
 	if err := db.SaveState(ctx, beaconState); err != nil {
 		t.Fatalf("could not save state: %v", err)
@@ -594,6 +595,7 @@ func TestWaitForActivation_ContextClosed(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()
 	mockStream := internal.NewMockValidatorService_WaitForActivationServer(ctrl)
+	mockStream.EXPECT().Context().Return(context.Background())
 	exitRoutine := make(chan bool)
 	go func(tt *testing.T) {
 		want := "context closed"
@@ -654,6 +656,7 @@ func TestWaitForActivation_ValidatorOriginallyExists(t *testing.T) {
 	defer ctrl.Finish()
 	mockStream := internal.NewMockValidatorService_WaitForActivationServer(ctrl)
 	mockStream.EXPECT().Context().Return(context.Background())
+	mockStream.EXPECT().Context().Return(context.Background())
 	mockStream.EXPECT().Send(
 		&pb.ValidatorActivationResponse{
 			ActivatedPublicKeys: pubKeys,