From b6b6f52ec8608e1a694357357c3f1fde669f1e6d Mon Sep 17 00:00:00 2001
From: rjl493456442 <garyrong0905@gmail.com>
Date: Wed, 14 Mar 2018 20:15:30 +0800
Subject: [PATCH] cmd: implement preimage dump and import cmds

---
 cmd/geth/chaincmd.go | 113 +++++++++++++++++++++++++++++++++++++++++++
 cmd/geth/main.go     |   2 +
 ethdb/database.go    |   6 +++
 3 files changed, 121 insertions(+)

diff --git a/cmd/geth/chaincmd.go b/cmd/geth/chaincmd.go
index c9ab72b6d..692cc2d8d 100644
--- a/cmd/geth/chaincmd.go
+++ b/cmd/geth/chaincmd.go
@@ -40,6 +40,11 @@ import (
 	"gopkg.in/urfave/cli.v1"
 )
 
+var (
+	// secureKeyPrefix is the database key prefix used to store trie node preimages.
+	secureKeyPrefix = []byte("secure-key-")
+)
+
 var (
 	initCommand = cli.Command{
 		Action:    utils.MigrateFlags(initGenesis),
@@ -141,6 +146,34 @@ Remove blockchain and state databases`,
 The arguments are interpreted as block numbers or hashes.
 Use "ethereum dump 0" to dump the genesis block.`,
 	}
+	preimageDumpCommand = cli.Command{
+		Action:    utils.MigrateFlags(dumpPreimage),
+		Name:      "preimagedump",
+		Usage:     "Dump the preimage database in json format",
+		ArgsUsage: "<dumpfile>",
+		Flags: []cli.Flag{
+			utils.DataDirFlag,
+			utils.CacheFlag,
+			utils.LightModeFlag,
+		},
+		Category: "BLOCKCHAIN COMMANDS",
+		Description: `
+Dump the preimage database in json format`,
+	}
+	preimageImportCommand = cli.Command{
+		Action:    utils.MigrateFlags(importPreimage),
+		Name:      "preimageimport",
+		Usage:     "Import the preimage data from the specified file",
+		ArgsUsage: "<datafile>",
+		Flags: []cli.Flag{
+			utils.DataDirFlag,
+			utils.CacheFlag,
+			utils.LightModeFlag,
+		},
+		Category: "BLOCKCHAIN COMMANDS",
+		Description: `
+Import the preimage data from the specified file`,
+	}
 )
 
 // initGenesis will initialise the given JSON format genesis file and writes it as
@@ -406,6 +439,86 @@ func dump(ctx *cli.Context) error {
 	return nil
 }
 
+// PreimageEntry represents a map between preimage and hash.
+type PreimageEntry struct {
+	Hash     string `json:"hash"`
+	Preimage string `json:"preimage"`
+}
+
+// dumpPreimage dumps the preimage data to specified json file in streaming way.
+func dumpPreimage(ctx *cli.Context) error {
+	// Make sure the export json file has been specified.
+	if len(ctx.Args()) < 1 {
+		utils.Fatalf("This command requires an argument.")
+	}
+
+	// Encode preimage data to json file in streaming way.
+	file, err := os.Create(ctx.Args().First())
+	if err != nil {
+		return err
+	}
+	encoder := json.NewEncoder(file)
+
+	stack := makeFullNode(ctx)
+	db := utils.MakeChainDatabase(ctx, stack)
+
+	// Dump all preimage entries.
+	it := db.(*ethdb.LDBDatabase).NewIteratorByPrefix(secureKeyPrefix)
+	for it.Next() {
+		hash := it.Key()[len(secureKeyPrefix):]
+		if err := encoder.Encode(PreimageEntry{common.Bytes2Hex(hash), common.Bytes2Hex(it.Value())}); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// importPreimages imports preimage data from the specified file.
+func importPreimage(ctx *cli.Context) error {
+	// Make sure the export json file has been specified.
+	if len(ctx.Args()) < 1 {
+		utils.Fatalf("This command requires an argument.")
+	}
+
+	// Decode the preimage data in streaming way.
+	file, err := os.Open(ctx.Args().First())
+	if err != nil {
+		return err
+	}
+	decoder := json.NewDecoder(file)
+
+	stack := makeFullNode(ctx)
+	db := utils.MakeChainDatabase(ctx, stack)
+
+	var (
+		entry     PreimageEntry
+		preimages = make(map[common.Hash][]byte)
+	)
+
+	for decoder.More() {
+		if err := decoder.Decode(&entry); err != nil {
+			return err
+		}
+		preimages[common.HexToHash(entry.Hash)] = common.Hex2Bytes(entry.Preimage)
+		// Flush to database in batch
+		if len(preimages) > 1024 {
+			err := core.WritePreimages(db, 0, preimages)
+			if err != nil {
+				return err
+			}
+			preimages = make(map[common.Hash][]byte)
+		}
+	}
+	// Flush the last batch preimage data
+	if len(preimages) > 0 {
+		err := core.WritePreimages(db, 0, preimages)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 // hashish returns true for strings that look like hashes.
 func hashish(x string) bool {
 	_, err := strconv.Atoi(x)
diff --git a/cmd/geth/main.go b/cmd/geth/main.go
index f5a3fa941..6e234a704 100644
--- a/cmd/geth/main.go
+++ b/cmd/geth/main.go
@@ -158,6 +158,8 @@ func init() {
 		copydbCommand,
 		removedbCommand,
 		dumpCommand,
+		preimageDumpCommand,
+		preimageImportCommand,
 		// See monitorcmd.go:
 		monitorCommand,
 		// See accountcmd.go:
diff --git a/ethdb/database.go b/ethdb/database.go
index 8c557e482..d0256c56f 100644
--- a/ethdb/database.go
+++ b/ethdb/database.go
@@ -29,6 +29,7 @@ import (
 	"github.com/syndtr/goleveldb/leveldb/filter"
 	"github.com/syndtr/goleveldb/leveldb/iterator"
 	"github.com/syndtr/goleveldb/leveldb/opt"
+	"github.com/syndtr/goleveldb/leveldb/util"
 )
 
 var OpenFileLimit = 64
@@ -121,6 +122,11 @@ func (db *LDBDatabase) NewIterator() iterator.Iterator {
 	return db.db.NewIterator(nil, nil)
 }
 
+// NewIteratorByPrefix returns a iterator to iterate over subset of database content with a particular prefix.
+func (db *LDBDatabase) NewIteratorByPrefix(prefix []byte) iterator.Iterator {
+	return db.db.NewIterator(util.BytesPrefix(prefix), nil)
+}
+
 func (db *LDBDatabase) Close() {
 	// Stop the metrics collection to avoid internal database races
 	db.quitLock.Lock()