Prepare validator DB for attester protection implementation (#4584)

* Add flag for attester protection
* Merge branch 'master' of https://github.com/prysmaticlabs/Prysm into protecc-attester-db
* Merge branch 'master' of https://github.com/prysmaticlabs/Prysm into protecc-attester-db
* Remove flags
* Add attestation history DB functions to validator client
* Fix comments
* Update interface to new funcs
* Merge branch 'master' of https://github.com/prysmaticlabs/Prysm into protecc-attester-db
* Fix test
* Merge branch 'master' into protecc-attester-db

validator/db/BUILD.bazel

@@ -3,6 +3,7 @@ load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 go_library(
     name = "go_default_library",
     srcs = [
+        "attestation_history.go",
         "db.go",
         "proposal_history.go",
         "schema.go",
@@ -26,6 +27,7 @@ go_library(
 go_test(
     name = "go_default_test",
     srcs = [
+        "attestation_history_test.go",
         "proposal_history_test.go",
         "setup_db_test.go",
     ],

validator/db/attestation_history.go

@@ -0,0 +1,71 @@
+package db
+
+import (
+	"context"
+
+	"github.com/boltdb/bolt"
+	"github.com/gogo/protobuf/proto"
+	"github.com/pkg/errors"
+	slashpb "github.com/prysmaticlabs/prysm/proto/slashing"
+	"go.opencensus.io/trace"
+)
+
+func unmarshalAttestationHistory(enc []byte) (*slashpb.AttestationHistory, error) {
+	history := &slashpb.AttestationHistory{}
+	err := proto.Unmarshal(enc, history)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to unmarshal encoding")
+	}
+	return history, nil
+}
+
+// AttestationHistory accepts a validator public key and returns the corresponding attestation history.
+// Returns nil if there is no attestation history for the validator.
+func (db *Store) AttestationHistory(ctx context.Context, publicKey []byte) (*slashpb.AttestationHistory, error) {
+	ctx, span := trace.StartSpan(ctx, "Validator.AttestationHistory")
+	defer span.End()
+
+	var err error
+	var attestationHistory *slashpb.AttestationHistory
+	err = db.view(func(tx *bolt.Tx) error {
+		bucket := tx.Bucket(historicAttestationsBucket)
+		enc := bucket.Get(publicKey)
+		if enc == nil {
+			return nil
+		}
+		attestationHistory, err = unmarshalAttestationHistory(enc)
+		return err
+	})
+	return attestationHistory, err
+}
+
+// SaveAttestationHistory returns the attestation history for the requested validator public key.
+func (db *Store) SaveAttestationHistory(ctx context.Context, pubKey []byte, attestationHistory *slashpb.AttestationHistory) error {
+	ctx, span := trace.StartSpan(ctx, "Validator.SaveAttestationHistory")
+	defer span.End()
+
+	enc, err := proto.Marshal(attestationHistory)
+	if err != nil {
+		return errors.Wrap(err, "failed to encode attestation history")
+	}
+
+	err = db.update(func(tx *bolt.Tx) error {
+		bucket := tx.Bucket(historicAttestationsBucket)
+		return bucket.Put(pubKey, enc)
+	})
+	return err
+}
+
+// DeleteAttestationHistory deletes the attestation history for the corresponding validator public key.
+func (db *Store) DeleteAttestationHistory(ctx context.Context, pubkey []byte) error {
+	ctx, span := trace.StartSpan(ctx, "Validator.DeleteAttestationHistory")
+	defer span.End()
+
+	return db.update(func(tx *bolt.Tx) error {
+		bucket := tx.Bucket(historicAttestationsBucket)
+		if err := bucket.Delete(pubkey); err != nil {
+			return errors.Wrap(err, "failed to delete the attestation history")
+		}
+		return nil
+	})
+}

validator/db/attestation_history_test.go

@@ -0,0 +1,193 @@
+package db
+
+import (
+	"context"
+	"reflect"
+	"testing"
+
+	slashpb "github.com/prysmaticlabs/prysm/proto/slashing"
+	"github.com/prysmaticlabs/prysm/shared/params"
+)
+
+func TestAttestationHistory_InitializesNewPubKeys(t *testing.T) {
+	pubkeys := [][48]byte{[48]byte{30}, [48]byte{25}, [48]byte{20}}
+	db := SetupDB(t, pubkeys)
+	defer TeardownDB(t, db)
+
+	for _, pub := range pubkeys {
+		attestationHistory, err := db.AttestationHistory(context.Background(), pub[:])
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		newMap := make(map[uint64]uint64)
+		newMap[0] = params.BeaconConfig().FarFutureEpoch
+		clean := &slashpb.AttestationHistory{
+			TargetToSource: newMap,
+		}
+		if !reflect.DeepEqual(attestationHistory, clean) {
+			t.Fatalf("Expected attestation history epoch bits to be empty, received %v", attestationHistory)
+		}
+	}
+}
+
+func TestAttestationHistory_NilDB(t *testing.T) {
+	db := SetupDB(t, [][48]byte{})
+	defer TeardownDB(t, db)
+
+	valPubkey := []byte{1, 2, 3}
+
+	attestationHistory, err := db.AttestationHistory(context.Background(), valPubkey)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if attestationHistory != nil {
+		t.Fatalf("Expected attestation history to be nil, received: %v", attestationHistory)
+	}
+}
+
+func TestSaveAttestationHistory_OK(t *testing.T) {
+	db := SetupDB(t, [][48]byte{})
+	defer TeardownDB(t, db)
+
+	pubkey := []byte{3}
+	epoch := uint64(2)
+	farFuture := params.BeaconConfig().FarFutureEpoch
+	newMap := make(map[uint64]uint64)
+	// The validator attested at target epoch 2 but had no attestations for target epochs 0 and 1.
+	newMap[0] = farFuture
+	newMap[1] = farFuture
+	newMap[epoch] = 1
+	history := &slashpb.AttestationHistory{
+		TargetToSource:      newMap,
+		LatestEpochWritten: 2,
+	}
+
+	if err := db.SaveAttestationHistory(context.Background(), pubkey, history); err != nil {
+		t.Fatalf("Saving attestation history failed: %v", err)
+	}
+	savedHistory, err := db.AttestationHistory(context.Background(), pubkey)
+	if err != nil {
+		t.Fatalf("Failed to get attestation history: %v", err)
+	}
+
+	if savedHistory == nil || !reflect.DeepEqual(history, savedHistory) {
+		t.Fatalf("Expected DB to keep object the same, received: %v", history)
+	}
+	if savedHistory.TargetToSource[epoch] != newMap[epoch]{
+		t.Fatalf("Expected target epoch %d to have the same marked source epoch, received %d", epoch, savedHistory.TargetToSource[epoch])
+	}
+	if savedHistory.TargetToSource[epoch-1] != farFuture {
+		t.Fatalf("Expected target epoch %d to not be marked as attested for, received %d ", epoch-1, savedHistory.TargetToSource[epoch-1])
+	}
+	if savedHistory.TargetToSource[epoch-2] != farFuture {
+		t.Fatalf("Expected target epoch %d to not be marked as attested for, received %d", epoch-2, savedHistory.TargetToSource[epoch-2])
+	}
+}
+
+func TestSaveAttestationHistory_Overwrites(t *testing.T) {
+	db := SetupDB(t, [][48]byte{})
+	defer TeardownDB(t, db)
+	farFuture := params.BeaconConfig().FarFutureEpoch
+	newMap1 := make(map[uint64]uint64)
+	newMap1[0] = farFuture
+	newMap1[1] = 0
+	newMap2 := make(map[uint64]uint64)
+	newMap2[0] = farFuture
+	newMap2[1] = farFuture
+	newMap2[2] = 1
+	newMap3 := make(map[uint64]uint64)
+	newMap3[0] = farFuture
+	newMap3[1] = farFuture
+	newMap3[2] = farFuture
+	newMap3[3] = 2
+	tests := []struct {
+		pubkey  []byte
+		epoch   uint64
+		history *slashpb.AttestationHistory
+	}{
+		{
+			pubkey: []byte{0},
+			epoch:  uint64(1),
+			history: &slashpb.AttestationHistory{
+				TargetToSource: newMap1,
+				LatestEpochWritten: 1,
+			},
+		},
+		{
+			pubkey: []byte{0},
+			epoch:  uint64(2),
+			history: &slashpb.AttestationHistory{
+				TargetToSource: newMap2,
+				LatestEpochWritten: 2,
+			},
+		},
+		{
+			pubkey: []byte{0},
+			epoch:  uint64(3),
+			history: &slashpb.AttestationHistory{
+				TargetToSource:     newMap3,
+				LatestEpochWritten: 3,
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		if err := db.SaveAttestationHistory(context.Background(), tt.pubkey, tt.history); err != nil {
+			t.Fatalf("Saving attestation history failed: %v", err)
+		}
+		history, err := db.AttestationHistory(context.Background(), tt.pubkey)
+		if err != nil {
+			t.Fatalf("Failed to get attestation history: %v", err)
+		}
+
+		if history == nil || !reflect.DeepEqual(history, tt.history) {
+			t.Fatalf("Expected DB to keep object the same, received: %v", history)
+		}
+		if history.TargetToSource[tt.epoch] != tt.epoch - 1 {
+			t.Fatalf("Expected target epoch %d to be marked with correct source epoch %d", tt.epoch, history.TargetToSource[tt.epoch])
+		}
+		if history.TargetToSource[tt.epoch - 1] != farFuture {
+			t.Fatalf("Expected target epoch %d to not be marked as attested for, received %d", tt.epoch-1, history.TargetToSource[tt.epoch - 1])
+		}
+	}
+}
+
+func TestDeleteAttestationHistory_OK(t *testing.T) {
+	db := SetupDB(t, [][48]byte{})
+	defer TeardownDB(t, db)
+
+	pubkey := []byte{2}
+	newMap := make(map[uint64]uint64)
+	newMap[0] = params.BeaconConfig().FarFutureEpoch
+	newMap[1] = 0
+	history := &slashpb.AttestationHistory{
+		TargetToSource:     newMap,
+		LatestEpochWritten: 1,
+	}
+
+	if err := db.SaveAttestationHistory(context.Background(), pubkey, history); err != nil {
+		t.Fatalf("Save attestation history failed: %v", err)
+	}
+	// Making sure everything is saved.
+	savedHistory, err := db.AttestationHistory(context.Background(), pubkey)
+	if err != nil {
+		t.Fatalf("Failed to get attestation history: %v", err)
+	}
+	if savedHistory == nil || !reflect.DeepEqual(savedHistory, history) {
+		t.Fatalf("Expected DB to keep object the same, received: %v, expected %v", savedHistory, history)
+	}
+	if err := db.DeleteAttestationHistory(context.Background(), pubkey); err != nil {
+		t.Fatal(err)
+	}
+
+	// Check after deleting from DB.
+	savedHistory, err = db.AttestationHistory(context.Background(), pubkey)
+	if err != nil {
+		t.Fatalf("Failed to get attestation history: %v", err)
+	}
+	if savedHistory != nil {
+		t.Fatalf("Expected attestation history to be nil, received %v", savedHistory)
+	}
+}

validator/db/db.go

@@ -87,26 +87,42 @@ func NewKVStore(dirPath string, pubkeys [][48]byte) (*Store, error) {
 		return createBuckets(
 			tx,
 			historicProposalsBucket,
-			validatorsMinMaxSpanBucket,
+			historicAttestationsBucket,
 		)
 	}); err != nil {
 		return nil, err
 	}
 
 	// Initialize the required pubkeys into the DB to ensure they're not empty.
 	for _, pubkey := range pubkeys {
-		history, err := kv.ProposalHistory(context.Background(), pubkey[:])
+		proHistory, err := kv.ProposalHistory(context.Background(), pubkey[:])
 		if err != nil {
 			return nil, err
 		}
-		if history == nil {
+		if proHistory == nil {
 			cleanHistory := &slashpb.ProposalHistory{
 				EpochBits: bitfield.NewBitlist(params.BeaconConfig().WeakSubjectivityPeriod),
 			}
 			if err := kv.SaveProposalHistory(context.Background(), pubkey[:], cleanHistory); err != nil {
 				return nil, err
 			}
 		}
+
+		attHistory, err := kv.AttestationHistory(context.Background(), pubkey[:])
+		if err != nil {
+			return nil, err
+		}
+		if attHistory == nil {
+			newMap := make(map[uint64]uint64)
+			newMap[0] = params.BeaconConfig().FarFutureEpoch
+			cleanHistory := &slashpb.AttestationHistory{
+				TargetToSource: newMap,
+			}
+			if err := kv.SaveAttestationHistory(context.Background(), pubkey[:], cleanHistory); err != nil {
+				return nil, err
+			}
+		}
+
 	}
 
 	return kv, err

validator/db/iface/interface.go

@@ -17,4 +17,8 @@ type ValidatorDB interface {
 	ProposalHistory(ctx context.Context, publicKey []byte) (*slashpb.ProposalHistory, error)
 	SaveProposalHistory(ctx context.Context, publicKey []byte, history *slashpb.ProposalHistory) error
 	DeleteProposalHistory(ctx context.Context, publicKey []byte) error
+	// Attester protection related methods.
+	AttestationHistory(ctx context.Context, publicKey []byte) (*slashpb.AttestationHistory, error)
+	SaveAttestationHistory(ctx context.Context, publicKey []byte, history *slashpb.AttestationHistory) error
+	DeleteAttestationHistory(ctx context.Context, publicKey []byte) error
 }

validator/db/proposal_history_test.go

@@ -34,9 +34,9 @@ func TestProposalHistory_NilDB(t *testing.T) {
 	db := SetupDB(t, [][48]byte{})
 	defer TeardownDB(t, db)
 
-	balPubkey := []byte{1, 2, 3}
+	valPubkey := []byte{1, 2, 3}
 
-	proposalHistory, err := db.ProposalHistory(context.Background(), balPubkey)
+	proposalHistory, err := db.ProposalHistory(context.Background(), valPubkey)
 	if err != nil {
 		t.Fatal(err)
 	}

validator/db/schema.go

@@ -3,8 +3,6 @@ package db
 var (
 	// Validator slashing protection from double proposals.
 	historicProposalsBucket = []byte("proposal-history-bucket")
-	// In order to quickly detect surround and surrounded attestations we need to store
-	// the min and max span for each validator for each epoch.
-	// see https://github.com/protolambda/eth2-surround/blob/master/README.md#min-max-surround
-	validatorsMinMaxSpanBucket = []byte("validators-min-max-span-bucket")
+	// Validator slashing protection from slashable attestations.
+	historicAttestationsBucket = []byte("attestation-history-bucket")
 )