From fb0226d293d3b44de5afcb06b6a5e37d0efa0bb5 Mon Sep 17 00:00:00 2001
From: battlmonstr <battlmonstr@users.noreply.github.com>
Date: Mon, 15 Jan 2024 12:37:37 +0100
Subject: [PATCH] polygon/sync: move PeersWithBlockNumInfo and mocks, refactor
 bor.GetRootHash (#9232)

---
 polygon/bor/bor.go                            | 22 +++++++++-----
 polygon/sync/canonical_chain_builder.go       |  2 +-
 .../canonical_chain_builder_mock.go           |  4 +--
 polygon/sync/db.go                            |  2 +-
 polygon/sync/{mock => }/db_mock.go            |  4 +--
 polygon/sync/header_downloader.go             |  9 +++---
 polygon/sync/header_downloader_test.go        | 29 +++++++++----------
 polygon/sync/header_verifier.go               |  5 ----
 polygon/sync/heimdall.go                      |  2 +-
 polygon/sync/{mock => }/heimdall_mock.go      |  4 +--
 .../peer_with_block_num_info.go               |  2 +-
 polygon/sync/sentry.go                        |  5 ++--
 polygon/sync/{mock => }/sentry_mock.go        |  9 +++---
 polygon/sync/state_point_headers_verifier.go  | 22 ++++++++++++++
 14 files changed, 70 insertions(+), 51 deletions(-)
 rename polygon/sync/{mock => }/canonical_chain_builder_mock.go (98%)
 rename polygon/sync/{mock => }/db_mock.go (95%)
 delete mode 100644 polygon/sync/header_verifier.go
 rename polygon/sync/{mock => }/heimdall_mock.go (98%)
 rename polygon/sync/{peerinfo => }/peer_with_block_num_info.go (96%)
 rename polygon/sync/{mock => }/sentry_mock.go (91%)
 create mode 100644 polygon/sync/state_point_headers_verifier.go

diff --git a/polygon/bor/bor.go b/polygon/bor/bor.go
index df251b81d..90fb15965 100644
--- a/polygon/bor/bor.go
+++ b/polygon/bor/bor.go
@@ -1392,12 +1392,23 @@ func (c *Bor) GetRootHash(ctx context.Context, tx kv.Tx, start, end uint64) (str
 	if start > end || end > currentHeaderNumber {
 		return "", &valset.InvalidStartEndBlockError{Start: start, End: end, CurrentHeader: currentHeaderNumber}
 	}
-	blockHeaders := make([]*types.Header, end-start+1)
+	blockHeaders := make([]*types.Header, length)
 	for number := start; number <= end; number++ {
 		blockHeaders[number-start], _ = c.getHeaderByNumber(ctx, tx, number)
 	}
 
-	headers := make([][32]byte, NextPowerOfTwo(length))
+	hash, err := ComputeHeadersRootHash(blockHeaders)
+	if err != nil {
+		return "", err
+	}
+
+	hashStr := hex.EncodeToString(hash)
+	c.rootHashCache.Add(cacheKey, hashStr)
+	return hashStr, nil
+}
+
+func ComputeHeadersRootHash(blockHeaders []*types.Header) ([]byte, error) {
+	headers := make([][32]byte, NextPowerOfTwo(uint64(len(blockHeaders))))
 	for i := 0; i < len(blockHeaders); i++ {
 		blockHeader := blockHeaders[i]
 		header := crypto.Keccak256(AppendBytes32(
@@ -1413,13 +1424,10 @@ func (c *Bor) GetRootHash(ctx context.Context, tx kv.Tx, start, end uint64) (str
 	}
 	tree := merkle.NewTreeWithOpts(merkle.TreeOptions{EnableHashSorting: false, DisableHashLeaves: true})
 	if err := tree.Generate(Convert(headers), sha3.NewLegacyKeccak256()); err != nil {
-		return "", err
+		return nil, err
 	}
-	root := hex.EncodeToString(tree.Root().Hash)
 
-	c.rootHashCache.Add(cacheKey, root)
-
-	return root, nil
+	return tree.Root().Hash, nil
 }
 
 func (c *Bor) getHeaderByNumber(ctx context.Context, tx kv.Tx, number uint64) (*types.Header, error) {
diff --git a/polygon/sync/canonical_chain_builder.go b/polygon/sync/canonical_chain_builder.go
index b6575544b..541100c9b 100644
--- a/polygon/sync/canonical_chain_builder.go
+++ b/polygon/sync/canonical_chain_builder.go
@@ -11,7 +11,7 @@ import (
 	"github.com/ledgerwatch/erigon/polygon/bor"
 )
 
-//go:generate mockgen -destination=./mock/canonical_chain_builder_mock.go -package=mock . CanonicalChainBuilder
+//go:generate mockgen -destination=./canonical_chain_builder_mock.go -package=sync . CanonicalChainBuilder
 type CanonicalChainBuilder interface {
 	Reset(root *types.Header)
 	ContainsHash(hash libcommon.Hash) bool
diff --git a/polygon/sync/mock/canonical_chain_builder_mock.go b/polygon/sync/canonical_chain_builder_mock.go
similarity index 98%
rename from polygon/sync/mock/canonical_chain_builder_mock.go
rename to polygon/sync/canonical_chain_builder_mock.go
index 2985c4467..0ee6e49e3 100644
--- a/polygon/sync/mock/canonical_chain_builder_mock.go
+++ b/polygon/sync/canonical_chain_builder_mock.go
@@ -1,8 +1,8 @@
 // Code generated by MockGen. DO NOT EDIT.
 // Source: github.com/ledgerwatch/erigon/polygon/sync (interfaces: CanonicalChainBuilder)
 
-// Package mock is a generated GoMock package.
-package mock
+// Package sync is a generated GoMock package.
+package sync
 
 import (
 	reflect "reflect"
diff --git a/polygon/sync/db.go b/polygon/sync/db.go
index 560ab2bc1..9fc3ed9bc 100644
--- a/polygon/sync/db.go
+++ b/polygon/sync/db.go
@@ -2,7 +2,7 @@ package sync
 
 import "github.com/ledgerwatch/erigon/core/types"
 
-//go:generate mockgen -destination=./mock/db_mock.go -package=mock . DB
+//go:generate mockgen -destination=./db_mock.go -package=sync . DB
 type DB interface {
 	WriteHeaders(headers []*types.Header) error
 }
diff --git a/polygon/sync/mock/db_mock.go b/polygon/sync/db_mock.go
similarity index 95%
rename from polygon/sync/mock/db_mock.go
rename to polygon/sync/db_mock.go
index 22e6fa6b4..2993c959e 100644
--- a/polygon/sync/mock/db_mock.go
+++ b/polygon/sync/db_mock.go
@@ -1,8 +1,8 @@
 // Code generated by MockGen. DO NOT EDIT.
 // Source: github.com/ledgerwatch/erigon/polygon/sync (interfaces: DB)
 
-// Package mock is a generated GoMock package.
-package mock
+// Package sync is a generated GoMock package.
+package sync
 
 import (
 	reflect "reflect"
diff --git a/polygon/sync/header_downloader.go b/polygon/sync/header_downloader.go
index 76a8f29f3..2f67e2684 100644
--- a/polygon/sync/header_downloader.go
+++ b/polygon/sync/header_downloader.go
@@ -13,12 +13,11 @@ import (
 
 	"github.com/ledgerwatch/erigon-lib/common"
 	"github.com/ledgerwatch/erigon/core/types"
-	"github.com/ledgerwatch/erigon/polygon/sync/peerinfo"
 )
 
 const headerDownloaderLogPrefix = "HeaderDownloader"
 
-func NewHeaderDownloader(logger log.Logger, sentry Sentry, db DB, heimdall Heimdall, verify HeaderVerifier) *HeaderDownloader {
+func NewHeaderDownloader(logger log.Logger, sentry Sentry, db DB, heimdall Heimdall, verify StatePointHeadersVerifier) *HeaderDownloader {
 	statePointHeadersMemo, err := lru.New[common.Hash, []*types.Header](sentry.MaxPeers())
 	if err != nil {
 		panic(err)
@@ -39,7 +38,7 @@ type HeaderDownloader struct {
 	sentry                Sentry
 	db                    DB
 	heimdall              Heimdall
-	verify                HeaderVerifier
+	verify                StatePointHeadersVerifier
 	statePointHeadersMemo *lru.Cache[common.Hash, []*types.Header] // statePoint.rootHash->[headers part of state point]
 }
 
@@ -194,9 +193,9 @@ func (hd *HeaderDownloader) downloadUsingStatePoints(ctx context.Context, stateP
 }
 
 // choosePeers assumes peers are sorted in ascending order based on block num
-func (hd *HeaderDownloader) choosePeers(peers peerinfo.PeersWithBlockNumInfo, statePoints statePoints) peerinfo.PeersWithBlockNumInfo {
+func (hd *HeaderDownloader) choosePeers(peers PeersWithBlockNumInfo, statePoints statePoints) PeersWithBlockNumInfo {
 	var peersIdx int
-	chosenPeers := make(peerinfo.PeersWithBlockNumInfo, 0, len(peers))
+	chosenPeers := make(PeersWithBlockNumInfo, 0, len(peers))
 	for _, statePoint := range statePoints {
 		if peersIdx >= len(peers) {
 			break
diff --git a/polygon/sync/header_downloader_test.go b/polygon/sync/header_downloader_test.go
index ae650902a..7a24a6ab5 100644
--- a/polygon/sync/header_downloader_test.go
+++ b/polygon/sync/header_downloader_test.go
@@ -12,13 +12,10 @@ import (
 	"github.com/ledgerwatch/log/v3"
 	"github.com/stretchr/testify/require"
 
-	"github.com/ledgerwatch/erigon/polygon/heimdall/checkpoint"
-	"github.com/ledgerwatch/erigon/polygon/heimdall/milestone"
-
 	"github.com/ledgerwatch/erigon-lib/common"
 	"github.com/ledgerwatch/erigon/core/types"
-	"github.com/ledgerwatch/erigon/polygon/sync/mock"
-	"github.com/ledgerwatch/erigon/polygon/sync/peerinfo"
+	"github.com/ledgerwatch/erigon/polygon/heimdall/checkpoint"
+	"github.com/ledgerwatch/erigon/polygon/heimdall/milestone"
 	"github.com/ledgerwatch/erigon/turbo/testlog"
 )
 
@@ -28,10 +25,10 @@ func newHeaderDownloaderTest(t *testing.T) *headerDownloaderTest {
 
 func newHeaderDownloaderTestWithOpts(t *testing.T, opts headerDownloaderTestOpts) *headerDownloaderTest {
 	ctrl := gomock.NewController(t)
-	heimdall := mock.NewMockHeimdall(ctrl)
-	sentry := mock.NewMockSentry(ctrl)
+	heimdall := NewMockHeimdall(ctrl)
+	sentry := NewMockSentry(ctrl)
 	sentry.EXPECT().MaxPeers().Return(100).Times(1)
-	db := mock.NewMockDB(ctrl)
+	db := NewMockDB(ctrl)
 	logger := testlog.Logger(t, log.LvlDebug)
 	headerVerifier := opts.getOrCreateDefaultHeaderVerifier()
 	headerDownloader := NewHeaderDownloader(logger, sentry, db, heimdall, headerVerifier)
@@ -44,10 +41,10 @@ func newHeaderDownloaderTestWithOpts(t *testing.T, opts headerDownloaderTestOpts
 }
 
 type headerDownloaderTestOpts struct {
-	headerVerifier HeaderVerifier
+	headerVerifier StatePointHeadersVerifier
 }
 
-func (opts headerDownloaderTestOpts) getOrCreateDefaultHeaderVerifier() HeaderVerifier {
+func (opts headerDownloaderTestOpts) getOrCreateDefaultHeaderVerifier() StatePointHeadersVerifier {
 	if opts.headerVerifier == nil {
 		return func(_ *statePoint, _ []*types.Header) error {
 			return nil
@@ -58,14 +55,14 @@ func (opts headerDownloaderTestOpts) getOrCreateDefaultHeaderVerifier() HeaderVe
 }
 
 type headerDownloaderTest struct {
-	heimdall         *mock.MockHeimdall
-	sentry           *mock.MockSentry
-	db               *mock.MockDB
+	heimdall         *MockHeimdall
+	sentry           *MockSentry
+	db               *MockDB
 	headerDownloader *HeaderDownloader
 }
 
-func (hdt headerDownloaderTest) fakePeers(count int, blockNums ...*big.Int) peerinfo.PeersWithBlockNumInfo {
-	peers := make(peerinfo.PeersWithBlockNumInfo, count)
+func (hdt headerDownloaderTest) fakePeers(count int, blockNums ...*big.Int) PeersWithBlockNumInfo {
+	peers := make(PeersWithBlockNumInfo, count)
 	for i := range peers {
 		var blockNum *big.Int
 		if i < len(blockNums) {
@@ -74,7 +71,7 @@ func (hdt headerDownloaderTest) fakePeers(count int, blockNums ...*big.Int) peer
 			blockNum = new(big.Int).SetUint64(math.MaxUint64)
 		}
 
-		peers[i] = &peerinfo.PeerWithBlockNumInfo{
+		peers[i] = &PeerWithBlockNumInfo{
 			ID:       fmt.Sprintf("peer%d", i+1),
 			BlockNum: blockNum,
 		}
diff --git a/polygon/sync/header_verifier.go b/polygon/sync/header_verifier.go
deleted file mode 100644
index 6898f3849..000000000
--- a/polygon/sync/header_verifier.go
+++ /dev/null
@@ -1,5 +0,0 @@
-package sync
-
-import "github.com/ledgerwatch/erigon/core/types"
-
-type HeaderVerifier func(statePoint *statePoint, headers []*types.Header) error
diff --git a/polygon/sync/heimdall.go b/polygon/sync/heimdall.go
index b3e3f1e95..b77d58c8e 100644
--- a/polygon/sync/heimdall.go
+++ b/polygon/sync/heimdall.go
@@ -18,7 +18,7 @@ import (
 
 // Heimdall is a wrapper of Heimdall HTTP API
 //
-//go:generate mockgen -destination=./mock/heimdall_mock.go -package=mock . Heimdall
+//go:generate mockgen -destination=./heimdall_mock.go -package=sync . Heimdall
 type Heimdall interface {
 	FetchCheckpoints(ctx context.Context, start uint64) ([]*checkpoint.Checkpoint, error)
 	FetchMilestones(ctx context.Context, start uint64) ([]*milestone.Milestone, error)
diff --git a/polygon/sync/mock/heimdall_mock.go b/polygon/sync/heimdall_mock.go
similarity index 98%
rename from polygon/sync/mock/heimdall_mock.go
rename to polygon/sync/heimdall_mock.go
index fe887f30b..8604e623e 100644
--- a/polygon/sync/mock/heimdall_mock.go
+++ b/polygon/sync/heimdall_mock.go
@@ -1,8 +1,8 @@
 // Code generated by MockGen. DO NOT EDIT.
 // Source: github.com/ledgerwatch/erigon/polygon/sync (interfaces: Heimdall)
 
-// Package mock is a generated GoMock package.
-package mock
+// Package sync is a generated GoMock package.
+package sync
 
 import (
 	context "context"
diff --git a/polygon/sync/peerinfo/peer_with_block_num_info.go b/polygon/sync/peer_with_block_num_info.go
similarity index 96%
rename from polygon/sync/peerinfo/peer_with_block_num_info.go
rename to polygon/sync/peer_with_block_num_info.go
index 643aa078d..1cde2bbf7 100644
--- a/polygon/sync/peerinfo/peer_with_block_num_info.go
+++ b/polygon/sync/peer_with_block_num_info.go
@@ -1,4 +1,4 @@
-package peerinfo
+package sync
 
 import "math/big"
 
diff --git a/polygon/sync/sentry.go b/polygon/sync/sentry.go
index aa8e5198c..ffea66b08 100644
--- a/polygon/sync/sentry.go
+++ b/polygon/sync/sentry.go
@@ -5,13 +5,12 @@ import (
 	"math/big"
 
 	"github.com/ledgerwatch/erigon/core/types"
-	"github.com/ledgerwatch/erigon/polygon/sync/peerinfo"
 )
 
-//go:generate mockgen -destination=./mock/sentry_mock.go -package=mock . Sentry
+//go:generate mockgen -destination=./sentry_mock.go -package=sync . Sentry
 type Sentry interface {
 	MaxPeers() int
-	PeersWithBlockNumInfo() peerinfo.PeersWithBlockNumInfo
+	PeersWithBlockNumInfo() PeersWithBlockNumInfo
 	DownloadHeaders(ctx context.Context, start *big.Int, end *big.Int, peerID string) ([]*types.Header, error)
 	Penalize(peerID string)
 }
diff --git a/polygon/sync/mock/sentry_mock.go b/polygon/sync/sentry_mock.go
similarity index 91%
rename from polygon/sync/mock/sentry_mock.go
rename to polygon/sync/sentry_mock.go
index 09da63358..7070b6297 100644
--- a/polygon/sync/mock/sentry_mock.go
+++ b/polygon/sync/sentry_mock.go
@@ -1,8 +1,8 @@
 // Code generated by MockGen. DO NOT EDIT.
 // Source: github.com/ledgerwatch/erigon/polygon/sync (interfaces: Sentry)
 
-// Package mock is a generated GoMock package.
-package mock
+// Package sync is a generated GoMock package.
+package sync
 
 import (
 	context "context"
@@ -11,7 +11,6 @@ import (
 
 	gomock "github.com/golang/mock/gomock"
 	types "github.com/ledgerwatch/erigon/core/types"
-	peerinfo "github.com/ledgerwatch/erigon/polygon/sync/peerinfo"
 )
 
 // MockSentry is a mock of Sentry interface.
@@ -67,10 +66,10 @@ func (mr *MockSentryMockRecorder) MaxPeers() *gomock.Call {
 }
 
 // PeersWithBlockNumInfo mocks base method.
-func (m *MockSentry) PeersWithBlockNumInfo() peerinfo.PeersWithBlockNumInfo {
+func (m *MockSentry) PeersWithBlockNumInfo() PeersWithBlockNumInfo {
 	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "PeersWithBlockNumInfo")
-	ret0, _ := ret[0].(peerinfo.PeersWithBlockNumInfo)
+	ret0, _ := ret[0].(PeersWithBlockNumInfo)
 	return ret0
 }
 
diff --git a/polygon/sync/state_point_headers_verifier.go b/polygon/sync/state_point_headers_verifier.go
new file mode 100644
index 000000000..5ee550c92
--- /dev/null
+++ b/polygon/sync/state_point_headers_verifier.go
@@ -0,0 +1,22 @@
+package sync
+
+import (
+	"bytes"
+	"fmt"
+
+	"github.com/ledgerwatch/erigon/core/types"
+	"github.com/ledgerwatch/erigon/polygon/bor"
+)
+
+type StatePointHeadersVerifier func(statePoint *statePoint, headers []*types.Header) error
+
+func VerifyStatePointHeaders(statePoint *statePoint, headers []*types.Header) error {
+	rootHash, err := bor.ComputeHeadersRootHash(headers)
+	if err != nil {
+		return fmt.Errorf("VerifyStatePointHeaders: failed to compute headers root hash %w", err)
+	}
+	if !bytes.Equal(rootHash, statePoint.rootHash[:]) {
+		return fmt.Errorf("VerifyStatePointHeaders: bad headers root hash")
+	}
+	return nil
+}