From 7d0c94d777ece98900419abb574aeaba62d07158 Mon Sep 17 00:00:00 2001
From: "Leo Zhang (zhangchiqing)" <zhangchiqing@gmail.com>
Date: Fri, 7 Jun 2024 20:51:52 -0700
Subject: [PATCH] find trie root

---
 cmd/util/cmd/find-trie-root/cmd.go | 145 +++++++++++++++++++++++++++++
 cmd/util/cmd/root.go               |   2 +
 2 files changed, 147 insertions(+)
 create mode 100644 cmd/util/cmd/find-trie-root/cmd.go

diff --git a/cmd/util/cmd/find-trie-root/cmd.go b/cmd/util/cmd/find-trie-root/cmd.go
new file mode 100644
index 00000000000..8ebb65c03f2
--- /dev/null
+++ b/cmd/util/cmd/find-trie-root/cmd.go
@@ -0,0 +1,145 @@
+package find_trie_root
+
+import (
+	"encoding/hex"
+	"fmt"
+	"math"
+
+	prometheusWAL "github.com/onflow/wal/wal"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/rs/zerolog"
+	"github.com/rs/zerolog/log"
+	"github.com/spf13/cobra"
+
+	"github.com/onflow/flow-go/ledger"
+	"github.com/onflow/flow-go/ledger/common/hash"
+	"github.com/onflow/flow-go/ledger/complete/wal"
+)
+
+var (
+	flagExecutionStateDir string
+	flagRootHash          string
+	flagFrom              int
+	flagTo                int
+)
+
+var Cmd = &cobra.Command{
+	Use:   "find-trie-root",
+	Short: "find trie root",
+	Run:   run,
+}
+
+func init() {
+	Cmd.Flags().StringVarP(&flagExecutionStateDir, "execution-state-dir", "e", "/var/flow/data/execution", "directory to the execution state")
+	_ = Cmd.MarkFlagRequired("execution-state-dir")
+
+	Cmd.Flags().StringVar(&flagRootHash, "root-hash", "",
+		"ledger root hash (hex-encoded, 64 characters)")
+	_ = Cmd.MarkFlagRequired("root-hash")
+
+	Cmd.Flags().IntVar(&flagFrom, "from", 0, "from segment")
+	Cmd.Flags().IntVar(&flagTo, "to", math.MaxInt32, "to segment")
+}
+
+func run(*cobra.Command, []string) {
+	rootHash, err := parseInput(flagRootHash)
+	if err != nil {
+		log.Fatal().Err(err).Msg("cannot parse input")
+	}
+
+	segment, offset, err := searchRootHashInSegments(rootHash, flagExecutionStateDir, flagFrom, flagTo)
+	if err != nil {
+		log.Fatal().Err(err).Msg("cannot find root hash in segments")
+	}
+	log.Info().Msgf("found root hash in segment %d at offset %d", segment, offset)
+}
+
+func parseInput(rootHashStr string) (ledger.RootHash, error) {
+	rootHashBytes, err := hex.DecodeString(rootHashStr)
+	if err != nil {
+		return ledger.RootHash(hash.DummyHash), fmt.Errorf("cannot decode root hash: %w", err)
+	}
+	rootHash, err := ledger.ToRootHash(rootHashBytes)
+	if err != nil {
+		return ledger.RootHash(hash.DummyHash), fmt.Errorf("invalid root hash: %w", err)
+	}
+	return rootHash, nil
+}
+
+func searchRootHashInSegments(
+	expectedHash ledger.RootHash,
+	dir string,
+	wantFrom, wantTo int,
+) (int, int64, error) {
+	log := zerolog.Logger{}
+	w, err := prometheusWAL.NewSize(log, prometheus.DefaultRegisterer, dir, wal.SegmentSize, false)
+	if err != nil {
+		return 0, 0, fmt.Errorf("cannot create WAL: %w", err)
+	}
+
+	from, to, err := prometheusWAL.Segments(dir)
+	if err != nil {
+		return 0, 0, fmt.Errorf("cannot get segments: %w", err)
+	}
+
+	if wantFrom > to {
+		return 0, 0, fmt.Errorf("from segment %d is greater than the last segment %d", wantFrom, to)
+	}
+
+	if wantTo < from {
+		return 0, 0, fmt.Errorf("to segment %d is less than the first segment %d", wantTo, from)
+	}
+
+	if wantFrom > from {
+		from = wantFrom
+	}
+
+	if wantTo < to {
+		to = wantTo
+	}
+
+	log.Info().
+		Str("dir", dir).
+		Int("from", from).
+		Int("to", to).
+		Int("want-from", wantFrom).
+		Int("want-to", wantTo).
+		Msgf("searching for trie root hash %x in segments [%d,%d]", expectedHash, wantFrom, wantTo)
+
+	sr, err := prometheusWAL.NewSegmentsRangeReader(log, prometheusWAL.SegmentRange{
+		Dir:   w.Dir(),
+		First: from,
+		Last:  to,
+	})
+
+	if err != nil {
+		return 0, 0, fmt.Errorf("cannot create WAL segments reader: %w", err)
+	}
+
+	defer sr.Close()
+
+	reader := prometheusWAL.NewReader(sr)
+
+	for reader.Next() {
+		record := reader.Record()
+		operation, rootHash, _, err := wal.Decode(record)
+		if err != nil {
+			return 0, 0, fmt.Errorf("cannot decode LedgerWAL record: %w", err)
+		}
+		switch operation {
+		case wal.WALUpdate:
+			if rootHash == expectedHash {
+				log.Info().Msgf("found expected trie root hash %x", rootHash)
+				return reader.Segment(), reader.Offset(), nil
+			}
+		default:
+		}
+
+		err = reader.Err()
+		if err != nil {
+			return 0, 0, fmt.Errorf("cannot read LedgerWAL: %w", err)
+		}
+	}
+
+	return 0, 0, fmt.Errorf("finish reading all segment files from %d to %d, but not found", from, to)
+}
diff --git a/cmd/util/cmd/root.go b/cmd/util/cmd/root.go
index ba18f16d1f2..4ed05b17e18 100644
--- a/cmd/util/cmd/root.go
+++ b/cmd/util/cmd/root.go
@@ -18,6 +18,7 @@ import (
 	ledger_json_exporter "github.com/onflow/flow-go/cmd/util/cmd/export-json-execution-state"
 	export_json_transactions "github.com/onflow/flow-go/cmd/util/cmd/export-json-transactions"
 	find_inconsistent_result "github.com/onflow/flow-go/cmd/util/cmd/find-inconsistent-result"
+	find_trie_root "github.com/onflow/flow-go/cmd/util/cmd/find-trie-root"
 	read_badger "github.com/onflow/flow-go/cmd/util/cmd/read-badger/cmd"
 	read_execution_state "github.com/onflow/flow-go/cmd/util/cmd/read-execution-state"
 	read_hotstuff "github.com/onflow/flow-go/cmd/util/cmd/read-hotstuff/cmd"
@@ -81,6 +82,7 @@ func addCommands() {
 	rootCmd.AddCommand(export_json_transactions.Cmd)
 	rootCmd.AddCommand(read_hotstuff.RootCmd)
 	rootCmd.AddCommand(find_inconsistent_result.Cmd)
+	rootCmd.AddCommand(find_trie_root.Cmd)
 	rootCmd.AddCommand(update_commitment.Cmd)
 }