Fix not in selected chain crash (#2082)

* Avoid creating the chain iterator if high hash is actually low hash

* Always use iterator in nextPruningPointAndCandidateByBlockHash

Co-authored-by: Ori Newman <orinewman1@gmail.com>

domain/consensus/model/testapi/test_consensus.go

@@ -47,6 +47,7 @@ type TestConsensus interface {
 
 	AddUTXOInvalidBlock(parentHashes []*externalapi.DomainHash) (*externalapi.DomainHash,
 		*externalapi.VirtualChangeSet, error)
+	UpdatePruningPointByVirtual() error
 
 	MineJSON(r io.Reader, blockType MineJSONBlockType) (tips []*externalapi.DomainHash, err error)
 	ToJSON(w io.Writer) error

domain/consensus/processes/blockprocessor/validate_and_insert_imported_pruning_point_test.go

@@ -44,7 +44,7 @@ func TestValidateAndInsertImportedPruningPoint(t *testing.T) {
 		consensusConfig.K = 0
 		consensusConfig.PruningProofM = 1
 
-		syncConsensuses := func(tcSyncerRef, tcSynceeRef *testapi.TestConsensus) {
+		syncConsensuses := func(tcSyncerRef, tcSynceeRef *testapi.TestConsensus, updatePruningPointJustAfterImportingPruningPoint bool) {
 			tcSyncer, tcSyncee := *tcSyncerRef, *tcSynceeRef
 			pruningPointProof, err := tcSyncer.BuildPruningPointProof()
 			if err != nil {
@@ -236,6 +236,13 @@ func TestValidateAndInsertImportedPruningPoint(t *testing.T) {
 				t.Fatalf("ValidateAndInsertImportedPruningPoint: %+v", err)
 			}
 
+			if updatePruningPointJustAfterImportingPruningPoint {
+				err = synceeStaging.UpdatePruningPointByVirtual()
+				if err != nil {
+					t.Fatal(err)
+				}
+			}
+
 			emptyCoinbase := &externalapi.DomainCoinbaseData{
 				ScriptPublicKey: &externalapi.ScriptPublicKey{
 					Script:  nil,
@@ -386,7 +393,7 @@ func TestValidateAndInsertImportedPruningPoint(t *testing.T) {
 		}
 
 		tcSyncee1Ref := &tcSyncee1
-		syncConsensuses(&tcSyncer, tcSyncee1Ref)
+		syncConsensuses(&tcSyncer, tcSyncee1Ref, false)
 
 		// Test a situation where a consensus with pruned headers syncs another fresh consensus.
 		tcSyncee2, teardownSyncee2, err := factory.NewTestConsensus(consensusConfig, "TestValidateAndInsertPruningPointSyncee2")
@@ -395,7 +402,17 @@ func TestValidateAndInsertImportedPruningPoint(t *testing.T) {
 		}
 		defer teardownSyncee2(false)
 
-		syncConsensuses(tcSyncee1Ref, &tcSyncee2)
+		syncConsensuses(tcSyncee1Ref, &tcSyncee2, false)
+
+		// Check the regular sync but try to update the pruning point after the pruning point was imported. It tests a situation where the node
+		// was restarted before the virtual was resolved and then it calls UpdatePruningPointByVirtual on init.
+		tcSyncee3, teardownSyncee3, err := factory.NewTestConsensus(consensusConfig, "TestValidateAndInsertPruningPointSyncee3")
+		if err != nil {
+			t.Fatalf("Error setting up tcSyncee1: %+v", err)
+		}
+		defer teardownSyncee3(false)
+
+		syncConsensuses(&tcSyncer, &tcSyncee3, true)
 	})
 }
 

domain/consensus/processes/pruningmanager/pruningmanager.go

@@ -191,6 +191,46 @@ func (pm *pruningManager) UpdatePruningPointByVirtual(stagingArea *model.Staging
 	return nil
 }
 
+type blockIteratorFromOneBlock struct {
+	done, isClosed bool
+	hash           *externalapi.DomainHash
+}
+
+func (b *blockIteratorFromOneBlock) First() bool {
+	if b.isClosed {
+		panic("Tried using a closed blockIteratorFromOneBlock")
+	}
+
+	b.done = false
+	return true
+}
+
+func (b *blockIteratorFromOneBlock) Next() bool {
+	if b.isClosed {
+		panic("Tried using a closed blockIteratorFromOneBlock")
+	}
+
+	b.done = true
+	return false
+}
+
+func (b *blockIteratorFromOneBlock) Get() (*externalapi.DomainHash, error) {
+	if b.isClosed {
+		panic("Tried using a closed blockIteratorFromOneBlock")
+	}
+
+	return b.hash, nil
+}
+
+func (b *blockIteratorFromOneBlock) Close() error {
+	if b.isClosed {
+		panic("Tried using a closed blockIteratorFromOneBlock")
+	}
+
+	b.isClosed = true
+	return nil
+}
+
 func (pm *pruningManager) nextPruningPointAndCandidateByBlockHash(stagingArea *model.StagingArea,
 	blockHash, suggestedLowHash *externalapi.DomainHash) (*externalapi.DomainHash, *externalapi.DomainHash, error) {
 
@@ -222,12 +262,12 @@ func (pm *pruningManager) nextPruningPointAndCandidateByBlockHash(stagingArea *m
 		}
 	}
 
-	ghostdagData, err := pm.ghostdagDataStore.Get(pm.databaseContext, stagingArea, blockHash, false)
+	currentPruningPoint, err := pm.pruningStore.PruningPoint(pm.databaseContext, stagingArea)
 	if err != nil {
 		return nil, nil, err
 	}
 
-	currentPruningPoint, err := pm.pruningStore.PruningPoint(pm.databaseContext, stagingArea)
+	ghostdagData, err := pm.ghostdagDataStore.Get(pm.databaseContext, stagingArea, blockHash, false)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -240,9 +280,14 @@ func (pm *pruningManager) nextPruningPointAndCandidateByBlockHash(stagingArea *m
 	// We iterate until the selected parent of the given block, in order to allow a situation where the given block hash
 	// belongs to the virtual. This shouldn't change anything since the max blue score difference between a block and its
 	// selected parent is K, and K << pm.pruningDepth.
-	iterator, err := pm.dagTraversalManager.SelectedChildIterator(stagingArea, ghostdagData.SelectedParent(), lowHash, true)
-	if err != nil {
-		return nil, nil, err
+	var iterator model.BlockIterator
+	if blockHash.Equal(lowHash) {
+		iterator = &blockIteratorFromOneBlock{hash: lowHash}
+	} else {
+		iterator, err = pm.dagTraversalManager.SelectedChildIterator(stagingArea, ghostdagData.SelectedParent(), lowHash, true)
+		if err != nil {
+			return nil, nil, err
+		}
 	}
 	defer iterator.Close()
 

domain/consensus/test_consensus.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"github.com/kaspanet/kaspad/domain/consensus/model"
 	"github.com/kaspanet/kaspad/domain/consensus/utils/hashset"
+	"github.com/kaspanet/kaspad/util/staging"
 	"io"
 
 	"github.com/kaspanet/kaspad/domain/consensus/model/externalapi"
@@ -258,3 +259,16 @@ func (tc *testConsensus) BuildHeaderWithParents(parentHashes []*externalapi.Doma
 
 	return tc.testBlockBuilder.BuildUTXOInvalidHeader(parentHashes)
 }
+
+func (tc *testConsensus) UpdatePruningPointByVirtual() error {
+	tc.lock.Lock()
+	defer tc.lock.Unlock()
+
+	stagingArea := model.NewStagingArea()
+	err := tc.pruningManager.UpdatePruningPointByVirtual(stagingArea)
+	if err != nil {
+		return err
+	}
+
+	return staging.CommitAllChanges(tc.databaseContext, stagingArea)
+}