From aa52e693ce0a6307bbb68af66b60cc9519810992 Mon Sep 17 00:00:00 2001
From: Preston Van Loon <preston@prysmaticlabs.com>
Date: Thu, 2 Jul 2020 13:24:30 -0700
Subject: [PATCH] Sync: do not remove attestations from pool on failed block
 processing (#6474)

* Add test and return error on failed block processing
* use atts in body
* Merge branch 'master' into p2p-potential-error-block
---
 beacon-chain/sync/subscriber_beacon_blocks.go |  1 +
 .../sync/subscriber_beacon_blocks_test.go     | 85 +++++++++++++++++++
 2 files changed, 86 insertions(+)

diff --git a/beacon-chain/sync/subscriber_beacon_blocks.go b/beacon-chain/sync/subscriber_beacon_blocks.go
index 21f5c53456..b6e8d4c05e 100644
--- a/beacon-chain/sync/subscriber_beacon_blocks.go
+++ b/beacon-chain/sync/subscriber_beacon_blocks.go
@@ -32,6 +32,7 @@ func (s *Service) beaconBlockSubscriber(ctx context.Context, msg proto.Message)
 
 	if err := s.chain.ReceiveBlockNoPubsub(ctx, signed, root); err != nil {
 		interop.WriteBlockToDisk(signed, true /*failed*/)
+		return err
 	}
 
 	// Delete attestations from the block in the pool to avoid inclusion in future block.
diff --git a/beacon-chain/sync/subscriber_beacon_blocks_test.go b/beacon-chain/sync/subscriber_beacon_blocks_test.go
index 25ee06ee6b..b6207e79d8 100644
--- a/beacon-chain/sync/subscriber_beacon_blocks_test.go
+++ b/beacon-chain/sync/subscriber_beacon_blocks_test.go
@@ -1,11 +1,16 @@
 package sync
 
 import (
+	"context"
 	"reflect"
 	"testing"
 
+	"github.com/gogo/protobuf/proto"
 	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
 	"github.com/prysmaticlabs/go-bitfield"
+	chainMock "github.com/prysmaticlabs/prysm/beacon-chain/blockchain/testing"
+	"github.com/prysmaticlabs/prysm/beacon-chain/core/helpers"
+	dbtest "github.com/prysmaticlabs/prysm/beacon-chain/db/testing"
 	"github.com/prysmaticlabs/prysm/beacon-chain/operations/attestations"
 )
 
@@ -41,3 +46,83 @@ func TestDeleteAttsInPool(t *testing.T) {
 		t.Error("Did not get wanted attestation from pool")
 	}
 }
+
+func TestService_beaconBlockSubscriber(t *testing.T) {
+	pooledAttestations := []*ethpb.Attestation{
+		// Aggregated.
+		{
+			AggregationBits: bitfield.Bitlist{0b00011111},
+			Data:            &ethpb.AttestationData{},
+		},
+		// Unaggregated.
+		{
+			AggregationBits: bitfield.Bitlist{0b00010001},
+			Data:            &ethpb.AttestationData{},
+		},
+	}
+
+	type args struct {
+		msg proto.Message
+	}
+	tests := []struct {
+		name    string
+		args    args
+		wantErr bool
+		check   func(*testing.T, *Service)
+	}{
+		{
+			name: "invalid block does not remove attestations",
+			args: args{
+				msg: &ethpb.SignedBeaconBlock{
+					Block: &ethpb.BeaconBlock{
+						// An empty block will return an err in mocked chainService.ReceiveBlockNoPubsub.
+						Body: &ethpb.BeaconBlockBody{Attestations: pooledAttestations},
+					},
+				},
+			},
+			wantErr: true,
+			check: func(t *testing.T, s *Service) {
+				if s.attPool.AggregatedAttestationCount() == 0 {
+					t.Error("Expected at least 1 aggregated attestation in the pool")
+				}
+				if s.attPool.UnaggregatedAttestationCount() == 0 {
+					t.Error("Expected at least 1 unaggregated attestation in the pool")
+				}
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			db, _ := dbtest.SetupDB(t)
+			s := &Service{
+				chain: &chainMock.ChainService{
+					DB: db,
+				},
+				attPool: attestations.NewPool(),
+			}
+			if err := s.initCaches(); err != nil {
+				t.Error(err)
+			}
+			// Set up attestation pool.
+			for _, att := range pooledAttestations {
+				if helpers.IsAggregated(att) {
+					if err := s.attPool.SaveAggregatedAttestation(att); err != nil {
+						t.Error(err)
+					}
+				} else {
+					if err := s.attPool.SaveUnaggregatedAttestation(att); err != nil {
+						t.Error(err)
+					}
+				}
+			}
+			// Perform method under test call.
+			if err := s.beaconBlockSubscriber(context.Background(), tt.args.msg); (err != nil) != tt.wantErr {
+				t.Errorf("beaconBlockSubscriber(ctx, msg) error = %v, wantErr %v", err, tt.wantErr)
+			}
+			// Perform any test check.
+			if tt.check != nil {
+				tt.check(t, s)
+			}
+		})
+	}
+}