Part 2 of block chain service refactor - move process attestation (#4672)

beacon-chain/blockchain/BUILD.bazel

@@ -7,6 +7,8 @@ go_library(
         "info.go",
         "log.go",
         "metrics.go",
+        "process_attestation.go",
+        "process_attestation_helpers.go",
         "process_block.go",
         "process_block_helpers.go",
         "receive_attestation.go",
@@ -17,6 +19,7 @@ go_library(
     visibility = ["//beacon-chain:__subpackages__"],
     deps = [
         "//beacon-chain/blockchain/forkchoice:go_default_library",
+        "//beacon-chain/cache:go_default_library",
         "//beacon-chain/cache/depositcache:go_default_library",
         "//beacon-chain/core/blocks:go_default_library",
         "//beacon-chain/core/epoch/precompute:go_default_library",
@@ -64,6 +67,7 @@ go_test(
     size = "medium",
     srcs = [
         "chain_info_test.go",
+        "process_attestation_test.go",
         "process_block_test.go",
         "receive_attestation_test.go",
         "receive_block_test.go",

beacon-chain/blockchain/process_attestation.go

@@ -0,0 +1,128 @@
+package blockchain
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/gogo/protobuf/proto"
+	"github.com/pkg/errors"
+	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
+	"github.com/prysmaticlabs/prysm/beacon-chain/core/helpers"
+	"github.com/prysmaticlabs/prysm/shared/bytesutil"
+	"go.opencensus.io/trace"
+)
+
+// ErrTargetRootNotInDB returns when the target block root of an attestation cannot be found in the
+// beacon database.
+var ErrTargetRootNotInDB = errors.New("target root does not exist in db")
+
+// onAttestation is called whenever an attestation is received, verifies the attestation is valid and saves
+/// it to the DB.
+//
+// Spec pseudocode definition:
+//   def on_attestation(store: Service, attestation: Attestation) -> None:
+//    """
+//    Run ``on_attestation`` upon receiving a new ``attestation`` from either within a block or directly on the wire.
+//
+//    An ``attestation`` that is asserted as invalid may be valid at a later time,
+//    consider scheduling it for later processing in such case.
+//    """
+//    target = attestation.data.target
+//
+//    # Attestations must be from the current or previous epoch
+//    current_epoch = compute_epoch_at_slot(get_current_slot(store))
+//    # Use GENESIS_EPOCH for previous when genesis to avoid underflow
+//    previous_epoch = current_epoch - 1 if current_epoch > GENESIS_EPOCH else GENESIS_EPOCH
+//    assert target.epoch in [current_epoch, previous_epoch]
+//    assert target.epoch == compute_epoch_at_slot(attestation.data.slot)
+//
+//    # Attestations target be for a known block. If target block is unknown, delay consideration until the block is found
+//    assert target.root in store.blocks
+//    # Attestations cannot be from future epochs. If they are, delay consideration until the epoch arrives
+//    base_state = store.block_states[target.root].copy()
+//    assert store.time >= base_state.genesis_time + compute_start_slot_at_epoch(target.epoch) * SECONDS_PER_SLOT
+//
+//    # Attestations must be for a known block. If block is unknown, delay consideration until the block is found
+//    assert attestation.data.beacon_block_root in store.blocks
+//    # Attestations must not be for blocks in the future. If not, the attestation should not be considered
+//    assert store.blocks[attestation.data.beacon_block_root].slot <= attestation.data.slot
+//
+//    # Service target checkpoint state if not yet seen
+//    if target not in store.checkpoint_states:
+//        process_slots(base_state, compute_start_slot_at_epoch(target.epoch))
+//        store.checkpoint_states[target] = base_state
+//    target_state = store.checkpoint_states[target]
+//
+//    # Attestations can only affect the fork choice of subsequent slots.
+//    # Delay consideration in the fork choice until their slot is in the past.
+//    assert store.time >= (attestation.data.slot + 1) * SECONDS_PER_SLOT
+//
+//    # Get state at the `target` to validate attestation and calculate the committees
+//    indexed_attestation = get_indexed_attestation(target_state, attestation)
+//    assert is_valid_indexed_attestation(target_state, indexed_attestation)
+//
+//    # Update latest messages
+//    for i in indexed_attestation.attesting_indices:
+//        if i not in store.latest_messages or target.epoch > store.latest_messages[i].epoch:
+//            store.latest_messages[i] = LatestMessage(epoch=target.epoch, root=attestation.data.beacon_block_root)
+func (s *Service) onAttestation(ctx context.Context, a *ethpb.Attestation) ([]uint64, error) {
+	ctx, span := trace.StartSpan(ctx, "blockchain.onAttestation")
+	defer span.End()
+
+	tgt := proto.Clone(a.Data.Target).(*ethpb.Checkpoint)
+	tgtSlot := helpers.StartSlot(tgt.Epoch)
+
+	if helpers.SlotToEpoch(a.Data.Slot) != a.Data.Target.Epoch {
+		return nil, fmt.Errorf("data slot is not in the same epoch as target %d != %d", helpers.SlotToEpoch(a.Data.Slot), a.Data.Target.Epoch)
+	}
+
+	// Verify beacon node has seen the target block before.
+	if !s.beaconDB.HasBlock(ctx, bytesutil.ToBytes32(tgt.Root)) {
+		return nil, ErrTargetRootNotInDB
+	}
+
+	// Verify attestation target has had a valid pre state produced by the target block.
+	baseState, err := s.verifyAttPreState(ctx, tgt)
+	if err != nil {
+		return nil, err
+	}
+
+	// Verify attestation target is from current epoch or previous epoch.
+	if err := s.verifyAttTargetEpoch(ctx, baseState.GenesisTime, uint64(time.Now().Unix()), tgt); err != nil {
+		return nil, err
+	}
+
+	// Verify Attestations cannot be from future epochs.
+	if err := helpers.VerifySlotTime(baseState.GenesisTime, tgtSlot); err != nil {
+		return nil, errors.Wrap(err, "could not verify attestation target slot")
+	}
+
+	// Verify attestation beacon block is known and not from the future.
+	if err := s.verifyBeaconBlock(ctx, a.Data); err != nil {
+		return nil, errors.Wrap(err, "could not verify attestation beacon block")
+	}
+
+	// Service target checkpoint state if not yet seen.
+	baseState, err = s.saveCheckpointState(ctx, baseState, tgt)
+	if err != nil {
+		return nil, err
+	}
+
+	// Verify attestations can only affect the fork choice of subsequent slots.
+	if err := helpers.VerifySlotTime(baseState.GenesisTime, a.Data.Slot+1); err != nil {
+		return nil, err
+	}
+
+	// Use the target state to to validate attestation and calculate the committees.
+	indexedAtt, err := s.verifyAttestation(ctx, baseState, a)
+	if err != nil {
+		return nil, err
+	}
+
+	if err := s.beaconDB.SaveAttestation(ctx, a); err != nil {
+		return nil, err
+	}
+
+	return indexedAtt.AttestingIndices, nil
+}

beacon-chain/blockchain/process_attestation_helpers.go

@@ -0,0 +1,135 @@
+package blockchain
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/gogo/protobuf/proto"
+	"github.com/pkg/errors"
+	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
+	"github.com/prysmaticlabs/prysm/beacon-chain/cache"
+	"github.com/prysmaticlabs/prysm/beacon-chain/core/blocks"
+	"github.com/prysmaticlabs/prysm/beacon-chain/core/helpers"
+	"github.com/prysmaticlabs/prysm/beacon-chain/core/state"
+	pb "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
+	"github.com/prysmaticlabs/prysm/shared/bytesutil"
+	"github.com/prysmaticlabs/prysm/shared/params"
+	"go.opencensus.io/trace"
+)
+
+// verifyAttPreState validates input attested check point has a valid pre-state.
+func (s *Service) verifyAttPreState(ctx context.Context, c *ethpb.Checkpoint) (*pb.BeaconState, error) {
+	baseState, err := s.beaconDB.State(ctx, bytesutil.ToBytes32(c.Root))
+	if err != nil {
+		return nil, errors.Wrapf(err, "could not get pre state for slot %d", helpers.StartSlot(c.Epoch))
+	}
+	if baseState == nil {
+		return nil, fmt.Errorf("pre state of target block %d does not exist", helpers.StartSlot(c.Epoch))
+	}
+	return baseState, nil
+}
+
+// verifyAttTargetEpoch validates attestation is from the current or previous epoch.
+func (s *Service) verifyAttTargetEpoch(ctx context.Context, genesisTime uint64, nowTime uint64, c *ethpb.Checkpoint) error {
+	currentSlot := (nowTime - genesisTime) / params.BeaconConfig().SecondsPerSlot
+	currentEpoch := helpers.SlotToEpoch(currentSlot)
+	var prevEpoch uint64
+	// Prevents previous epoch under flow
+	if currentEpoch > 1 {
+		prevEpoch = currentEpoch - 1
+	}
+	if c.Epoch != prevEpoch && c.Epoch != currentEpoch {
+		return fmt.Errorf("target epoch %d does not match current epoch %d or prev epoch %d", c.Epoch, currentEpoch, prevEpoch)
+	}
+	return nil
+}
+
+// verifyBeaconBlock verifies beacon head block is known and not from the future.
+func (s *Service) verifyBeaconBlock(ctx context.Context, data *ethpb.AttestationData) error {
+	b, err := s.beaconDB.Block(ctx, bytesutil.ToBytes32(data.BeaconBlockRoot))
+	if err != nil {
+		return err
+	}
+	if b == nil || b.Block == nil {
+		return fmt.Errorf("beacon block %#x does not exist", bytesutil.Trunc(data.BeaconBlockRoot))
+	}
+	if b.Block.Slot > data.Slot {
+		return fmt.Errorf("could not process attestation for future block, %d > %d", b.Block.Slot, data.Slot)
+	}
+	return nil
+}
+
+// saveCheckpointState saves and returns the processed state with the associated check point.
+func (s *Service) saveCheckpointState(ctx context.Context, baseState *pb.BeaconState, c *ethpb.Checkpoint) (*pb.BeaconState, error) {
+	ctx, span := trace.StartSpan(ctx, "blockchain.saveCheckpointState")
+	defer span.End()
+
+	s.checkpointStateLock.Lock()
+	defer s.checkpointStateLock.Unlock()
+	cachedState, err := s.checkpointState.StateByCheckpoint(c)
+	if err != nil {
+		return nil, errors.Wrap(err, "could not get cached checkpoint state")
+	}
+	if cachedState != nil {
+		return cachedState, nil
+	}
+
+	// Advance slots only when it's higher than current state slot.
+	if helpers.StartSlot(c.Epoch) > baseState.Slot {
+		stateCopy := proto.Clone(baseState).(*pb.BeaconState)
+		stateCopy, err = state.ProcessSlots(ctx, stateCopy, helpers.StartSlot(c.Epoch))
+		if err != nil {
+			return nil, errors.Wrapf(err, "could not process slots up to %d", helpers.StartSlot(c.Epoch))
+		}
+
+		if err := s.checkpointState.AddCheckpointState(&cache.CheckpointState{
+			Checkpoint: c,
+			State:      stateCopy,
+		}); err != nil {
+			return nil, errors.Wrap(err, "could not saved checkpoint state to cache")
+		}
+
+		return stateCopy, nil
+	}
+
+	return baseState, nil
+}
+
+// verifyAttestation validates input attestation is valid.
+func (s *Service) verifyAttestation(ctx context.Context, baseState *pb.BeaconState, a *ethpb.Attestation) (*ethpb.IndexedAttestation, error) {
+	committee, err := helpers.BeaconCommitteeFromState(baseState, a.Data.Slot, a.Data.CommitteeIndex)
+	if err != nil {
+		return nil, err
+	}
+	indexedAtt, err := blocks.ConvertToIndexed(ctx, a, committee)
+	if err != nil {
+		return nil, errors.Wrap(err, "could not convert attestation to indexed attestation")
+	}
+
+	if err := blocks.VerifyIndexedAttestation(ctx, baseState, indexedAtt); err != nil {
+
+		// TODO(3603): Delete the following signature verify fallback when issue 3603 closes.
+		// When signature fails to verify with committee cache enabled at run time,
+		// the following re-runs the same signature verify routine without cache in play.
+		// This provides extra assurance that committee cache can't break run time.
+		if err == blocks.ErrSigFailedToVerify {
+			committee, err = helpers.BeaconCommitteeWithoutCache(baseState, a.Data.Slot, a.Data.CommitteeIndex)
+			if err != nil {
+				return nil, errors.Wrap(err, "could not convert attestation to indexed attestation without cache")
+			}
+			indexedAtt, err = blocks.ConvertToIndexed(ctx, a, committee)
+			if err != nil {
+				return nil, errors.Wrap(err, "could not convert attestation to indexed attestation")
+			}
+			if err := blocks.VerifyIndexedAttestation(ctx, baseState, indexedAtt); err != nil {
+				return nil, errors.Wrap(err, "could not verify indexed attestation without cache")
+			}
+			sigFailsToVerify.Inc()
+			return indexedAtt, nil
+		}
+
+		return nil, errors.Wrap(err, "could not verify indexed attestation")
+	}
+
+	return indexedAtt, nil
+}

beacon-chain/blockchain/process_attestation_test.go

@@ -0,0 +1,383 @@
+package blockchain
+
+import (
+	"context"
+	"reflect"
+	"strings"
+	"testing"
+
+	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
+	"github.com/prysmaticlabs/go-ssz"
+	"github.com/prysmaticlabs/prysm/beacon-chain/core/helpers"
+	"github.com/prysmaticlabs/prysm/beacon-chain/core/state"
+	testDB "github.com/prysmaticlabs/prysm/beacon-chain/db/testing"
+	pb "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
+	"github.com/prysmaticlabs/prysm/shared/params"
+	"github.com/prysmaticlabs/prysm/shared/testutil"
+)
+
+func TestStore_OnAttestation(t *testing.T) {
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+
+	cfg := &Config{BeaconDB: db}
+	service, err := NewService(ctx, cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	_, err = blockTree1(db, []byte{'g'})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	BlkWithOutState := &ethpb.SignedBeaconBlock{Block: &ethpb.BeaconBlock{Slot: 0}}
+	if err := db.SaveBlock(ctx, BlkWithOutState); err != nil {
+		t.Fatal(err)
+	}
+	BlkWithOutStateRoot, _ := ssz.HashTreeRoot(BlkWithOutState.Block)
+
+	BlkWithStateBadAtt := &ethpb.SignedBeaconBlock{Block: &ethpb.BeaconBlock{Slot: 1}}
+	if err := db.SaveBlock(ctx, BlkWithStateBadAtt); err != nil {
+		t.Fatal(err)
+	}
+	BlkWithStateBadAttRoot, _ := ssz.HashTreeRoot(BlkWithStateBadAtt.Block)
+	if err := service.beaconDB.SaveState(ctx, &pb.BeaconState{}, BlkWithStateBadAttRoot); err != nil {
+		t.Fatal(err)
+	}
+
+	BlkWithValidState := &ethpb.SignedBeaconBlock{Block: &ethpb.BeaconBlock{Slot: 2}}
+	if err := db.SaveBlock(ctx, BlkWithValidState); err != nil {
+		t.Fatal(err)
+	}
+	BlkWithValidStateRoot, _ := ssz.HashTreeRoot(BlkWithValidState.Block)
+	if err := service.beaconDB.SaveState(ctx, &pb.BeaconState{
+		Fork: &pb.Fork{
+			Epoch:           0,
+			CurrentVersion:  params.BeaconConfig().GenesisForkVersion,
+			PreviousVersion: params.BeaconConfig().GenesisForkVersion,
+		},
+		RandaoMixes: make([][]byte, params.BeaconConfig().EpochsPerHistoricalVector),
+	}, BlkWithValidStateRoot); err != nil {
+		t.Fatal(err)
+	}
+
+	tests := []struct {
+		name          string
+		a             *ethpb.Attestation
+		s             *pb.BeaconState
+		wantErr       bool
+		wantErrString string
+	}{
+		{
+			name:          "attestation's data slot not aligned with target vote",
+			a:             &ethpb.Attestation{Data: &ethpb.AttestationData{Slot: params.BeaconConfig().SlotsPerEpoch, Target: &ethpb.Checkpoint{}}},
+			s:             &pb.BeaconState{},
+			wantErr:       true,
+			wantErrString: "data slot is not in the same epoch as target 1 != 0",
+		},
+		{
+			name:          "attestation's target root not in db",
+			a:             &ethpb.Attestation{Data: &ethpb.AttestationData{Target: &ethpb.Checkpoint{Root: []byte{'A'}}}},
+			s:             &pb.BeaconState{},
+			wantErr:       true,
+			wantErrString: "target root does not exist in db",
+		},
+		{
+			name:          "no pre state for attestations's target block",
+			a:             &ethpb.Attestation{Data: &ethpb.AttestationData{Target: &ethpb.Checkpoint{Root: BlkWithOutStateRoot[:]}}},
+			s:             &pb.BeaconState{},
+			wantErr:       true,
+			wantErrString: "pre state of target block 0 does not exist",
+		},
+		{
+			name: "process attestation doesn't match current epoch",
+			a: &ethpb.Attestation{Data: &ethpb.AttestationData{Slot: 100 * params.BeaconConfig().SlotsPerEpoch, Target: &ethpb.Checkpoint{Epoch: 100,
+				Root: BlkWithStateBadAttRoot[:]}}},
+			s:             &pb.BeaconState{},
+			wantErr:       true,
+			wantErrString: "does not match current epoch",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			_, err := service.onAttestation(ctx, tt.a)
+			if tt.wantErr {
+				if !strings.Contains(err.Error(), tt.wantErrString) {
+					t.Errorf("Store.onAttestation() error = %v, wantErr = %v", err, tt.wantErrString)
+				}
+			} else {
+				t.Error(err)
+			}
+		})
+	}
+}
+
+func TestStore_SaveCheckpointState(t *testing.T) {
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+	params.UseDemoBeaconConfig()
+
+	cfg := &Config{BeaconDB: db}
+	service, err := NewService(ctx, cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	s := &pb.BeaconState{
+		Fork: &pb.Fork{
+			Epoch:           0,
+			CurrentVersion:  params.BeaconConfig().GenesisForkVersion,
+			PreviousVersion: params.BeaconConfig().GenesisForkVersion,
+		},
+		RandaoMixes:         make([][]byte, params.BeaconConfig().EpochsPerHistoricalVector),
+		StateRoots:          make([][]byte, params.BeaconConfig().EpochsPerHistoricalVector),
+		BlockRoots:          make([][]byte, params.BeaconConfig().SlotsPerHistoricalRoot),
+		LatestBlockHeader:   &ethpb.BeaconBlockHeader{},
+		JustificationBits:   []byte{0},
+		Slashings:           make([]uint64, params.BeaconConfig().EpochsPerSlashingsVector),
+		FinalizedCheckpoint: &ethpb.Checkpoint{},
+	}
+	r := [32]byte{'g'}
+	if err := service.beaconDB.SaveState(ctx, s, r); err != nil {
+		t.Fatal(err)
+	}
+	service.justifiedCheckpt = &ethpb.Checkpoint{Root: r[:]}
+	service.bestJustifiedCheckpt = &ethpb.Checkpoint{Root: r[:]}
+	service.finalizedCheckpt = &ethpb.Checkpoint{Root: r[:]}
+	service.prevFinalizedCheckpt = &ethpb.Checkpoint{Root: r[:]}
+
+	cp1 := &ethpb.Checkpoint{Epoch: 1, Root: []byte{'A'}}
+	s1, err := service.saveCheckpointState(ctx, s, cp1)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if s1.Slot != 1*params.BeaconConfig().SlotsPerEpoch {
+		t.Errorf("Wanted state slot: %d, got: %d", 1*params.BeaconConfig().SlotsPerEpoch, s1.Slot)
+	}
+
+	cp2 := &ethpb.Checkpoint{Epoch: 2, Root: []byte{'B'}}
+	s2, err := service.saveCheckpointState(ctx, s, cp2)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if s2.Slot != 2*params.BeaconConfig().SlotsPerEpoch {
+		t.Errorf("Wanted state slot: %d, got: %d", 2*params.BeaconConfig().SlotsPerEpoch, s2.Slot)
+	}
+
+	s1, err = service.saveCheckpointState(ctx, nil, cp1)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if s1.Slot != 1*params.BeaconConfig().SlotsPerEpoch {
+		t.Errorf("Wanted state slot: %d, got: %d", 1*params.BeaconConfig().SlotsPerEpoch, s1.Slot)
+	}
+
+	s1, err = service.checkpointState.StateByCheckpoint(cp1)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if s1.Slot != 1*params.BeaconConfig().SlotsPerEpoch {
+		t.Errorf("Wanted state slot: %d, got: %d", 1*params.BeaconConfig().SlotsPerEpoch, s1.Slot)
+	}
+
+	s2, err = service.checkpointState.StateByCheckpoint(cp2)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if s2.Slot != 2*params.BeaconConfig().SlotsPerEpoch {
+		t.Errorf("Wanted state slot: %d, got: %d", 2*params.BeaconConfig().SlotsPerEpoch, s2.Slot)
+	}
+
+	s.Slot = params.BeaconConfig().SlotsPerEpoch + 1
+	service.justifiedCheckpt = &ethpb.Checkpoint{Root: r[:]}
+	service.bestJustifiedCheckpt = &ethpb.Checkpoint{Root: r[:]}
+	service.finalizedCheckpt = &ethpb.Checkpoint{Root: r[:]}
+	service.prevFinalizedCheckpt = &ethpb.Checkpoint{Root: r[:]}
+	cp3 := &ethpb.Checkpoint{Epoch: 1, Root: []byte{'C'}}
+
+	s3, err := service.saveCheckpointState(ctx, s, cp3)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if s3.Slot != s.Slot {
+		t.Errorf("Wanted state slot: %d, got: %d", s.Slot, s3.Slot)
+	}
+}
+
+func TestStore_UpdateCheckpointState(t *testing.T) {
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+
+	cfg := &Config{BeaconDB: db}
+	service, err := NewService(ctx, cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	epoch := uint64(1)
+	baseState, _ := testutil.DeterministicGenesisState(t, 1)
+	baseState.Slot = epoch * params.BeaconConfig().SlotsPerEpoch
+	checkpoint := &ethpb.Checkpoint{Epoch: epoch}
+	returned, err := service.saveCheckpointState(ctx, baseState, checkpoint)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !reflect.DeepEqual(baseState, returned) {
+		t.Error("Incorrectly returned base state")
+	}
+
+	cached, err := service.checkpointState.StateByCheckpoint(checkpoint)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if cached != nil {
+		t.Error("State shouldn't have been cached")
+	}
+
+	epoch = uint64(2)
+	newCheckpoint := &ethpb.Checkpoint{Epoch: epoch}
+	returned, err = service.saveCheckpointState(ctx, baseState, newCheckpoint)
+	if err != nil {
+		t.Fatal(err)
+	}
+	baseState, err = state.ProcessSlots(ctx, baseState, helpers.StartSlot(newCheckpoint.Epoch))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !reflect.DeepEqual(baseState, returned) {
+		t.Error("Incorrectly returned base state")
+	}
+
+	cached, err = service.checkpointState.StateByCheckpoint(newCheckpoint)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !reflect.DeepEqual(returned, cached) {
+		t.Error("Incorrectly cached base state")
+	}
+}
+
+func TestAttEpoch_MatchPrevEpoch(t *testing.T) {
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+
+	cfg := &Config{BeaconDB: db}
+	service, err := NewService(ctx, cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if err := service.verifyAttTargetEpoch(
+		ctx,
+		0,
+		params.BeaconConfig().SlotsPerEpoch*params.BeaconConfig().SecondsPerSlot,
+		&ethpb.Checkpoint{}); err != nil {
+		t.Error(err)
+	}
+}
+
+func TestAttEpoch_MatchCurrentEpoch(t *testing.T) {
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+
+	cfg := &Config{BeaconDB: db}
+	service, err := NewService(ctx, cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if err := service.verifyAttTargetEpoch(
+		ctx,
+		0,
+		params.BeaconConfig().SlotsPerEpoch*params.BeaconConfig().SecondsPerSlot,
+		&ethpb.Checkpoint{Epoch: 1}); err != nil {
+		t.Error(err)
+	}
+}
+
+func TestAttEpoch_NotMatch(t *testing.T) {
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+
+	cfg := &Config{BeaconDB: db}
+	service, err := NewService(ctx, cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if err := service.verifyAttTargetEpoch(
+		ctx,
+		0,
+		2*params.BeaconConfig().SlotsPerEpoch*params.BeaconConfig().SecondsPerSlot,
+		&ethpb.Checkpoint{}); !strings.Contains(err.Error(),
+		"target epoch 0 does not match current epoch 2 or prev epoch 1") {
+		t.Error("Did not receive wanted error")
+	}
+}
+
+func TestVerifyBeaconBlock_NoBlock(t *testing.T) {
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+
+	cfg := &Config{BeaconDB: db}
+	service, err := NewService(ctx, cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	d := &ethpb.AttestationData{}
+	if err := service.verifyBeaconBlock(ctx, d); !strings.Contains(err.Error(), "beacon block  does not exist") {
+		t.Error("Did not receive the wanted error")
+	}
+}
+
+func TestVerifyBeaconBlock_futureBlock(t *testing.T) {
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+
+	cfg := &Config{BeaconDB: db}
+	service, err := NewService(ctx, cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	b := &ethpb.SignedBeaconBlock{Block: &ethpb.BeaconBlock{Slot: 2}}
+	service.beaconDB.SaveBlock(ctx, b)
+	r, _ := ssz.HashTreeRoot(b.Block)
+	d := &ethpb.AttestationData{Slot: 1, BeaconBlockRoot: r[:]}
+
+	if err := service.verifyBeaconBlock(ctx, d); !strings.Contains(err.Error(), "could not process attestation for future block") {
+		t.Error("Did not receive the wanted error")
+	}
+}
+
+func TestVerifyBeaconBlock_OK(t *testing.T) {
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+
+	cfg := &Config{BeaconDB: db}
+	service, err := NewService(ctx, cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	b := &ethpb.SignedBeaconBlock{Block: &ethpb.BeaconBlock{Slot: 2}}
+	service.beaconDB.SaveBlock(ctx, b)
+	r, _ := ssz.HashTreeRoot(b.Block)
+	d := &ethpb.AttestationData{Slot: 2, BeaconBlockRoot: r[:]}
+
+	if err := service.verifyBeaconBlock(ctx, d); err != nil {
+		t.Error("Did not receive the wanted error")
+	}
+}

beacon-chain/blockchain/process_block.go

@@ -48,7 +48,7 @@ import (
 //    if state.finalized_checkpoint.epoch > store.finalized_checkpoint.epoch:
 //        store.finalized_checkpoint = state.finalized_checkpoint
 func (s *Service) onBlock(ctx context.Context, signed *ethpb.SignedBeaconBlock) (*pb.BeaconState, error) {
-	ctx, span := trace.StartSpan(ctx, "forkchoice.onBlock")
+	ctx, span := trace.StartSpan(ctx, "blockchain.onBlock")
 	defer span.End()
 
 	if signed == nil || signed.Block == nil {
@@ -139,7 +139,7 @@ func (s *Service) onBlock(ctx context.Context, signed *ethpb.SignedBeaconBlock)
 // It runs state transition on the block and without any BLS verification. The excluded BLS verification
 // includes attestation's aggregated signature. It also does not save attestations.
 func (s *Service) onBlockInitialSyncStateTransition(ctx context.Context, signed *ethpb.SignedBeaconBlock) (*pb.BeaconState, error) {
-	ctx, span := trace.StartSpan(ctx, "forkchoice.onBlock")
+	ctx, span := trace.StartSpan(ctx, "blockchain.onBlock")
 	defer span.End()
 
 	if signed == nil || signed.Block == nil {

beacon-chain/blockchain/receive_attestation.go

@@ -33,13 +33,21 @@ func (s *Service) ReceiveAttestationNoPubsub(ctx context.Context, att *ethpb.Att
 	defer span.End()
 
 	// Update forkchoice store for the new attestation
-	indices, err := s.forkChoiceStoreOld.OnAttestation(ctx, att)
-	if err != nil {
-		return errors.Wrap(err, "could not process attestation from fork choice service")
-	}
-
+	indices := make([]uint64, 0)
+	var err error
 	if featureconfig.Get().ProtoArrayForkChoice {
+		indices, err = s.onAttestation(ctx, att)
+		if err != nil {
+			return errors.Wrap(err, "could not process attestation from fork choice service")
+		}
+
 		s.forkChoiceStore.ProcessAttestation(ctx, indices, bytesutil.ToBytes32(att.Data.BeaconBlockRoot), att.Data.Target.Epoch)
+
+	} else {
+		indices, err = s.forkChoiceStoreOld.OnAttestation(ctx, att)
+		if err != nil {
+			return errors.Wrap(err, "could not process attestation from fork choice service")
+		}
 	}
 
 	// Run fork choice for head block after updating fork choice store.

beacon-chain/blockchain/service.go

@@ -15,6 +15,7 @@ import (
 	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
 	"github.com/prysmaticlabs/go-ssz"
 	"github.com/prysmaticlabs/prysm/beacon-chain/blockchain/forkchoice"
+	"github.com/prysmaticlabs/prysm/beacon-chain/cache"
 	"github.com/prysmaticlabs/prysm/beacon-chain/cache/depositcache"
 	"github.com/prysmaticlabs/prysm/beacon-chain/core/blocks"
 	"github.com/prysmaticlabs/prysm/beacon-chain/core/epoch/precompute"
@@ -70,6 +71,8 @@ type Service struct {
 	voteLock               sync.RWMutex
 	initSyncState          map[[32]byte]*pb.BeaconState
 	initSyncStateLock      sync.RWMutex
+	checkpointState        *cache.CheckpointStateCache
+	checkpointStateLock    sync.Mutex
 }
 
 // Config options for the service.
@@ -107,6 +110,7 @@ func NewService(ctx context.Context, cfg *Config) (*Service, error) {
 		epochParticipation: make(map[uint64]*precompute.Balance),
 		forkChoiceStore:    cfg.ForkChoiceStore,
 		initSyncState:      make(map[[32]byte]*pb.BeaconState),
+		checkpointState:    cache.NewCheckpointStateCache(),
 	}, nil
 }