From b5a6532d3f7e6ac3e84d57c8246c0b44cdfbe8cf Mon Sep 17 00:00:00 2001
From: Dhruba Basu <7675102+dhrubabasu@users.noreply.github.com>
Date: Thu, 1 Aug 2024 17:36:24 -0400
Subject: [PATCH] [vms/platformvm] Add `VerifyWithContext` to `Block`s (#3236)

---
 vms/platformvm/block/executor/block.go       | 32 +++++++++++++++---
 vms/platformvm/block/executor/block_state.go |  7 ++--
 vms/platformvm/block/executor/verifier.go    | 34 ++++++++++++--------
 3 files changed, 52 insertions(+), 21 deletions(-)

diff --git a/vms/platformvm/block/executor/block.go b/vms/platformvm/block/executor/block.go
index d4bc1d8ace..76cf6ff907 100644
--- a/vms/platformvm/block/executor/block.go
+++ b/vms/platformvm/block/executor/block.go
@@ -9,11 +9,14 @@ import (
 
 	"github.com/ava-labs/avalanchego/snow/consensus/snowman"
 	"github.com/ava-labs/avalanchego/vms/platformvm/block"
+
+	smblock "github.com/ava-labs/avalanchego/snow/engine/snowman/block"
 )
 
 var (
-	_ snowman.Block       = (*Block)(nil)
-	_ snowman.OracleBlock = (*Block)(nil)
+	_ snowman.Block             = (*Block)(nil)
+	_ snowman.OracleBlock       = (*Block)(nil)
+	_ smblock.WithVerifyContext = (*Block)(nil)
 )
 
 // Exported for testing in platformvm package.
@@ -22,9 +25,25 @@ type Block struct {
 	manager *manager
 }
 
-func (b *Block) Verify(context.Context) error {
+func (*Block) ShouldVerifyWithContext(context.Context) (bool, error) {
+	return true, nil
+}
+
+func (b *Block) VerifyWithContext(_ context.Context, ctx *smblock.Context) error {
+	pChainHeight := uint64(0)
+	if ctx != nil {
+		pChainHeight = ctx.PChainHeight
+	}
+
 	blkID := b.ID()
-	if _, ok := b.manager.blkIDToState[blkID]; ok {
+	if blkState, ok := b.manager.blkIDToState[blkID]; ok {
+		if !blkState.verifiedHeights.Contains(pChainHeight) {
+			// PlatformVM blocks are currently valid regardless of the ProposerVM's
+			// PChainHeight. If this changes, those validity checks should be done prior
+			// to adding [pChainHeight] to [verifiedHeights].
+			blkState.verifiedHeights.Add(pChainHeight)
+		}
+
 		// This block has already been verified.
 		return nil
 	}
@@ -32,9 +51,14 @@ func (b *Block) Verify(context.Context) error {
 	return b.Visit(&verifier{
 		backend:           b.manager.backend,
 		txExecutorBackend: b.manager.txExecutorBackend,
+		pChainHeight:      pChainHeight,
 	})
 }
 
+func (b *Block) Verify(ctx context.Context) error {
+	return b.VerifyWithContext(ctx, nil)
+}
+
 func (b *Block) Accept(context.Context) error {
 	return b.Visit(b.manager.acceptor)
 }
diff --git a/vms/platformvm/block/executor/block_state.go b/vms/platformvm/block/executor/block_state.go
index 9d6b377c26..5f6801d3c1 100644
--- a/vms/platformvm/block/executor/block_state.go
+++ b/vms/platformvm/block/executor/block_state.go
@@ -28,7 +28,8 @@ type blockState struct {
 	onAcceptState state.Diff
 	onAcceptFunc  func()
 
-	inputs         set.Set[ids.ID]
-	timestamp      time.Time
-	atomicRequests map[ids.ID]*atomic.Requests
+	inputs          set.Set[ids.ID]
+	timestamp       time.Time
+	atomicRequests  map[ids.ID]*atomic.Requests
+	verifiedHeights set.Set[uint64]
 }
diff --git a/vms/platformvm/block/executor/verifier.go b/vms/platformvm/block/executor/verifier.go
index a8f2aeaf1d..e238a35c10 100644
--- a/vms/platformvm/block/executor/verifier.go
+++ b/vms/platformvm/block/executor/verifier.go
@@ -34,6 +34,7 @@ var (
 type verifier struct {
 	*backend
 	txExecutorBackend *executor.Backend
+	pChainHeight      uint64
 }
 
 func (v *verifier) BanffAbortBlock(b *block.BanffAbortBlock) error {
@@ -233,9 +234,10 @@ func (v *verifier) ApricotAtomicBlock(b *block.ApricotAtomicBlock) error {
 
 		onAcceptState: atomicExecutor.OnAccept,
 
-		inputs:         atomicExecutor.Inputs,
-		timestamp:      atomicExecutor.OnAccept.GetTimestamp(),
-		atomicRequests: atomicExecutor.AtomicRequests,
+		inputs:          atomicExecutor.Inputs,
+		timestamp:       atomicExecutor.OnAccept.GetTimestamp(),
+		atomicRequests:  atomicExecutor.AtomicRequests,
+		verifiedHeights: set.Of(v.pChainHeight),
 	}
 	return nil
 }
@@ -345,9 +347,10 @@ func (v *verifier) abortBlock(b block.Block) error {
 
 	blkID := b.ID()
 	v.blkIDToState[blkID] = &blockState{
-		statelessBlock: b,
-		onAcceptState:  onAbortState,
-		timestamp:      onAbortState.GetTimestamp(),
+		statelessBlock:  b,
+		onAcceptState:   onAbortState,
+		timestamp:       onAbortState.GetTimestamp(),
+		verifiedHeights: set.Of(v.pChainHeight),
 	}
 	return nil
 }
@@ -362,9 +365,10 @@ func (v *verifier) commitBlock(b block.Block) error {
 
 	blkID := b.ID()
 	v.blkIDToState[blkID] = &blockState{
-		statelessBlock: b,
-		onAcceptState:  onCommitState,
-		timestamp:      onCommitState.GetTimestamp(),
+		statelessBlock:  b,
+		onAcceptState:   onCommitState,
+		timestamp:       onCommitState.GetTimestamp(),
+		verifiedHeights: set.Of(v.pChainHeight),
 	}
 	return nil
 }
@@ -415,8 +419,9 @@ func (v *verifier) proposalBlock(
 		// It is safe to use [b.onAbortState] here because the timestamp will
 		// never be modified by an Apricot Abort block and the timestamp will
 		// always be the same as the Banff Proposal Block.
-		timestamp:      onAbortState.GetTimestamp(),
-		atomicRequests: atomicRequests,
+		timestamp:       onAbortState.GetTimestamp(),
+		atomicRequests:  atomicRequests,
+		verifiedHeights: set.Of(v.pChainHeight),
 	}
 	return nil
 }
@@ -441,9 +446,10 @@ func (v *verifier) standardBlock(
 		onAcceptState: onAcceptState,
 		onAcceptFunc:  onAcceptFunc,
 
-		timestamp:      onAcceptState.GetTimestamp(),
-		inputs:         inputs,
-		atomicRequests: atomicRequests,
+		timestamp:       onAcceptState.GetTimestamp(),
+		inputs:          inputs,
+		atomicRequests:  atomicRequests,
+		verifiedHeights: set.Of(v.pChainHeight),
 	}
 	return nil
 }