feat: query witnesses over wit/1 protocol (#634)

* feat: query witnesses

* fix: update request id

* fix: disable tx fetching

* fix: print enode

* fix: return when EOF error

* fix: lint

* feat: use witness pagination

* feat: request witness pages

* fix: requesting multiple witness pages

* fix: logging

* fix: logging

* fix: undo comments

* feat: ping private key flags

* fix: flag set

* feat: add DialOpts

* fix: remove duplicate listen flag

* fix: remove unnecessary int cast

* fix: listen flag desc

* docs: make gen-doc

* fix: remove unused log

* fix: remove log import

* feat: add key file flags to p2p query

* fix: mark query key flags mutex

* fix: p2p query conflicting `a` shorthand

* docs: make gen-doc

* fix: clean up logs

* feat: handle enabling caps

---------

Co-authored-by: Léo Vincent <[email protected]>
Co-authored-by: Claude <[email protected]>

Author: 3 people

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

@@ -1,6 +1,8 @@
 package ping
 
 import (
+	"crypto/ecdsa"
+	"net"
 	"sync"
 	"time"
 
@@ -17,6 +19,13 @@ type (
 		OutputFile string
 		NodesFile  string
 		Listen     bool
+		Port       int
+		Addr       net.IP
+		KeyFile    string
+		PrivateKey string
+		EnableWit  bool
+
+		privateKey *ecdsa.PrivateKey
 	}
 )
 
@@ -34,6 +43,14 @@ Status messages and output JSON. If providing a enode/enr rather than a nodes
 file, then the connection will remain open by default (--listen=true), and you
 can see other messages the peer sends (e.g. blocks, transactions, etc.).`,
 	Args: cobra.MinimumNArgs(1),
+	PreRunE: func(cmd *cobra.Command, args []string) (err error) {
+		inputPingParams.privateKey, err = p2p.ParsePrivateKey(inputPingParams.KeyFile, inputPingParams.PrivateKey)
+		if err != nil {
+			return err
+		}
+
+		return nil
+	},
 	RunE: func(cmd *cobra.Command, args []string) error {
 		nodes := []*enode.Node{}
 		if input, err := p2p.ReadNodeSet(args[0]); err == nil {
@@ -81,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:       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 {
@@ -133,4 +157,10 @@ func init() {
 	f.BoolVarP(&inputPingParams.Listen, "listen", "l", true,
 		`keep connection open and listen to peer. This only works if first
 argument is an enode/enr, not a nodes file`)
+	f.BoolVarP(&inputPingParams.EnableWit, "wit", "w", false, "enable wit/1 capability")
+	f.IntVarP(&inputPingParams.Port, "port", "P", 30303, "port for discovery protocol")
+	f.IPVarP(&inputPingParams.Addr, "addr", "a", net.ParseIP("127.0.0.1"), "address to bind discovery listener")
+	f.StringVarP(&inputPingParams.KeyFile, "key-file", "k", "", "private key file (cannot be set with --key)")
+	f.StringVar(&inputPingParams.PrivateKey, "key", "", "hex-encoded private key (cannot be set with --key-file)")
+	PingCmd.MarkFlagsMutuallyExclusive("key-file", "key")
 }

cmd/p2p/query/query.go

@@ -1,7 +1,9 @@
 package query
 
 import (
+	"crypto/ecdsa"
 	"fmt"
+	"net"
 
 	"github.com/ethereum/go-ethereum/eth/protocols/eth"
 	"github.com/rs/zerolog/log"
@@ -14,6 +16,12 @@ type (
 	queryParams struct {
 		StartBlock uint64
 		Amount     uint64
+		Port       int
+		Addr       net.IP
+		KeyFile    string
+		PrivateKey string
+
+		privateKey *ecdsa.PrivateKey
 	}
 )
 
@@ -30,10 +38,16 @@ This command will initially establish a handshake and exchange status message
 from the peer. Then, it will query the node for block(s) given the start block
 and the amount of blocks to query and print the results.`,
 	Args: cobra.MinimumNArgs(1),
-	PreRunE: func(cmd *cobra.Command, args []string) error {
+	PreRunE: func(cmd *cobra.Command, args []string) (err error) {
 		if inputQueryParams.Amount < 1 {
 			return fmt.Errorf("amount must be greater than 0")
 		}
+
+		inputQueryParams.privateKey, err = p2p.ParsePrivateKey(inputQueryParams.KeyFile, inputQueryParams.PrivateKey)
+		if err != nil {
+			return err
+		}
+
 		return nil
 	},
 	Run: func(cmd *cobra.Command, args []string) {
@@ -50,7 +64,13 @@ and the amount of blocks to query and print the results.`,
 			amount uint64 = inputQueryParams.Amount
 		)
 
-		conn, err := p2p.Dial(node)
+		opts := p2p.DialOpts{
+			Port:       inputQueryParams.Port,
+			Addr:       inputQueryParams.Addr,
+			PrivateKey: inputQueryParams.privateKey,
+		}
+
+		conn, err := p2p.Dial(node, opts)
 		if err != nil {
 			log.Error().Err(err).Msg("Dial failed")
 			return
@@ -104,8 +124,14 @@ func print(headers eth.BlockHeadersRequest) {
 }
 
 func init() {
-	QueryCmd.Flags().Uint64VarP(&inputQueryParams.StartBlock, "start-block", "s", 0, "block number to start querying from")
-	QueryCmd.Flags().Uint64VarP(&inputQueryParams.Amount, "amount", "a", 1, "amount of blocks to query")
+	f := QueryCmd.Flags()
+	f.Uint64VarP(&inputQueryParams.StartBlock, "start-block", "s", 0, "block number to start querying from")
+	f.Uint64VarP(&inputQueryParams.Amount, "amount", "a", 1, "amount of blocks to query")
+	f.IntVarP(&inputQueryParams.Port, "port", "P", 30303, "port for discovery protocol")
+	f.IPVar(&inputQueryParams.Addr, "addr", net.ParseIP("127.0.0.1"), "address to bind discovery listener")
+	f.StringVarP(&inputQueryParams.KeyFile, "key-file", "k", "", "private key file (cannot be set with --key)")
+	f.StringVar(&inputQueryParams.PrivateKey, "key", "", "hex-encoded private key (cannot be set with --key-file)")
+	QueryCmd.MarkFlagsMutuallyExclusive("key-file", "key")
 	if err := QueryCmd.MarkFlagRequired("start-block"); err != nil {
 		log.Error().Err(err).Msg("Failed to mark start-block as required flag")
 	}

doc/polycli_p2p_ping.md

@@ -28,11 +28,16 @@ can see other messages the peer sends (e.g. blocks, transactions, etc.).
 ## Flags
 
 ```bash
-  -h, --help            help for ping
-  -l, --listen          keep connection open and listen to peer. This only works if first
-                        argument is an enode/enr, not a nodes file (default true)
-  -o, --output string   write ping results to output file (default stdout)
-  -p, --parallel int    how many parallel pings to attempt (default 16)
+  -a, --addr ip           address to bind discovery listener (default 127.0.0.1)
+  -h, --help              help for ping
+      --key string        hex-encoded private key (cannot be set with --key-file)
+  -k, --key-file string   private key file (cannot be set with --key)
+  -l, --listen            keep connection open and listen to peer. This only works if first
+                          argument is an enode/enr, not a nodes file (default true)
+  -o, --output string     write ping results to output file (default stdout)
+  -p, --parallel int      how many parallel pings to attempt (default 16)
+  -P, --port int          port for discovery protocol (default 30303)
+  -w, --wit               enable wit/1 capability
 ```
 
 The command also inherits flags from parent commands.

doc/polycli_p2p_query.md

@@ -27,8 +27,12 @@ and the amount of blocks to query and print the results.
 ## Flags
 
 ```bash
+      --addr ip            address to bind discovery listener (default 127.0.0.1)
   -a, --amount uint        amount of blocks to query (default 1)
   -h, --help               help for query
+      --key string         hex-encoded private key (cannot be set with --key-file)
+  -k, --key-file string    private key file (cannot be set with --key)
+  -P, --port int           port for discovery protocol (default 30303)
   -s, --start-block uint   block number to start querying from
 ```
 

p2p/log.go

@@ -21,6 +21,10 @@ type MessageCount struct {
 	Pings               int64 `json:"pings,omitempty"`
 	Errors              int64 `json:"errors,omitempty"`
 	Disconnects         int64 `json:"disconnects,omitempty"`
+	NewWitness          int64 `json:"new_witness,omitempty"`
+	NewWitnessHashes    int64 `json:"new_witness_hashes,omitempty"`
+	GetWitnessRequest   int64 `json:"get_witness_request,omitempty"`
+	Witness             int64 `json:"witness,omitempty"`
 }
 
 // Load takes a snapshot of all the counts in a thread-safe manner. Make sure
@@ -39,6 +43,10 @@ func (count *MessageCount) Load() MessageCount {
 		Pings:               atomic.LoadInt64(&count.Pings),
 		Errors:              atomic.LoadInt64(&count.Errors),
 		Disconnects:         atomic.LoadInt64(&count.Disconnects),
+		NewWitness:          atomic.LoadInt64(&count.NewWitness),
+		NewWitnessHashes:    atomic.LoadInt64(&count.NewWitnessHashes),
+		GetWitnessRequest:   atomic.LoadInt64(&count.GetWitnessRequest),
+		Witness:             atomic.LoadInt64(&count.Witness),
 	}
 }
 
@@ -56,6 +64,10 @@ func (count *MessageCount) Clear() {
 	atomic.StoreInt64(&count.Pings, 0)
 	atomic.StoreInt64(&count.Errors, 0)
 	atomic.StoreInt64(&count.Disconnects, 0)
+	atomic.StoreInt64(&count.NewWitness, 0)
+	atomic.StoreInt64(&count.NewWitnessHashes, 0)
+	atomic.StoreInt64(&count.GetWitnessRequest, 0)
+	atomic.StoreInt64(&count.Witness, 0)
 }
 
 // IsEmpty checks whether the sum of all the counts is empty. Make sure to call

p2p/p2p.go

@@ -2,15 +2,18 @@ package p2p
 
 import (
 	"bytes"
+	"crypto/ecdsa"
 	"encoding/base64"
 	"encoding/hex"
 	"fmt"
 	"net"
 	"strings"
 
+	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/p2p/enode"
 	"github.com/ethereum/go-ethereum/p2p/enr"
 	"github.com/ethereum/go-ethereum/rlp"
+	"github.com/rs/zerolog/log"
 )
 
 func Listen(ln *enode.LocalNode) (*net.UDPConn, error) {
@@ -101,3 +104,43 @@ func ParseBootnodes(bootnodes string) ([]*enode.Node, error) {
 
 	return nodes, nil
 }
+
+// ParsePrivateKey loads a private key from a file path or hex string.
+// If file is provided, it attempts to load from the file path first.
+// If the file doesn't exist, it generates a new key and saves it to the file.
+// If key is provided instead, it parses the hex-encoded private key string.
+// If neither file nor key is provided, it generates and returns a new private key.
+func ParsePrivateKey(file string, key string) (*ecdsa.PrivateKey, error) {
+	if len(file) > 0 {
+		privateKey, err := crypto.LoadECDSA(file)
+		if err == nil {
+			return privateKey, nil
+		}
+
+		log.Warn().Err(err).Msg("Key file was not found, generating a new key file")
+
+		privateKey, err = crypto.GenerateKey()
+		if err != nil {
+			log.Error().Err(err).Msg("Failed to generate new private key")
+			return nil, err
+		}
+
+		if err := crypto.SaveECDSA(file, privateKey); err != nil {
+			log.Error().Err(err).Msg("Failed to save private key to file")
+			return nil, err
+		}
+
+		return privateKey, nil
+	}
+
+	if len(key) > 0 {
+		privateKey, err := crypto.HexToECDSA(key)
+		if err != nil {
+			log.Error().Err(err).Msg("Failed to parse private key")
+			return nil, err
+		}
+		return privateKey, nil
+	}
+
+	return crypto.GenerateKey()
+}