From 02e44e714ba76ee01708da1a8a6392d9ba9122a0 Mon Sep 17 00:00:00 2001
From: Steven Allen <steven@stebalien.com>
Date: Mon, 28 Oct 2024 09:57:13 -0700
Subject: [PATCH] test(chain): test forking before/after finality

---
 chain/gen/gen.go               | 21 +++++++++++++-----
 chain/store/checkpoint_test.go | 39 ++++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 5 deletions(-)

diff --git a/chain/gen/gen.go b/chain/gen/gen.go
index 435d942dc18..753e0824bd4 100644
--- a/chain/gen/gen.go
+++ b/chain/gen/gen.go
@@ -77,6 +77,10 @@ type ChainGen struct {
 
 	GetMessages func(*ChainGen) ([]*types.SignedMessage, error)
 
+	// Set to false to allow the chain to advance without updating the state-tree (e.g., this
+	// allows one to keep the power without having to post).
+	AdvanceState bool
+
 	w *wallet.LocalWallet
 
 	eppProvs  map[address.Address]WinningPoStProver
@@ -272,11 +276,12 @@ func NewGeneratorWithSectorsAndUpgradeSchedule(numSectors int, us stmgr.UpgradeS
 		beacon:       beac,
 		w:            w,
 
-		GetMessages: getRandomMessages,
-		Miners:      miners,
-		eppProvs:    mgen,
-		banker:      banker,
-		receivers:   receievers,
+		GetMessages:  getRandomMessages,
+		AdvanceState: true,
+		Miners:       miners,
+		eppProvs:     mgen,
+		banker:       banker,
+		receivers:    receievers,
 
 		CurTipset: gents,
 
@@ -355,6 +360,8 @@ func (cg *ChainGen) nextBlockProof(ctx context.Context, pts *types.TipSet, m add
 	mbi, err := mc.MinerGetBaseInfo(ctx, m, round, pts.Key())
 	if err != nil {
 		return nil, nil, nil, xerrors.Errorf("get miner base info: %w", err)
+	} else if mbi == nil {
+		return nil, nil, nil, nil
 	}
 
 	entries := mbi.BeaconEntries
@@ -523,6 +530,10 @@ func (cg *ChainGen) makeBlock(parents *types.TipSet, m address.Address, vrfticke
 		return nil, err
 	}
 
+	if !cg.AdvanceState {
+		fblk.Header.ParentStateRoot = parents.ParentState()
+	}
+
 	return fblk, err
 }
 
diff --git a/chain/store/checkpoint_test.go b/chain/store/checkpoint_test.go
index 8b775d334a0..4629affdaed 100644
--- a/chain/store/checkpoint_test.go
+++ b/chain/store/checkpoint_test.go
@@ -7,6 +7,7 @@ import (
 
 	"github.com/stretchr/testify/require"
 
+	"github.com/filecoin-project/lotus/chain/actors/policy"
 	"github.com/filecoin-project/lotus/chain/gen"
 )
 
@@ -20,6 +21,7 @@ func TestChainCheckpoint(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+	cg.AdvanceState = false
 
 	// Let the first miner mine some blocks.
 	last := cg.CurTipset.TipSet()
@@ -83,4 +85,41 @@ func TestChainCheckpoint(t *testing.T) {
 
 	head = cs.GetHeaviestTipSet()
 	require.True(t, head.Equals(checkpoint))
+
+	// Now extend the fork 900 epochs into the future.
+	for i := 0; i < int(policy.ChainFinality)+10; i++ {
+		ts, err := cg.NextTipSetFromMiners(last, cg.Miners[1:], 0)
+		require.NoError(t, err)
+
+		last = ts.TipSet.TipSet()
+	}
+
+	// Try to re-checkpoint to the long fork. This will work because we only have to revert a
+	// single epoch.
+	err = cs.SetCheckpoint(ctx, last)
+	require.NoError(t, err)
+
+	head = cs.GetHeaviestTipSet()
+	require.True(t, head.Equals(last))
+
+	// Now try to go back to the checkpoint. This should fail because it's too far in the past
+	// on the wrong fork.
+	err = cs.SetCheckpoint(ctx, checkpoint)
+	require.Error(t, err)
+
+	// Now extend the checkpoint chain to the same tipset.
+	for checkpoint.Height() < last.Height() {
+		ts, err := cg.NextTipSetFromMiners(checkpoint, cg.Miners[1:], 0)
+		require.NoError(t, err)
+
+		checkpoint = ts.TipSet.TipSet()
+	}
+
+	// We should still refuse to switch.
+	err = cs.SetCheckpoint(ctx, checkpoint)
+	require.Error(t, err)
+
+	// But it should be possible to set a checkpoint on a common chain
+	err = cs.SetCheckpoint(ctx, checkpointParents)
+	require.NoError(t, err)
 }