From ea10fc679b094ae8173e3402f04708f2d5953bfb Mon Sep 17 00:00:00 2001
From: Ekaterina Pavlova <ekt@morphbits.io>
Date: Mon, 29 Jul 2024 17:26:38 +0800
Subject: [PATCH] cli: add command for traversing MPT for the current state

Traverse MPT for the current state and dump key/value pairs into file.
Example usage: `./bin/neo-go db traverse --config-file config/protocol
.mainnet.yml`.

Close #3519

Signed-off-by: Ekaterina Pavlova <ekt@morphbits.io>
---
 cli/server/mptdump.go | 87 +++++++++++++++++++++++++++++++++++++++++++
 cli/server/server.go  | 13 +++++++
 2 files changed, 100 insertions(+)
 create mode 100644 cli/server/mptdump.go

diff --git a/cli/server/mptdump.go b/cli/server/mptdump.go
new file mode 100644
index 0000000000..cab3486a6c
--- /dev/null
+++ b/cli/server/mptdump.go
@@ -0,0 +1,87 @@
+package server
+
+import (
+	"encoding/hex"
+	"encoding/json"
+	"fmt"
+	"os"
+	"time"
+
+	"github.com/nspcc-dev/neo-go/cli/options"
+	"github.com/nspcc-dev/neo-go/pkg/core/mpt"
+	"github.com/nspcc-dev/neo-go/pkg/core/storage"
+	"github.com/urfave/cli/v2"
+	"go.uber.org/zap"
+)
+
+// KVPair represents a key-value pair.
+type KVPair struct {
+	Key   string `json:"key"`
+	Value string `json:"value"`
+}
+
+// TraverseMPT collects key-value pairs from the TrieStore and returns them.
+func TraverseMPT(store *mpt.TrieStore, encoder *json.Encoder) error {
+	prefix := []byte{byte(storage.STStorage)}
+	rng := storage.SeekRange{Prefix: prefix}
+
+	store.Seek(rng, func(k, v []byte) bool {
+		kvPair := KVPair{
+			Key:   hex.EncodeToString(k),
+			Value: hex.EncodeToString(v),
+		}
+		if err := encoder.Encode(kvPair); err != nil {
+			fmt.Printf("error encoding key-value pair: %v\n", err)
+			return false
+		}
+		return true
+	})
+	return nil
+}
+
+// traverseMPT handles the CLI command to traverse the MPT and dump key-value pairs.
+func traverseMPT(ctx *cli.Context) error {
+	logger := zap.NewExample()
+	cfg, err := options.GetConfigFromContext(ctx)
+	if err != nil {
+		return cli.Exit(err, 1)
+	}
+
+	chain, store, err := initBlockChain(cfg, logger)
+	if err != nil {
+		return cli.Exit(err, 1)
+	}
+	defer store.Close()
+	defer chain.Close()
+
+	stateModule := chain.GetStateModule()
+	stateRoot := stateModule.CurrentLocalStateRoot()
+	stateRootHash := stateRoot
+
+	trieStore := mpt.NewTrieStore(stateRootHash, mpt.ModeAll, store)
+
+	outputFile := ctx.String("out")
+	if outputFile == "" {
+		outputFile = "kv_pairs.json"
+	}
+
+	file, err := os.Create(outputFile)
+	if err != nil {
+		return fmt.Errorf("error creating file: %w", err)
+	}
+	defer file.Close()
+
+	encoder := json.NewEncoder(file)
+	encoder.SetIndent("", "  ")
+
+	startTime := time.Now()
+	fmt.Println(startTime)
+	err = TraverseMPT(trieStore, encoder)
+	if err != nil {
+		return cli.Exit(err, 1)
+	}
+
+	duration := time.Since(startTime)
+	fmt.Printf("MPT key-value pairs successfully dumped to %s in %s\n", outputFile, duration)
+	return nil
+}
diff --git a/cli/server/server.go b/cli/server/server.go
index 882423405b..962ce6fde4 100644
--- a/cli/server/server.go
+++ b/cli/server/server.go
@@ -119,6 +119,19 @@ func NewCommands() []*cli.Command {
 					Action:    resetDB,
 					Flags:     cfgHeightFlags,
 				},
+				{
+					Name:      "traverse",
+					Usage:     "Traverse the MPT and dump key-value pairs to a file",
+					UsageText: "neo-go db traverse [--out file] [--config-path path] [-p/-m/-t] [--config-file file]",
+					Action:    traverseMPT,
+					Flags: append(cfgFlags,
+						&cli.StringFlag{
+							Name:    "out",
+							Aliases: []string{"o"},
+							Usage:   "Output file (default: kv_pairs.json)",
+						},
+					),
+				},
 			},
 		},
 	}