datasync/downloader: bind staking info recovery to the requested peer

datasync/downloader/downloader.go

@@ -104,6 +104,7 @@ type Downloader struct {
 	mux  *event.TypeMux // Event multiplexer to announce sync operation events
 
 	isStakingInfoRecovery     bool
+	stakingInfoRecoveryPeer   string
 	stakingInfoRecoveryTotal  int
 	stakingInfoRecoveryCh     chan []*staking.P2PStakingInfo
 	stakingInfoRecoveryBlocks []uint64
@@ -636,6 +637,7 @@ func (d *Downloader) SyncStakingInfo(id string, from, to uint64) error {
 		return errors.New("the given peer is not registered")
 	}
 
+	d.stakingInfoRecoveryPeer = id
 	d.stakingInfoRecoveryBlocks = blockNums
 	d.stakingInfoRecoveryTotal = len(blockNums)
 	d.stakingInfoRecoveryCh = make(chan []*staking.P2PStakingInfo, 1)
@@ -669,15 +671,21 @@ func (d *Downloader) SyncStakingInfo(id string, from, to uint64) error {
 				return
 			case stakingInfos := <-d.stakingInfoRecoveryCh:
 				logger.Info("received stakinginfos", "len", len(stakingInfos))
-				for _, stakingInfo := range stakingInfos {
-					if d.stakingInfoRecoveryBlocks[0] != stakingInfo.BlockNum {
-						logger.Error("failed to receive expected block", "expected", d.stakingInfoRecoveryBlocks[0], "actual", stakingInfo.BlockNum)
+				if len(stakingInfos) > len(d.stakingInfoRecoveryBlocks) {
+					logger.Error("received more staking infos than requested", "received", len(stakingInfos), "pending", len(d.stakingInfoRecoveryBlocks))
+					return
+				}
+				for i, stakingInfo := range stakingInfos {
+					if d.stakingInfoRecoveryBlocks[i] != stakingInfo.BlockNum {
+						logger.Error("failed to receive expected block", "expected", d.stakingInfoRecoveryBlocks[i], "actual", stakingInfo.BlockNum)
 						return
 					}
-					d.stakingModule.PutStakingInfoToDB(stakingInfo.BlockNum, staking.ToStakingInfo(stakingInfo))
+				}
+				for i, stakingInfo := range stakingInfos {
+					d.stakingModule.PutStakingInfoToDB(d.stakingInfoRecoveryBlocks[i], staking.ToStakingInfo(stakingInfo))
 					fixed++
-					d.stakingInfoRecoveryBlocks = d.stakingInfoRecoveryBlocks[1:]
 				}
+				d.stakingInfoRecoveryBlocks = d.stakingInfoRecoveryBlocks[len(stakingInfos):]
 
 				if len(d.stakingInfoRecoveryBlocks) == 0 {
 					logger.Info("syncing staking info is finished", "fixed", fixed)
@@ -1934,8 +1942,11 @@ func (d *Downloader) DeliverReceipts(id string, receipts [][]*types.Receipt) (er
 
 // DeliverStakingInfos injects a new batch of staking information received from a remote node.
 func (d *Downloader) DeliverStakingInfos(id string, stakingInfos []*staking.P2PStakingInfo) error {
-	if d.isStakingInfoRecovery {
-		d.stakingInfoRecoveryCh <- stakingInfos
+	if d.isStakingInfoRecovery && id == d.stakingInfoRecoveryPeer {
+		select {
+		case d.stakingInfoRecoveryCh <- stakingInfos:
+		default:
+		}
 	}
 	return d.deliver(id, d.stakingInfoCh, &stakingInfoPack{id, stakingInfos}, stakingInfoInMeter, stakingInfoDropMeter)
 }

datasync/downloader/downloader_test.go

@@ -1992,6 +1992,57 @@ func testStakingInfoSync(t *testing.T, protocol int) {
 	}
 }
 
+// newRecoveryTester starts a staking info recovery from "requested" and returns the
+// tester. The peer answers slowly, so the recovery is still pending on return.
+func newRecoveryTester(t *testing.T) *downloadTester {
+	config := params.TestKaiaConfig("cancun")
+	config.Governance.Reward.StakingUpdateInterval = testInterval
+
+	tester := newTesterWithConfig(t, config)
+	hashes, headers, blocks, receipts, stakingInfos := tester.makeChain(16, 0, tester.genesis, nil, false)
+	if err := tester.newSlowPeer("requested", 65, hashes, headers, blocks, receipts, stakingInfos, 500*time.Millisecond); err != nil {
+		t.Fatal(err)
+	}
+	for hash, info := range stakingInfos {
+		tester.stateDb.WriteCanonicalHash(hash, info.BlockNum)
+	}
+	// Pending blocks become [4, 8].
+	if err := tester.downloader.SyncStakingInfo("requested", 9, 9); err != nil {
+		t.Fatal(err)
+	}
+	return tester
+}
+
+func TestSyncStakingInfoIgnoresUnselectedPeer(t *testing.T) {
+	tester := newRecoveryTester(t)
+	defer tester.terminate()
+
+	forged := &staking.P2PStakingInfo{BlockNum: 4, KEFAddr: common.HexToAddress("0x2222222222222222222222222222222222222222")}
+	tester.downloader.DeliverStakingInfos("mallory", []*staking.P2PStakingInfo{forged})
+	time.Sleep(200 * time.Millisecond)
+
+	if si := staking_impl.ReadStakingInfo(tester.stateDb.GetMiscDB(), forged.BlockNum); si != nil {
+		t.Fatalf("staking info from an unselected peer was persisted: %+v", si)
+	}
+}
+
+func TestSyncStakingInfoRejectsOverlongBatch(t *testing.T) {
+	tester := newRecoveryTester(t)
+	defer tester.terminate()
+
+	// Three entries against two pending blocks: the batch must be dropped whole.
+	tester.downloader.DeliverStakingInfos("requested", []*staking.P2PStakingInfo{
+		{BlockNum: 4}, {BlockNum: 8}, {BlockNum: 12},
+	})
+	time.Sleep(200 * time.Millisecond)
+
+	for _, num := range []uint64{4, 8, 12} {
+		if si := staking_impl.ReadStakingInfo(tester.stateDb.GetMiscDB(), num); si != nil {
+			t.Fatalf("an entry of the overlong batch was persisted at %d: %+v", num, si)
+		}
+	}
+}
+
 func TestCalcStakingBlockNumber(t *testing.T) {
 	testCase := []struct {
 		interval uint64