From fff6472a040edf8e8d4eede7a92057a96d545a09 Mon Sep 17 00:00:00 2001
From: Raul Jordan <raul@prysmaticlabs.com>
Date: Wed, 13 Jan 2021 17:23:29 -0600
Subject: [PATCH] Allow Multiple Targets Per Source Epoch in Attester
 Protection (#8262)

---
 validator/db/kv/attester_protection.go      | 82 +++++++++++++--------
 validator/db/kv/attester_protection_test.go | 25 +++++++
 2 files changed, 76 insertions(+), 31 deletions(-)

diff --git a/validator/db/kv/attester_protection.go b/validator/db/kv/attester_protection.go
index 91bff5d90..2afb00625 100644
--- a/validator/db/kv/attester_protection.go
+++ b/validator/db/kv/attester_protection.go
@@ -78,38 +78,47 @@ func (store *Store) CheckSlashableAttestation(
 			return nil
 		}
 		// Check for surround votes.
-		return sourceEpochsBucket.ForEach(func(sourceEpochBytes []byte, targetEpochBytes []byte) error {
+		return sourceEpochsBucket.ForEach(func(sourceEpochBytes []byte, targetEpochsBytes []byte) error {
 			existingSourceEpoch := bytesutil.BytesToUint64BigEndian(sourceEpochBytes)
-			existingTargetEpoch := bytesutil.BytesToUint64BigEndian(targetEpochBytes)
-			existingAtt := &ethpb.IndexedAttestation{
-				Data: &ethpb.AttestationData{
-					Source: &ethpb.Checkpoint{Epoch: existingSourceEpoch},
-					Target: &ethpb.Checkpoint{Epoch: existingTargetEpoch},
-				},
+
+			// There can be multiple target epochs attested per source epoch.
+			attestedTargetEpochs := make([]uint64, 0)
+			for i := 0; i < len(targetEpochsBytes); i += 8 {
+				targetEpoch := bytesutil.BytesToUint64BigEndian(targetEpochsBytes[i : i+8])
+				attestedTargetEpochs = append(attestedTargetEpochs, targetEpoch)
 			}
-			// Checks if the incoming attestation is surrounding or
-			// is surrounded by an existing one.
-			surrounding := slashutil.IsSurround(att, existingAtt)
-			surrounded := slashutil.IsSurround(existingAtt, att)
-			if surrounding {
-				slashKind = SurroundingVote
-				return fmt.Errorf(
-					surroundingVoteMessage,
-					att.Data.Source.Epoch,
-					att.Data.Target.Epoch,
-					existingSourceEpoch,
-					existingTargetEpoch,
-				)
-			}
-			if surrounded {
-				slashKind = SurroundedVote
-				return fmt.Errorf(
-					surroundedVoteMessage,
-					att.Data.Source.Epoch,
-					att.Data.Target.Epoch,
-					existingSourceEpoch,
-					existingTargetEpoch,
-				)
+
+			for _, existingTargetEpoch := range attestedTargetEpochs {
+				existingAtt := &ethpb.IndexedAttestation{
+					Data: &ethpb.AttestationData{
+						Source: &ethpb.Checkpoint{Epoch: existingSourceEpoch},
+						Target: &ethpb.Checkpoint{Epoch: existingTargetEpoch},
+					},
+				}
+				// Checks if the incoming attestation is surrounding or
+				// is surrounded by an existing one.
+				surrounding := slashutil.IsSurround(att, existingAtt)
+				surrounded := slashutil.IsSurround(existingAtt, att)
+				if surrounding {
+					slashKind = SurroundingVote
+					return fmt.Errorf(
+						surroundingVoteMessage,
+						att.Data.Source.Epoch,
+						att.Data.Target.Epoch,
+						existingSourceEpoch,
+						existingTargetEpoch,
+					)
+				}
+				if surrounded {
+					slashKind = SurroundedVote
+					return fmt.Errorf(
+						surroundedVoteMessage,
+						att.Data.Source.Epoch,
+						att.Data.Target.Epoch,
+						existingSourceEpoch,
+						existingTargetEpoch,
+					)
+				}
 			}
 			return nil
 		})
@@ -219,7 +228,18 @@ func (store *Store) saveAttestationRecords(ctx context.Context, atts []*attestat
 			if err != nil {
 				return errors.Wrap(err, "could not create source epochs bucket")
 			}
-			if err := sourceEpochsBucket.Put(sourceEpochBytes, targetEpochBytes); err != nil {
+
+			// There can be multiple attested target epochs per source epoch.
+			// If a previous list exists, we append to that list with the incoming target epoch.
+			// Otherwise, we initialize it using the incoming target epoch.
+			var existingAttestedTargetsBytes []byte
+			if existing := sourceEpochsBucket.Get(sourceEpochBytes); existing != nil {
+				existingAttestedTargetsBytes = append(existing, targetEpochBytes...)
+			} else {
+				existingAttestedTargetsBytes = targetEpochBytes
+			}
+
+			if err := sourceEpochsBucket.Put(sourceEpochBytes, existingAttestedTargetsBytes); err != nil {
 				return errors.Wrapf(err, "could not save source epoch %d for epoch %d", att.source, att.target)
 			}
 			// Initialize buckets for the lowest target and source epochs.
diff --git a/validator/db/kv/attester_protection_test.go b/validator/db/kv/attester_protection_test.go
index 3801ebdd9..c8d82a72a 100644
--- a/validator/db/kv/attester_protection_test.go
+++ b/validator/db/kv/attester_protection_test.go
@@ -87,6 +87,31 @@ func TestStore_CheckSlashableAttestation_DoubleVote(t *testing.T) {
 	}
 }
 
+func TestStore_CheckSlashableAttestation_SurroundVote_MultipleTargetsPerSource(t *testing.T) {
+	ctx := context.Background()
+	numValidators := 1
+	pubKeys := make([][48]byte, numValidators)
+	validatorDB := setupDB(t, pubKeys)
+
+	// Create an attestation with source 1 and target 50, save it.
+	firstAtt := createAttestation(1, 50)
+	err := validatorDB.SaveAttestationForPubKey(ctx, pubKeys[0], [32]byte{0}, firstAtt)
+	require.NoError(t, err)
+
+	// Create an attestation with source 1 and target 100, save it.
+	secondAtt := createAttestation(1, 100)
+	err = validatorDB.SaveAttestationForPubKey(ctx, pubKeys[0], [32]byte{1}, secondAtt)
+	require.NoError(t, err)
+
+	// Create an attestation with source 0 and target 51, which should surround
+	// our first attestation. Given there can be multiple attested target epochs per
+	// source epoch, we expect our logic to be able to catch this slashable offense.
+	evilAtt := createAttestation(firstAtt.Data.Source.Epoch-1, firstAtt.Data.Target.Epoch+1)
+	slashable, err := validatorDB.CheckSlashableAttestation(ctx, pubKeys[0], [32]byte{2}, evilAtt)
+	require.NotNil(t, err)
+	assert.Equal(t, SurroundingVote, slashable)
+}
+
 func TestStore_CheckSlashableAttestation_SurroundVote_54kEpochs(t *testing.T) {
 	ctx := context.Background()
 	numValidators := 1