From 830a0a4bca16b6a9e2797aca7445b581a5d50c7d Mon Sep 17 00:00:00 2001
From: Preston Van Loon <preston@prysmaticlabs.com>
Date: Sat, 10 Aug 2019 15:49:58 -0400
Subject: [PATCH] Fix mid-epoch assignment requests (#3168)

* fix mid-epoch assignments

* add quick test comment
---
 beacon-chain/rpc/validator_server.go      | 11 +++--
 beacon-chain/rpc/validator_server_test.go | 58 +++++++++++++++++++++++
 2 files changed, 64 insertions(+), 5 deletions(-)

diff --git a/beacon-chain/rpc/validator_server.go b/beacon-chain/rpc/validator_server.go
index 69d50afb9..13b35a489 100644
--- a/beacon-chain/rpc/validator_server.go
+++ b/beacon-chain/rpc/validator_server.go
@@ -139,11 +139,12 @@ func (vs *ValidatorServer) CommitteeAssignment(ctx context.Context, req *pb.Assi
 		return nil, errors.Wrap(err, "could not fetch beacon state")
 	}
 
-	// Advance state with empty transitions up to the requested slot.
-	slotsToAdvance := req.EpochStart * params.BeaconConfig().SlotsPerEpoch
-	s, err = state.ProcessSlots(ctx, s, slotsToAdvance)
-	if err != nil {
-		return nil, fmt.Errorf("could not process slots up to %d", slotsToAdvance)
+	// Advance state with empty transitions up to the requested epoch start slot.
+	if epochStartSlot := helpers.StartSlot(req.EpochStart); s.Slot < epochStartSlot {
+		s, err = state.ProcessSlots(ctx, s, epochStartSlot)
+		if err != nil {
+			return nil, errors.Wrapf(err, "could not process slots up to %d", epochStartSlot)
+		}
 	}
 
 	validatorIndexMap := stateutils.ValidatorIndexMap(s)
diff --git a/beacon-chain/rpc/validator_server_test.go b/beacon-chain/rpc/validator_server_test.go
index 0fc786fee..6ae195473 100644
--- a/beacon-chain/rpc/validator_server_test.go
+++ b/beacon-chain/rpc/validator_server_test.go
@@ -237,6 +237,64 @@ func TestCommitteeAssignment_OK(t *testing.T) {
 	}
 }
 
+func TestCommitteeAssignment_CurrentEpoch_ShouldNotFail(t *testing.T) {
+	helpers.ClearAllCaches()
+
+	db := internal.SetupDB(t)
+	defer internal.TeardownDB(t, db)
+	ctx := context.Background()
+
+	genesis := blk.NewGenesisBlock([]byte{})
+	if err := db.SaveBlock(genesis); err != nil {
+		t.Fatalf("Could not save genesis block: %v", err)
+	}
+	depChainStart := params.BeaconConfig().MinGenesisActiveValidatorCount / 16
+
+	deposits, _ := testutil.SetupInitialDeposits(t, depChainStart)
+	state, err := state.GenesisBeaconState(deposits, 0, &ethpb.Eth1Data{})
+	if err != nil {
+		t.Fatalf("Could not setup genesis state: %v", err)
+	}
+	state.Slot = 5 // Set state to non-epoch start slot.
+	if err := db.UpdateChainHead(ctx, genesis, state); err != nil {
+		t.Fatalf("Could not save genesis state: %v", err)
+	}
+	var wg sync.WaitGroup
+	numOfValidators := int(depChainStart)
+	errs := make(chan error, numOfValidators)
+	for i := 0; i < len(deposits); i++ {
+		wg.Add(1)
+		go func(index int) {
+			errs <- db.SaveValidatorIndexBatch(deposits[index].Data.PublicKey, index)
+			wg.Done()
+		}(i)
+	}
+	wg.Wait()
+	close(errs)
+	for err := range errs {
+		if err != nil {
+			t.Fatalf("Could not save validator index: %v", err)
+		}
+	}
+
+	vs := &ValidatorServer{
+		beaconDB: db,
+	}
+
+	// Test the first validator in registry.
+	req := &pb.AssignmentRequest{
+		PublicKeys: [][]byte{deposits[0].Data.PublicKey},
+		EpochStart: 0,
+	}
+	res, err := vs.CommitteeAssignment(context.Background(), req)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(res.ValidatorAssignment) != 1 {
+		t.Error("Expected 1 assignment")
+	}
+}
+
 func TestCommitteeAssignment_multipleKeys_OK(t *testing.T) {
 	db := internal.SetupDB(t)
 	defer internal.TeardownDB(t, db)