feat: add DialOpts

Author: minhd-vu

cmd/p2p/crawl/crawl_util.go

@@ -181,7 +181,7 @@ func (c *crawler) updateNode(n *enode.Node) int {
 		c.mu.Unlock()
 	}()
 
-	conn, dialErr := p2p.Dial(n)
+	conn, dialErr := p2p.Dial(n, p2p.NewDialOpts())
 	if err = dialErr; err != nil {
 		log.Error().Err(err).Msg("Dial failed")
 		return nodeDialErr

cmd/p2p/ping/ping.go

@@ -2,6 +2,7 @@ package ping
 
 import (
 	"crypto/ecdsa"
+	"net"
 	"sync"
 	"time"
 
@@ -18,10 +19,11 @@ type (
 		OutputFile string
 		NodesFile  string
 		Listen     bool
-		Port       uint
-		Addr       bool
+		Port       int
+		Addr       net.IP
 		KeyFile    string
 		PrivateKey string
+		EnableWit  bool
 
 		privateKey *ecdsa.PrivateKey
 	}
@@ -96,7 +98,14 @@ can see other messages the peer sends (e.g. blocks, transactions, etc.).`,
 					status *p2p.Status
 				)
 
-				conn, err := p2p.Dial(node)
+				opts := p2p.DialOpts{
+					EnableWit:  inputPingParams.EnableWit,
+					Port:       int(inputPingParams.Port),
+					Addr:       inputPingParams.Addr,
+					PrivateKey: inputPingParams.privateKey,
+				}
+
+				conn, err := p2p.Dial(node, opts)
 				if err != nil {
 					log.Error().Err(err).Msg("Dial failed")
 				} else {
@@ -150,7 +159,9 @@ argument is an enode/enr, not a nodes file.`)
 	PingCmd.Flags().BoolVarP(&inputPingParams.Listen, "listen", "l", true,
 		`Keep the connection open and listen to the peer. This only works if the first
 argument is an enode/enr, not a nodes file.`)
-	PingCmd.Flags().UintVarP(&inputPingParams.Port, "port", "P", 30303, "The port to listen on")
+	PingCmd.Flags().BoolVarP(&inputPingParams.EnableWit, "wit", "w", false, "Whether to enable the wit/1 capability")
+	PingCmd.Flags().IntVarP(&inputPingParams.Port, "port", "P", 30303, "Port for discovery protocol")
+	PingCmd.Flags().IPVarP(&inputPingParams.Addr, "addr", "a", net.ParseIP("127.0.0.1"), "Address to bind discovery listener")
 	PingCmd.Flags().StringVarP(&inputPingParams.KeyFile, "key-file", "k", "", "Private key file (cannot be set with --key)")
 	PingCmd.Flags().StringVar(&inputPingParams.PrivateKey, "key", "", "Hex-encoded private key (cannot be set with --key-file)")
 	PingCmd.MarkFlagsMutuallyExclusive("key-file", "key")

cmd/p2p/query/query.go

@@ -50,7 +50,7 @@ and the amount of blocks to query and print the results.`,
 			amount uint64 = inputQueryParams.Amount
 		)
 
-		conn, err := p2p.Dial(node)
+		conn, err := p2p.Dial(node, p2p.NewDialOpts())
 		if err != nil {
 			log.Error().Err(err).Msg("Dial failed")
 			return

p2p/rlpx.go

@@ -1,6 +1,7 @@
 package p2p
 
 import (
+	"crypto/ecdsa"
 	"fmt"
 	"math/rand"
 	"net"
@@ -21,39 +22,54 @@ var (
 	timeout = 20 * time.Second
 )
 
+type DialOpts struct {
+	EnableWit  bool
+	Port       int
+	Addr       net.IP
+	PrivateKey *ecdsa.PrivateKey
+}
+
+func NewDialOpts() DialOpts {
+	privateKey, err := crypto.GenerateKey()
+	if err != nil {
+		log.Error().Err(err).Msg("Failed to generate private key")
+	}
+
+	return DialOpts{
+		EnableWit:  false,
+		Port:       30303,
+		Addr:       net.ParseIP("127.0.0.1"),
+		PrivateKey: privateKey,
+	}
+}
+
 // Dial attempts to Dial the given node and perform a handshake,
 // returning the created Conn if successful.
-func Dial(n *enode.Node) (*rlpxConn, error) {
+func Dial(n *enode.Node, opts DialOpts) (*rlpxConn, error) {
 	fd, err := net.Dial("tcp", fmt.Sprintf("%v:%d", n.IP(), n.TCP()))
 	if err != nil {
 		return nil, err
 	}
 
-	conn := rlpxConn{
-		Conn:   rlpx.NewConn(fd, n.Pubkey()),
-		node:   n,
-		logger: log.With().Str("peer", n.URLv4()).Logger(),
-		caps: []p2p.Cap{
-			{Name: "eth", Version: 66},
-			{Name: "eth", Version: 67},
-			{Name: "eth", Version: 68},
-			{Name: "wit", Version: 1},
-		},
+	caps := []p2p.Cap{
+		{Name: "eth", Version: 66},
+		{Name: "eth", Version: 67},
+		{Name: "eth", Version: 68},
 	}
 
-	if conn.ourKey, err = crypto.LoadECDSA("witness.key"); err != nil {
-		log.Warn().Msg("witness.key not found generating key")
-		if conn.ourKey, err = crypto.GenerateKey(); err != nil {
-			return nil, err
-		}
+	if opts.EnableWit {
+		caps = append(caps, p2p.Cap{Name: "wit", Version: 1})
 	}
 
-	if err = crypto.SaveECDSA("witness.key", conn.ourKey); err != nil {
-		return nil, err
+	conn := rlpxConn{
+		Conn:   rlpx.NewConn(fd, n.Pubkey()),
+		node:   n,
+		logger: log.With().Str("peer", n.URLv4()).Logger(),
+		caps:   caps,
+		ourKey: opts.PrivateKey,
 	}
 
-	ip := net.ParseIP("127.0.0.1")
-	v4 := enode.NewV4(&conn.ourKey.PublicKey, ip, 30303, 30303)
+	v4 := enode.NewV4(&conn.ourKey.PublicKey, opts.Addr, opts.Port, opts.Port)
 
 	log.Info().Any("enode", v4.String()).Send()
 
@@ -230,6 +246,7 @@ func (c *rlpxConn) ReadAndServe(count *MessageCount) error {
 					}
 
 					req := GetWitnessPacket{
+						RequestId: rand.Uint64(),
 						GetWitnessRequest: &GetWitnessRequest{
 							WitnessPages: []WitnessPageRequest{
 								{
@@ -238,9 +255,7 @@ func (c *rlpxConn) ReadAndServe(count *MessageCount) error {
 								},
 							},
 						},
-						RequestId: uint64(time.Now().Unix()),
 					}
-					log.Trace().Any("request", req).Msg("Writing GetWitnessPacket request")
 					if err := c.Write(req); err != nil {
 						log.Error().Err(err).Msg("Failed to write GetWitnessPacket request")
 					}
@@ -250,6 +265,7 @@ func (c *rlpxConn) ReadAndServe(count *MessageCount) error {
 				c.logger.Trace().Str("hash", msg.Block.Hash().Hex()).Msg("Received NewBlock")
 
 				req := GetWitnessPacket{
+					RequestId: rand.Uint64(),
 					GetWitnessRequest: &GetWitnessRequest{
 						WitnessPages: []WitnessPageRequest{
 							{
@@ -258,9 +274,7 @@ func (c *rlpxConn) ReadAndServe(count *MessageCount) error {
 							},
 						},
 					},
-					RequestId: uint64(time.Now().Unix()),
 				}
-				log.Trace().Any("request", req).Msg("Writing GetWitnessPacket request")
 				if err := c.Write(req); err != nil {
 					log.Error().Err(err).Msg("Failed to write GetWitnessPacket request")
 				}