From eba5366cefbdf870c5789bce634201f7deb7c0d7 Mon Sep 17 00:00:00 2001
From: Guillaume Ballet <3272758+gballet@users.noreply.github.com>
Date: Thu, 5 Oct 2023 09:47:10 +0200
Subject: [PATCH] cmd/{geth,utils}: add cmd to export preimages in snap
 enumeration order

Co-Authored-By: Ignacio Hagopian <jsign.uy@gmail.com>
---
 cmd/geth/chaincmd.go | 38 ++++++++++++++++++++++++
 cmd/geth/main.go     |  1 +
 cmd/utils/cmd.go     | 69 ++++++++++++++++++++++++++++++++++++++++++++
 cmd/utils/flags.go   |  5 ++++
 4 files changed, 113 insertions(+)

diff --git a/cmd/geth/chaincmd.go b/cmd/geth/chaincmd.go
index aebcc29eb8..2ef20e2afb 100644
--- a/cmd/geth/chaincmd.go
+++ b/cmd/geth/chaincmd.go
@@ -165,6 +165,17 @@ It's deprecated, please use "geth db export" instead.
 		}, utils.DatabaseFlags),
 		Description: `
 This command dumps out the state for a given block (or latest, if none provided).
+`,
+	}
+	exportSnapshotPreimagesCommand = &cli.Command{
+		Action:    exportSnapshotPreimages,
+		Name:      "export-snapshot-preimages",
+		Usage:     "Export the preimage in snapshot enumeration order",
+		ArgsUsage: "<dumpfile>",
+		Flags:     flags.Merge([]cli.Flag{utils.TreeRootFlag}, utils.DatabaseFlags),
+		Description: `
+The export-overlay-preimages command exports hash preimages to a flat file, in exactly
+the expected order for the overlay tree migration.
 `,
 	}
 )
@@ -406,6 +417,33 @@ func exportPreimages(ctx *cli.Context) error {
 	return nil
 }
 
+// exportSnapshotPreimages dumps the preimage data to a flat file.
+func exportSnapshotPreimages(ctx *cli.Context) error {
+	if ctx.Args().Len() < 1 {
+		utils.Fatalf("This command requires an argument.")
+	}
+	stack, _ := makeConfigNode(ctx)
+	defer stack.Close()
+
+	chain, _ := utils.MakeChain(ctx, stack, true)
+
+	var root common.Hash
+	if ctx.String(utils.TreeRootFlag.Name) != "" {
+		rootBytes := common.FromHex(ctx.String(utils.StartKeyFlag.Name))
+		if len(rootBytes) != common.HashLength {
+			return fmt.Errorf("invalid root hash length")
+		}
+		root = common.BytesToHash(rootBytes)
+	}
+
+	start := time.Now()
+	if err := utils.ExportSnapshotPreimages(chain, ctx.Args().First(), root); err != nil {
+		utils.Fatalf("Export error: %v\n", err)
+	}
+	fmt.Printf("Export done in %v\n", time.Since(start))
+	return nil
+}
+
 func parseDumpConfig(ctx *cli.Context, stack *node.Node) (*state.DumpConfig, ethdb.Database, common.Hash, error) {
 	db := utils.MakeChainDatabase(ctx, stack, true)
 	defer db.Close()
diff --git a/cmd/geth/main.go b/cmd/geth/main.go
index 4b26de05a9..a0c7d15b3d 100644
--- a/cmd/geth/main.go
+++ b/cmd/geth/main.go
@@ -209,6 +209,7 @@ func init() {
 		exportCommand,
 		importPreimagesCommand,
 		exportPreimagesCommand,
+		exportSnapshotPreimagesCommand,
 		removedbCommand,
 		dumpCommand,
 		dumpGenesisCommand,
diff --git a/cmd/utils/cmd.go b/cmd/utils/cmd.go
index 16b1260572..601096bfbc 100644
--- a/cmd/utils/cmd.go
+++ b/cmd/utils/cmd.go
@@ -374,6 +374,75 @@ func ExportPreimages(db ethdb.Database, fn string) error {
 	return nil
 }
 
+// ExportSnapshotPreimages exports the preimages corresponding to the enumeration of
+// the snapshot for a given root.
+func ExportSnapshotPreimages(chain *core.BlockChain, fn string, root common.Hash) error {
+	log.Info("Exporting preimages", "file", fn)
+
+	fh, err := os.OpenFile(fn, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, os.ModePerm)
+	if err != nil {
+		return err
+	}
+	defer fh.Close()
+
+	writer := bufio.NewWriter(fh)
+	defer writer.Flush()
+
+	statedb, err := chain.State()
+	if err != nil {
+		return fmt.Errorf("failed to open statedb: %w", err)
+	}
+
+	if root == (common.Hash{}) {
+		root = chain.CurrentBlock().Root
+	}
+
+	accIt, err := chain.Snapshots().AccountIterator(root, common.Hash{})
+	if err != nil {
+		return err
+	}
+	defer accIt.Release()
+
+	count := 0
+	for accIt.Next() {
+		acc, err := types.FullAccount(accIt.Account())
+		if err != nil {
+			return fmt.Errorf("invalid account encountered during traversal: %s", err)
+		}
+		addr := rawdb.ReadPreimage(statedb.Database().DiskDB(), accIt.Hash())
+		if len(addr) != 20 {
+			return fmt.Errorf("addr len is zero is not 32: %d", len(addr))
+		}
+		if _, err := writer.Write(addr); err != nil {
+			return fmt.Errorf("failed to write addr preimage: %w", err)
+		}
+
+		if acc.Root != (common.Hash{}) && acc.Root != types.EmptyRootHash {
+			stIt, err := chain.Snapshots().StorageIterator(root, accIt.Hash(), common.Hash{})
+			if err != nil {
+				return fmt.Errorf("failed to create storage iterator: %w", err)
+			}
+			for stIt.Next() {
+				slotnr := rawdb.ReadPreimage(statedb.Database().DiskDB(), stIt.Hash())
+				if len(slotnr) != 32 {
+					return fmt.Errorf("slotnr not 32 len")
+				}
+				if _, err := writer.Write(slotnr); err != nil {
+					return fmt.Errorf("failed to write slotnr preimage: %w", err)
+				}
+			}
+			stIt.Release()
+		}
+		count++
+		if count%100000 == 0 {
+			log.Info("Last exported account", "account", accIt.Hash())
+		}
+	}
+
+	log.Info("Exported preimages", "file", fn)
+	return nil
+}
+
 // exportHeader is used in the export/import flow. When we do an export,
 // the first element we output is the exportHeader.
 // Whenever a backwards-incompatible change is made, the Version header
diff --git a/cmd/utils/flags.go b/cmd/utils/flags.go
index 9743a7b9ca..c3dcf875b6 100644
--- a/cmd/utils/flags.go
+++ b/cmd/utils/flags.go
@@ -219,6 +219,11 @@ var (
 		Usage: "Max number of elements (0 = no limit)",
 		Value: 0,
 	}
+	TreeRootFlag = &cli.StringFlag{
+		Name:  "roothash",
+		Usage: "Root hash of the tree (if empty, use that of current block)",
+		Value: "",
+	}
 
 	defaultSyncMode = ethconfig.Defaults.SyncMode
 	SnapshotFlag    = &cli.BoolFlag{