From 2c4e012660f34ccdf09b80c377ad6665ddaffd0f Mon Sep 17 00:00:00 2001
From: Felix Lange <fjl@twurst.com>
Date: Sun, 8 Mar 2020 15:10:29 +0100
Subject: [PATCH] cmd/devp2p: add discv5 commands

---
 cmd/devp2p/crawl.go     |   9 ++-
 cmd/devp2p/discv5cmd.go | 127 ++++++++++++++++++++++++++++++++++++++++
 cmd/devp2p/main.go      |   1 +
 3 files changed, 134 insertions(+), 3 deletions(-)
 create mode 100644 cmd/devp2p/discv5cmd.go

diff --git a/cmd/devp2p/crawl.go b/cmd/devp2p/crawl.go
index 92aaad72a3..cc68dbcbfc 100644
--- a/cmd/devp2p/crawl.go
+++ b/cmd/devp2p/crawl.go
@@ -20,14 +20,13 @@ import (
 	"time"
 
 	"github.com/ethereum/go-ethereum/log"
-	"github.com/ethereum/go-ethereum/p2p/discover"
 	"github.com/ethereum/go-ethereum/p2p/enode"
 )
 
 type crawler struct {
 	input     nodeSet
 	output    nodeSet
-	disc      *discover.UDPv4
+	disc      resolver
 	iters     []enode.Iterator
 	inputIter enode.Iterator
 	ch        chan *enode.Node
@@ -37,7 +36,11 @@ type crawler struct {
 	revalidateInterval time.Duration
 }
 
-func newCrawler(input nodeSet, disc *discover.UDPv4, iters ...enode.Iterator) *crawler {
+type resolver interface {
+	RequestENR(*enode.Node) (*enode.Node, error)
+}
+
+func newCrawler(input nodeSet, disc resolver, iters ...enode.Iterator) *crawler {
 	c := &crawler{
 		input:     input,
 		output:    make(nodeSet, len(input)),
diff --git a/cmd/devp2p/discv5cmd.go b/cmd/devp2p/discv5cmd.go
new file mode 100644
index 0000000000..049f576b40
--- /dev/null
+++ b/cmd/devp2p/discv5cmd.go
@@ -0,0 +1,127 @@
+// Copyright 2019 The go-ethereum Authors
+// This file is part of go-ethereum.
+//
+// go-ethereum is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// go-ethereum is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU General Public License for more details.
+//
+// You should have received a copy of the GNU General Public License
+// along with go-ethereum. If not, see <http://www.gnu.org/licenses/>.
+
+package main
+
+import (
+	"fmt"
+	"time"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/p2p/discover"
+	"gopkg.in/urfave/cli.v1"
+)
+
+var (
+	discv5Command = cli.Command{
+		Name:  "discv5",
+		Usage: "Node Discovery v5 tools",
+		Subcommands: []cli.Command{
+			discv5PingCommand,
+			discv5ResolveCommand,
+			discv5CrawlCommand,
+			discv5ListenCommand,
+		},
+	}
+	discv5PingCommand = cli.Command{
+		Name:   "ping",
+		Usage:  "Sends ping to a node",
+		Action: discv5Ping,
+	}
+	discv5ResolveCommand = cli.Command{
+		Name:   "resolve",
+		Usage:  "Finds a node in the DHT",
+		Action: discv5Resolve,
+		Flags:  []cli.Flag{bootnodesFlag},
+	}
+	discv5CrawlCommand = cli.Command{
+		Name:   "crawl",
+		Usage:  "Updates a nodes.json file with random nodes found in the DHT",
+		Action: discv5Crawl,
+		Flags:  []cli.Flag{bootnodesFlag, crawlTimeoutFlag},
+	}
+	discv5ListenCommand = cli.Command{
+		Name:   "listen",
+		Usage:  "Runs a node",
+		Action: discv5Listen,
+		Flags:  []cli.Flag{bootnodesFlag},
+	}
+)
+
+func discv5Ping(ctx *cli.Context) error {
+	n := getNodeArg(ctx)
+	disc := startV5(nil)
+	defer disc.Close()
+
+	fmt.Println(disc.Ping(n))
+	return nil
+}
+
+func discv5Resolve(ctx *cli.Context) error {
+	n := getNodeArg(ctx)
+	disc := startV5(nil)
+	defer disc.Close()
+
+	fmt.Println(disc.Resolve(n))
+	return nil
+}
+
+func discv5Crawl(ctx *cli.Context) error {
+	if ctx.NArg() < 1 {
+		return fmt.Errorf("need nodes file as argument")
+	}
+	nodesFile := ctx.Args().First()
+	var inputSet nodeSet
+	if common.FileExist(nodesFile) {
+		inputSet = loadNodesJSON(nodesFile)
+	}
+
+	disc := startV5(ctx)
+	defer disc.Close()
+	c := newCrawler(inputSet, disc, disc.RandomNodes())
+	c.revalidateInterval = 10 * time.Minute
+	output := c.run(ctx.Duration(crawlTimeoutFlag.Name))
+	writeNodesJSON(nodesFile, output)
+	return nil
+}
+
+func discv5Listen(ctx *cli.Context) error {
+	disc := startV5(nil)
+	defer disc.Close()
+
+	fmt.Println(disc.Self())
+	select {}
+}
+
+// startV5 starts an ephemeral discovery v5 node.
+func startV5(ctx *cli.Context) *discover.UDPv5 {
+	socket, ln, cfg, err := listen()
+	if err != nil {
+		exit(err)
+	}
+	if commandHasFlag(ctx, bootnodesFlag) {
+		bn, err := parseBootnodes(ctx)
+		if err != nil {
+			exit(err)
+		}
+		cfg.Bootnodes = bn
+	}
+	disc, err := discover.ListenV5(socket, ln, cfg)
+	if err != nil {
+		exit(err)
+	}
+	return disc
+}
diff --git a/cmd/devp2p/main.go b/cmd/devp2p/main.go
index b895941f25..19aec77ed4 100644
--- a/cmd/devp2p/main.go
+++ b/cmd/devp2p/main.go
@@ -59,6 +59,7 @@ func init() {
 	app.Commands = []cli.Command{
 		enrdumpCommand,
 		discv4Command,
+		discv5Command,
 		dnsCommand,
 		nodesetCommand,
 	}