From 92e3c56e7be26aac4a25859f55f234aadeec7dbf Mon Sep 17 00:00:00 2001
From: Martin Holst Swende <martin@swende.se>
Date: Fri, 8 Apr 2022 15:08:46 +0200
Subject: [PATCH] cmd/geth: inspect snapshot dangling storage (#24643)

* cmd/geth: inspect snapshot dangling storage

* cmd/geth: make verify-state invoke verify-dangling
---
 cmd/geth/snapshot.go | 92 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 92 insertions(+)

diff --git a/cmd/geth/snapshot.go b/cmd/geth/snapshot.go
index d0539eeff7..9be50a20c5 100644
--- a/cmd/geth/snapshot.go
+++ b/cmd/geth/snapshot.go
@@ -20,6 +20,7 @@ import (
 	"bytes"
 	"encoding/json"
 	"errors"
+	"fmt"
 	"os"
 	"time"
 
@@ -31,6 +32,7 @@ import (
 	"github.com/ethereum/go-ethereum/core/state/snapshot"
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/crypto"
+	"github.com/ethereum/go-ethereum/ethdb"
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/rlp"
 	"github.com/ethereum/go-ethereum/trie"
@@ -102,6 +104,25 @@ geth snapshot verify-state <state-root>
 will traverse the whole accounts and storages set based on the specified
 snapshot and recalculate the root hash of state for verification.
 In other words, this command does the snapshot to trie conversion.
+`,
+			},
+			{
+				Name:      "check-dangling-storage",
+				Usage:     "Check that there is no 'dangling' snap storage",
+				ArgsUsage: "<root>",
+				Action:    utils.MigrateFlags(checkDanglingStorage),
+				Category:  "MISCELLANEOUS COMMANDS",
+				Flags: []cli.Flag{
+					utils.DataDirFlag,
+					utils.AncientFlag,
+					utils.RopstenFlag,
+					utils.SepoliaFlag,
+					utils.RinkebyFlag,
+					utils.GoerliFlag,
+				},
+				Description: `
+geth snapshot check-dangling-storage <state-root> traverses the snap storage 
+data, and verifies that all snapshot storage data has a corresponding account. 
 `,
 			},
 			{
@@ -242,6 +263,77 @@ func verifyState(ctx *cli.Context) error {
 		return err
 	}
 	log.Info("Verified the state", "root", root)
+	if err := checkDangling(chaindb, snaptree.Snapshot(root)); err != nil {
+		log.Error("Dangling snap storage check failed", "root", root, "err", err)
+		return err
+	}
+	return nil
+}
+
+// checkDanglingStorage iterates the snap storage data, and verifies that all
+// storage also has corresponding account data.
+func checkDanglingStorage(ctx *cli.Context) error {
+	stack, _ := makeConfigNode(ctx)
+	defer stack.Close()
+
+	chaindb := utils.MakeChainDatabase(ctx, stack, true)
+	headBlock := rawdb.ReadHeadBlock(chaindb)
+	if headBlock == nil {
+		log.Error("Failed to load head block")
+		return errors.New("no head block")
+	}
+	snaptree, err := snapshot.New(chaindb, trie.NewDatabase(chaindb), 256, headBlock.Root(), false, false, false)
+	if err != nil {
+		log.Error("Failed to open snapshot tree", "err", err)
+		return err
+	}
+	if ctx.NArg() > 1 {
+		log.Error("Too many arguments given")
+		return errors.New("too many arguments")
+	}
+	var root = headBlock.Root()
+	if ctx.NArg() == 1 {
+		root, err = parseRoot(ctx.Args()[0])
+		if err != nil {
+			log.Error("Failed to resolve state root", "err", err)
+			return err
+		}
+	}
+	return checkDangling(chaindb, snaptree.Snapshot(root))
+}
+
+func checkDangling(chaindb ethdb.Database, snap snapshot.Snapshot) error {
+	log.Info("Checking dangling snapshot storage")
+	var (
+		lastReport = time.Now()
+		start      = time.Now()
+		lastKey    []byte
+		it         = rawdb.NewKeyLengthIterator(chaindb.NewIterator(rawdb.SnapshotStoragePrefix, nil), 1+2*common.HashLength)
+	)
+	defer it.Release()
+	for it.Next() {
+		k := it.Key()
+		accKey := k[1:33]
+		if bytes.Equal(accKey, lastKey) {
+			// No need to look up for every slot
+			continue
+		}
+		lastKey = common.CopyBytes(accKey)
+		if time.Since(lastReport) > time.Second*8 {
+			log.Info("Iterating snap storage", "at", fmt.Sprintf("%#x", accKey), "elapsed", common.PrettyDuration(time.Since(start)))
+			lastReport = time.Now()
+		}
+		data, err := snap.AccountRLP(common.BytesToHash(accKey))
+		if err != nil {
+			log.Error("Error loading snap storage data", "account", fmt.Sprintf("%#x", accKey), "err", err)
+			return err
+		}
+		if len(data) == 0 {
+			log.Error("Dangling storage - missing account", "account", fmt.Sprintf("%#x", accKey), "storagekey", fmt.Sprintf("%#x", k))
+			return fmt.Errorf("dangling snapshot storage account %#x", accKey)
+		}
+	}
+	log.Info("Verified the snapshot storage", "root", snap.Root(), "time", common.PrettyDuration(time.Since(start)), "err", it.Error())
 	return nil
 }