traverse with --account

Signed-off-by: jsvisa <delweng@gmail.com>

cmd/geth/snapshot.go

@@ -103,17 +103,24 @@ information about the specified address.
 			{
 				Name:      "traverse-state",
 				Usage:     "Traverse the state with given root hash and perform quick verification",
-				ArgsUsage: "<root> [accountHash|accountAddress]",
+				ArgsUsage: "<root>",
 				Action:    traverseState,
-				Flags:     slices.Concat(utils.NetworkFlags, utils.DatabaseFlags),
+				Flags:     slices.Concat(utils.TraverseStateFlags, utils.NetworkFlags, utils.DatabaseFlags),
 				Description: `
-geth snapshot traverse-state <state-root> [accountHash|accountAddress]
+geth snapshot traverse-state [--account <account>] [--start <key>] [--limit <key>] <state-root>
-will traverse the whole state from the given state root and will abort if any
-referenced trie node or contract code is missing. This command can be used for
-state integrity verification. The default checking target is the HEAD state.
 
-If accountHash or accountAddress is provided, traversal will start from that specific account and continue through all subsequent accounts.
+1. Traverse the whole state from the given state root:
-The format is auto-detected: 40/42 chars for address, 64/66 chars for hash.
+- --start: starting account key (64/66 chars hex) [optional]
+- --limit: ending account key (64/66 chars hex) [optional]
+
+2. Traverse a specific account's storage:
+- --account: account address (40/42 chars) or hash (64/66 chars) [required]
+- --start: starting storage key (64/66 chars hex) [optional]
+- --limit: ending storage key (64/66 chars hex) [optional]
+
+The default checking state root is the HEAD state if not specified.
+The command will abort if any referenced trie node or contract code is missing.
+This can be used for state integrity verification. The default target is HEAD state.
 
 It's also usable without snapshot enabled.
 `,
@@ -121,18 +128,26 @@ It's also usable without snapshot enabled.
 			{
 				Name:      "traverse-rawstate",
 				Usage:     "Traverse the state with given root hash and perform detailed verification",
-				ArgsUsage: "<root> [accountHash|accountAddress]",
+				ArgsUsage: "<root>",
 				Action:    traverseRawState,
-				Flags:     slices.Concat(utils.NetworkFlags, utils.DatabaseFlags),
+				Flags:     slices.Concat(utils.TraverseStateFlags, utils.NetworkFlags, utils.DatabaseFlags),
 				Description: `
-geth snapshot traverse-rawstate <state-root>
+geth snapshot traverse-rawstate [--account <account>] [--start <key>] [--limit <key>] <state-root>
-will traverse the whole state from the given root and will abort if any referenced
-trie node or contract code is missing. This command can be used for state integrity
-verification. The default checking target is the HEAD state. It's basically identical
-to traverse-state, but the check granularity is smaller.
 
-If accountHash or accountAddress is provided, traversal will start from that specific account and continue through all subsequent accounts.
+Similar to traverse-state but with more detailed verification at the trie node level.
-The format is auto-detected: 40/42 chars for address, 64/66 chars for hash.
+
+1. Traverse the whole state from the given state root:
+- --start: starting account key (64/66 chars hex) [optional]
+- --limit: ending account key (64/66 chars hex) [optional]
+
+2. Traverse a specific account's storage:
+- --account: account address (40/42 chars) or hash (64/66 chars) [required]
+- --start: starting storage key (64/66 chars hex) [optional]
+- --limit: ending storage key (64/66 chars hex) [optional]
+
+The default checking state root is the HEAD state if not specified.
+The command will abort if any referenced trie node or contract code is missing.
+This can be used for state integrity verification. The default target is HEAD state.
 
 It's also usable without snapshot enabled.
 `,
@@ -297,35 +312,92 @@ func traverseState(ctx *cli.Context) error {
 		return errors.New("no head block")
 	}
 
-	root, startKey, err := parseTraverseArgs(ctx)
+	config, err := parseTraverseArgs(ctx)
 	if err != nil {
 		return err
 	}
-	if root == (common.Hash{}) {
+	if config.root == (common.Hash{}) {
-		root = headBlock.Root()
+		config.root = headBlock.Root()
 	}
 
-	log.Info("Start traversing the state", "root", root.Hex(), "startKey", common.Bytes2Hex(startKey))
+	t, err := trie.NewStateTrie(trie.StateTrieID(config.root), triedb)
-	t, err := trie.NewStateTrie(trie.StateTrieID(root), triedb)
 	if err != nil {
-		log.Error("Failed to open trie", "root", root, "err", err)
+		log.Error("Failed to open trie", "root", config.root, "err", err)
 		return err
 	}
+
 	var (
 		accounts int
 		slots    int
 		codes    int
-		lastReport time.Time
 		start    = time.Now()
 	)
 
-	acctIt, err := t.NodeIterator(startKey)
+	go func() {
+		timer := time.NewTicker(time.Second * 8)
+		defer timer.Stop()
+		for range timer.C {
+			log.Info("Traversing state", "accounts", accounts, "slots", slots, "codes", codes, "elapsed", common.PrettyDuration(time.Since(start)))
+		}
+	}()
+
+	if config.isAccount {
+		log.Info("Start traversing storage trie", "root", config.root.Hex(), "account", config.account.Hex(), "startKey", common.Bytes2Hex(config.startKey), "limitKey", common.Bytes2Hex(config.limitKey))
+
+		acc, err := t.GetAccountByHash(config.account)
 		if err != nil {
-		log.Error("Failed to open iterator", "root", root, "err", err)
+			log.Error("Get account failed", "account", config.account.Hex(), "err", err)
+			return err
+		}
+
+		if acc.Root == types.EmptyRootHash {
+			log.Info("Account has no storage")
+			return nil
+		}
+
+		id := trie.StorageTrieID(config.root, config.account, acc.Root)
+		storageTrie, err := trie.NewStateTrie(id, triedb)
+		if err != nil {
+			log.Error("Failed to open storage trie", "root", acc.Root, "err", err)
+			return err
+		}
+
+		storageIt, err := storageTrie.NodeIterator(config.startKey)
+		if err != nil {
+			log.Error("Failed to open storage iterator", "root", acc.Root, "err", err)
+			return err
+		}
+
+		storageIter := trie.NewIterator(storageIt)
+		for storageIter.Next() {
+			if config.limitKey != nil && bytes.Compare(storageIter.Key, config.limitKey) >= 0 {
+				break
+			}
+
+			slots += 1
+			log.Debug("Storage slot", "key", common.Bytes2Hex(storageIter.Key), "value", common.Bytes2Hex(storageIter.Value))
+		}
+		if storageIter.Err != nil {
+			log.Error("Failed to traverse storage trie", "root", acc.Root, "err", storageIter.Err)
+			return storageIter.Err
+		}
+
+		log.Info("Storage traversal complete", "slots", slots, "elapsed", common.PrettyDuration(time.Since(start)))
+		return nil
+	} else {
+		log.Info("Start traversing state trie", "root", config.root.Hex(), "startKey", common.Bytes2Hex(config.startKey), "limitKey", common.Bytes2Hex(config.limitKey))
+
+		acctIt, err := t.NodeIterator(config.startKey)
+		if err != nil {
+			log.Error("Failed to open iterator", "root", config.root, "err", err)
 			return err
 		}
 		accIter := trie.NewIterator(acctIt)
 		for accIter.Next() {
+			if config.limitKey != nil && bytes.Compare(accIter.Key, config.limitKey) >= 0 {
+				break
+			}
+
 			accounts += 1
 			var acc types.StateAccount
 			if err := rlp.DecodeBytes(accIter.Value, &acc); err != nil {
@@ -333,7 +405,7 @@ func traverseState(ctx *cli.Context) error {
 				return err
 			}
 			if acc.Root != types.EmptyRootHash {
-			id := trie.StorageTrieID(root, common.BytesToHash(accIter.Key), acc.Root)
+				id := trie.StorageTrieID(config.root, common.BytesToHash(accIter.Key), acc.Root)
 				storageTrie, err := trie.NewStateTrie(id, triedb)
 				if err != nil {
 					log.Error("Failed to open storage trie", "root", acc.Root, "err", err)
@@ -347,11 +419,6 @@ func traverseState(ctx *cli.Context) error {
 				storageIter := trie.NewIterator(storageIt)
 				for storageIter.Next() {
 					slots += 1
-
-				if time.Since(lastReport) > time.Second*8 {
-					log.Info("Traversing state", "accounts", accounts, "slots", slots, "codes", codes, "elapsed", common.PrettyDuration(time.Since(start)))
-					lastReport = time.Now()
-				}
 				}
 				if storageIter.Err != nil {
 					log.Error("Failed to traverse storage trie", "root", acc.Root, "err", storageIter.Err)
@@ -365,45 +432,68 @@ func traverseState(ctx *cli.Context) error {
 				}
 				codes += 1
 			}
-		if time.Since(lastReport) > time.Second*8 {
-			log.Info("Traversing state", "accounts", accounts, "slots", slots, "codes", codes, "elapsed", common.PrettyDuration(time.Since(start)))
-			lastReport = time.Now()
-		}
 		}
 		if accIter.Err != nil {
-		log.Error("Failed to traverse state trie", "root", root, "err", accIter.Err)
+			log.Error("Failed to traverse state trie", "root", config.root, "err", accIter.Err)
 			return accIter.Err
 		}
-	log.Info("State is complete", "accounts", accounts, "slots", slots, "codes", codes, "elapsed", common.PrettyDuration(time.Since(start)))
+		log.Info("State traversal complete", "accounts", accounts, "slots", slots, "codes", codes, "elapsed", common.PrettyDuration(time.Since(start)))
 		return nil
+	}
 }
 
-func parseTraverseArgs(ctx *cli.Context) (root common.Hash, startKey []byte, err error) {
+type traverseConfig struct {
-	if ctx.NArg() > 2 {
+	root      common.Hash
-		err = errors.New("too many arguments")
+	startKey  []byte
-		return
+	limitKey  []byte
+	account   common.Hash
+	isAccount bool
+}
+
+func parseTraverseArgs(ctx *cli.Context) (*traverseConfig, error) {
+	if ctx.NArg() > 1 {
+		return nil, errors.New("too many arguments, only <root> is required")
 	}
 
-	if ctx.NArg() >= 1 {
+	config := &traverseConfig{}
-		root, err = parseRoot(ctx.Args().First())
+	var err error
+
+	if ctx.NArg() == 1 {
+		config.root, err = parseRoot(ctx.Args().First())
 		if err != nil {
-			return
+			return nil, err
 		}
 	}
 
-	if ctx.NArg() == 2 {
+	if accountFlag := ctx.String("account"); accountFlag != "" {
-		arg := ctx.Args().Get(1)
+		config.isAccount = true
-		switch len(arg) {
+		switch len(accountFlag) {
 		case 40, 42:
-			startKey = crypto.Keccak256Hash(common.HexToAddress(arg).Bytes()).Bytes()
+			config.account = crypto.Keccak256Hash(common.HexToAddress(accountFlag).Bytes())
 		case 64, 66:
-			startKey = common.HexToHash(arg).Bytes()
+			config.account = common.HexToHash(accountFlag)
 		default:
-			err = errors.New("invalid account format: must be 40/42 chars for address or 64/66 chars for hash")
+			return nil, errors.New("account must be 40/42 chars for address or 64/66 chars for hash")
-			return
 		}
 	}
-	return root, startKey, nil
+
+	if startFlag := ctx.String("start"); startFlag != "" {
+		if len(startFlag) == 64 || len(startFlag) == 66 {
+			config.startKey = common.HexToHash(startFlag).Bytes()
+		} else {
+			return nil, errors.New("start key must be 64/66 chars hex")
+		}
+	}
+
+	if limitFlag := ctx.String("limit"); limitFlag != "" {
+		if len(limitFlag) == 64 || len(limitFlag) == 66 {
+			config.limitKey = common.HexToHash(limitFlag).Bytes()
+		} else {
+			return nil, errors.New("limit key must be 64/66 chars hex")
+		}
+	}
+
+	return config, nil
 }
 
 // traverseRawState is a helper function used for pruning verification.
@@ -426,36 +516,122 @@ func traverseRawState(ctx *cli.Context) error {
 		return errors.New("no head block")
 	}
 
-	root, startKey, err := parseTraverseArgs(ctx)
+	config, err := parseTraverseArgs(ctx)
 	if err != nil {
 		log.Error("Failed to parse arguments", "err", err)
 		return err
 	}
-
+	if config.root == (common.Hash{}) {
-	log.Info("Start traversing the state", "root", root.Hex(), "startKey", common.Bytes2Hex(startKey))
+		config.root = headBlock.Root()
-	t, err := trie.NewStateTrie(trie.StateTrieID(root), triedb)
+	}
+	t, err := trie.NewStateTrie(trie.StateTrieID(config.root), triedb)
 	if err != nil {
-		log.Error("Failed to open trie", "root", root, "err", err)
+		log.Error("Failed to open trie", "root", config.root, "err", err)
 		return err
 	}
+
 	var (
-		nodes      int
 		accounts int
+		nodes    int
 		slots    int
 		codes    int
-		lastReport time.Time
 		start    = time.Now()
 		hasher   = crypto.NewKeccakState()
 		got      = make([]byte, 32)
 	)
-	accIter, err := t.NodeIterator(startKey)
+
+	go func() {
+		timer := time.NewTicker(time.Second * 8)
+		defer timer.Stop()
+		for range timer.C {
+			log.Info("Traversing rawstate", "nodes", nodes, "accounts", accounts, "slots", slots, "codes", codes, "elapsed", common.PrettyDuration(time.Since(start)))
+		}
+	}()
+
+	if config.isAccount {
+		log.Info("Start traversing storage trie (raw)", "root", config.root.Hex(), "account", config.account.Hex(), "startKey", common.Bytes2Hex(config.startKey), "limitKey", common.Bytes2Hex(config.limitKey))
+
+		acc, err := t.GetAccountByHash(config.account)
 		if err != nil {
-		log.Error("Failed to open iterator", "root", root, "err", err)
+			log.Error("Get account failed", "account", config.account.Hex(), "err", err)
 			return err
 		}
-	reader, err := triedb.NodeReader(root)
+
+		if acc.Root == types.EmptyRootHash {
+			log.Info("Account has no storage")
+			return nil
+		}
+
+		// Traverse the storage trie with detailed verification
+		id := trie.StorageTrieID(config.root, config.account, acc.Root)
+		storageTrie, err := trie.NewStateTrie(id, triedb)
 		if err != nil {
-		log.Error("State is non-existent", "root", root)
+			log.Error("Failed to open storage trie", "root", acc.Root, "err", err)
+			return err
+		}
+
+		storageIter, err := storageTrie.NodeIterator(config.startKey)
+		if err != nil {
+			log.Error("Failed to open storage iterator", "root", acc.Root, "err", err)
+			return err
+		}
+
+		reader, err := triedb.NodeReader(config.root)
+		if err != nil {
+			log.Error("State is non-existent", "root", config.root)
+			return nil
+		}
+
+		for storageIter.Next(true) {
+			nodes += 1
+			node := storageIter.Hash()
+
+			// Check the presence for non-empty hash node(embedded node doesn't
+			// have their own hash).
+			if node != (common.Hash{}) {
+				blob, _ := reader.Node(config.account, storageIter.Path(), node)
+				if len(blob) == 0 {
+					log.Error("Missing trie node(storage)", "hash", node)
+					return errors.New("missing storage")
+				}
+				hasher.Reset()
+				hasher.Write(blob)
+				hasher.Read(got)
+				if !bytes.Equal(got, node.Bytes()) {
+					log.Error("Invalid trie node(storage)", "hash", node.Hex(), "value", blob)
+					return errors.New("invalid storage node")
+				}
+			}
+
+			// Bump the counter if it's leaf node.
+			if storageIter.Leaf() {
+				// Check if we've exceeded the limit key for storage
+				if config.limitKey != nil && bytes.Compare(storageIter.LeafKey(), config.limitKey) >= 0 {
+					break
+				}
+
+				slots += 1
+				log.Debug("Storage slot", "key", common.Bytes2Hex(storageIter.LeafKey()), "value", common.Bytes2Hex(storageIter.LeafBlob()))
+			}
+		}
+		if storageIter.Error() != nil {
+			log.Error("Failed to traverse storage trie", "root", acc.Root, "err", storageIter.Error())
+			return storageIter.Error()
+		}
+
+		log.Info("Storage traversal complete (raw)", "nodes", nodes, "slots", slots, "elapsed", common.PrettyDuration(time.Since(start)))
+		return nil
+	} else {
+		log.Info("Start traversing the state trie (raw)", "root", config.root.Hex(), "startKey", common.Bytes2Hex(config.startKey), "limitKey", common.Bytes2Hex(config.limitKey))
+
+		accIter, err := t.NodeIterator(config.startKey)
+		if err != nil {
+			log.Error("Failed to open iterator", "root", config.root, "err", err)
+			return err
+		}
+		reader, err := triedb.NodeReader(config.root)
+		if err != nil {
+			log.Error("State is non-existent", "root", config.root)
 			return nil
 		}
 		for accIter.Next(true) {
@@ -481,6 +657,11 @@ func traverseRawState(ctx *cli.Context) error {
 			// If it's a leaf node, yes we are touching an account,
 			// dig into the storage trie further.
 			if accIter.Leaf() {
+				// Check if we've exceeded the limit key for accounts
+				if config.limitKey != nil && bytes.Compare(accIter.LeafKey(), config.limitKey) >= 0 {
+					break
+				}
+
 				accounts += 1
 				var acc types.StateAccount
 				if err := rlp.DecodeBytes(accIter.LeafBlob(), &acc); err != nil {
@@ -488,7 +669,7 @@ func traverseRawState(ctx *cli.Context) error {
 					return errors.New("invalid account")
 				}
 				if acc.Root != types.EmptyRootHash {
-				id := trie.StorageTrieID(root, common.BytesToHash(accIter.LeafKey()), acc.Root)
+					id := trie.StorageTrieID(config.root, common.BytesToHash(accIter.LeafKey()), acc.Root)
 					storageTrie, err := trie.NewStateTrie(id, triedb)
 					if err != nil {
 						log.Error("Failed to open storage trie", "root", acc.Root, "err", err)
@@ -523,10 +704,6 @@ func traverseRawState(ctx *cli.Context) error {
 						if storageIter.Leaf() {
 							slots += 1
 						}
-					if time.Since(lastReport) > time.Second*8 {
-						log.Info("Traversing state", "nodes", nodes, "accounts", accounts, "slots", slots, "codes", codes, "elapsed", common.PrettyDuration(time.Since(start)))
-						lastReport = time.Now()
-					}
 					}
 					if storageIter.Error() != nil {
 						log.Error("Failed to traverse storage trie", "root", acc.Root, "err", storageIter.Error())
@@ -540,18 +717,15 @@ func traverseRawState(ctx *cli.Context) error {
 					}
 					codes += 1
 				}
-			if time.Since(lastReport) > time.Second*8 {
-				log.Info("Traversing state", "nodes", nodes, "accounts", accounts, "slots", slots, "codes", codes, "elapsed", common.PrettyDuration(time.Since(start)))
-				lastReport = time.Now()
-			}
 			}
 		}
 		if accIter.Error() != nil {
-		log.Error("Failed to traverse state trie", "root", root, "err", accIter.Error())
+			log.Error("Failed to traverse state trie", "root", config.root, "err", accIter.Error())
 			return accIter.Error()
 		}
-	log.Info("State is complete", "nodes", nodes, "accounts", accounts, "slots", slots, "codes", codes, "elapsed", common.PrettyDuration(time.Since(start)))
+		log.Info("State traversal complete (raw)", "nodes", nodes, "accounts", accounts, "slots", slots, "codes", codes, "elapsed", common.PrettyDuration(time.Since(start)))
 		return nil
+	}
 }
 
 func parseRoot(input string) (common.Hash, error) {

cmd/utils/flags.go

@@ -994,6 +994,22 @@ var (
 		StateSchemeFlag,
 		HttpHeaderFlag,
 	}
+
+	// TraverseStateFlags is the flag group of all state traversal flags.
+	TraverseStateFlags = []cli.Flag{
+		&cli.StringFlag{
+			Name:  "account",
+			Usage: "Account address or hash to traverse storage for (enables account mode)",
+		},
+		&cli.StringFlag{
+			Name:  "start",
+			Usage: "Starting key (account key for state mode, storage key for account mode)",
+		},
+		&cli.StringFlag{
+			Name:  "limit",
+			Usage: "Ending key (account key for state mode, storage key for account mode)",
+		},
+	}
 )
 
 // default account to prefund when running Geth in dev mode