Discovery: refactor public key to node ID conversions. (#3634)

Encode and hash logic was duplicated in multiple places. * Move encoding to p2p/discover/v4wire * Move hashing to p2p/enode/idscheme * Change newRandomLookup to create a proper random key on a curve.
erigontech · Mar 3, 2022 · 0242b5a · 0242b5a
1 parent c8a25d0
commit 0242b5a
Show file tree

Hide file tree

Showing 9 changed files with 55 additions and 87 deletions.
diff --git a/p2p/discover/node.go b/p2p/discover/node.go
@@ -17,15 +17,9 @@
 package discover
 
 import (
-	"crypto/ecdsa"
-	"crypto/elliptic"
-	"errors"
-	"math/big"
 	"net"
 	"time"
 
-	"github.com/ledgerwatch/erigon/common/math"
-	"github.com/ledgerwatch/erigon/crypto"
 	"github.com/ledgerwatch/erigon/p2p/enode"
 )
 
@@ -37,33 +31,6 @@ type node struct {
 	livenessChecks uint      // how often liveness was checked
 }
 
-type encPubkey [64]byte
-
-func encodePubkey(key *ecdsa.PublicKey) encPubkey {
-	var e encPubkey
-	math.ReadBits(key.X, e[:len(e)/2])
-	math.ReadBits(key.Y, e[len(e)/2:])
-	return e
-}
-
-func decodePubkey(curve elliptic.Curve, e []byte) (*ecdsa.PublicKey, error) {
-	if len(e) != len(encPubkey{}) {
-		return nil, errors.New("wrong size public key data")
-	}
-	p := &ecdsa.PublicKey{Curve: curve, X: new(big.Int), Y: new(big.Int)}
-	half := len(e) / 2
-	p.X.SetBytes(e[:half])
-	p.Y.SetBytes(e[half:])
-	if !p.Curve.IsOnCurve(p.X, p.Y) {
-		return nil, errors.New("invalid curve point")
-	}
-	return p, nil
-}
-
-func (e encPubkey) id() enode.ID {
-	return enode.ID(crypto.Keccak256Hash(e[:]))
-}
-
 func wrapNode(n *enode.Node) *node {
 	return &node{Node: *n}
 }

diff --git a/p2p/discover/table_util_test.go b/p2p/discover/table_util_test.go
@@ -244,7 +244,7 @@ func hexEncPrivkey(h string) *ecdsa.PrivateKey {
 }
 
 // hexEncPubkey decodes h as a public key.
-func hexEncPubkey(h string) (ret encPubkey) {
+func hexEncPubkey(h string) (ret enode.PubkeyEncoded) {
 	b, err := hex.DecodeString(h)
 	if err != nil {
 		panic(err)

diff --git a/p2p/discover/v4_lookup_test.go b/p2p/discover/v4_lookup_test.go
@@ -39,7 +39,7 @@ func TestUDPv4_Lookup(t *testing.T) {
 	test := newUDPTest(t)
 
 	// Lookup on empty table returns no nodes.
-	targetKey, _ := decodePubkey(crypto.S256(), lookupTestnet.target[:])
+	targetKey, _ := v4wire.DecodePubkey(crypto.S256(), v4wire.Pubkey(lookupTestnet.target))
 	if results := test.udp.LookupPubkey(targetKey); len(results) > 0 {
 		t.Fatalf("lookup on empty table returned %d results: %#v", len(results), results)
 	}
@@ -61,7 +61,7 @@ func TestUDPv4_Lookup(t *testing.T) {
 	results := <-resultC
 	t.Logf("results:")
 	for _, e := range results {
-		t.Logf("  ld=%d, %x", enode.LogDist(lookupTestnet.target.id(), e.ID()), e.ID().Bytes())
+		t.Logf("  ld=%d, %x", enode.LogDist(lookupTestnet.target.ID(), e.ID()), e.ID().Bytes())
 	}
 	if len(results) != bucketSize {
 		t.Errorf("wrong number of results: got %d, want %d", len(results), bucketSize)
@@ -150,7 +150,7 @@ func serveTestnet(test *udpTest, testnet *preminedTestnet) {
 			case *v4wire.Ping:
 				test.packetInFrom(nil, key, to, &v4wire.Pong{Expiration: futureExp, ReplyTok: hash})
 			case *v4wire.Findnode:
-				dist := enode.LogDist(n.ID(), testnet.target.id())
+				dist := enode.LogDist(n.ID(), testnet.target.ID())
 				nodes := testnet.nodesAtDistance(dist - 1)
 				test.packetInFrom(nil, key, to, &v4wire.Neighbors{Expiration: futureExp, Nodes: nodes})
 			}
@@ -164,12 +164,12 @@ func checkLookupResults(t *testing.T, tn *preminedTestnet, results []*enode.Node
 	t.Helper()
 	t.Logf("results:")
 	for _, e := range results {
-		t.Logf("  ld=%d, %x", enode.LogDist(tn.target.id(), e.ID()), e.ID().Bytes())
+		t.Logf("  ld=%d, %x", enode.LogDist(tn.target.ID(), e.ID()), e.ID().Bytes())
 	}
 	if hasDuplicates(wrapNodes(results)) {
 		t.Errorf("result set contains duplicate entries")
 	}
-	if !sortedByDistanceTo(tn.target.id(), wrapNodes(results)) {
+	if !sortedByDistanceTo(tn.target.ID(), wrapNodes(results)) {
 		t.Errorf("result set not sorted by distance to target")
 	}
 	wantNodes := tn.closest(len(results))
@@ -239,7 +239,7 @@ var lookupTestnet = &preminedTestnet{
 }
 
 type preminedTestnet struct {
-	target encPubkey
+	target enode.PubkeyEncoded
 	dists  [hashBits + 1][]*ecdsa.PrivateKey
 }
 
@@ -311,7 +311,7 @@ func (tn *preminedTestnet) closest(n int) (nodes []*enode.Node) {
 		}
 	}
 	sort.Slice(nodes, func(i, j int) bool {
-		return enode.DistCmp(tn.target.id(), nodes[i].ID(), nodes[j].ID()) < 0
+		return enode.DistCmp(tn.target.ID(), nodes[i].ID(), nodes[j].ID()) < 0
 	})
 	return nodes[:n]
 }
@@ -326,11 +326,11 @@ func (tn *preminedTestnet) mine() {
 		tn.dists[i] = nil
 	}
 
-	targetSha := tn.target.id()
+	targetSha := tn.target.ID()
 	found, need := 0, 40
 	for found < need {
 		k := newkey()
-		ld := enode.LogDist(targetSha, encodePubkey(&k.PublicKey).id())
+		ld := enode.LogDist(targetSha, enode.PubkeyToIDV4(&k.PublicKey))
 		if len(tn.dists[ld]) < 8 {
 			tn.dists[ld] = append(tn.dists[ld], k)
 			found++

diff --git a/p2p/discover/v4_udp.go b/p2p/discover/v4_udp.go
@@ -21,7 +21,6 @@ import (
 	"container/list"
 	"context"
 	"crypto/ecdsa"
-	crand "crypto/rand"
 	"errors"
 	"fmt"
 	"io"
@@ -266,7 +265,7 @@ func (t *UDPv4) LookupPubkey(key *ecdsa.PublicKey) []*enode.Node {
 		// case and run the bootstrapping logic.
 		<-t.tab.refresh()
 	}
-	return t.newLookup(t.closeCtx, encodePubkey(key)).run()
+	return t.newLookup(t.closeCtx, key).run()
 }
 
 // RandomNodes is an iterator yielding nodes from a random walk of the DHT.
@@ -281,20 +280,23 @@ func (t *UDPv4) lookupRandom() []*enode.Node {
 
 // lookupSelf implements transport.
 func (t *UDPv4) lookupSelf() []*enode.Node {
-	return t.newLookup(t.closeCtx, encodePubkey(&t.priv.PublicKey)).run()
+	return t.newLookup(t.closeCtx, &t.priv.PublicKey).run()
 }
 
 func (t *UDPv4) newRandomLookup(ctx context.Context) *lookup {
-	var target encPubkey
-	crand.Read(target[:])
-	return t.newLookup(ctx, target)
+	key, err := crypto.GenerateKey()
+	if err != nil {
+		t.log.Warn("Failed to generate a random node key for newRandomLookup", "err", err)
+		key = t.priv
+	}
+	return t.newLookup(ctx, &key.PublicKey)
 }
 
-func (t *UDPv4) newLookup(ctx context.Context, targetKey encPubkey) *lookup {
-	target := enode.ID(crypto.Keccak256Hash(targetKey[:]))
-	ekey := v4wire.Pubkey(targetKey)
+func (t *UDPv4) newLookup(ctx context.Context, targetKey *ecdsa.PublicKey) *lookup {
+	targetKeyEnc := v4wire.EncodePubkey(targetKey)
+	target := enode.PubkeyEncoded(targetKeyEnc).ID()
 	it := newLookup(ctx, t.tab, target, func(n *node) ([]*node, error) {
-		return t.findnode(n.ID(), n.addr(), ekey)
+		return t.findnode(n.ID(), n.addr(), targetKeyEnc)
 	})
 	return it
 }
@@ -565,7 +567,7 @@ func (t *UDPv4) handlePacket(from *net.UDPAddr, buf []byte) error {
 		return err
 	}
 	packet := t.wrapPacket(rawpacket)
-	fromID := fromKey.ID()
+	fromID := enode.PubkeyEncoded(fromKey).ID()
 	if packet.preverify != nil {
 		err = packet.preverify(packet, from, fromID, fromKey)
 	}
@@ -741,7 +743,7 @@ func (t *UDPv4) handleFindnode(h *packetHandlerV4, from *net.UDPAddr, fromID eno
 	req := h.Packet.(*v4wire.Findnode)
 
 	// Determine closest nodes.
-	target := enode.ID(crypto.Keccak256Hash(req.Target[:]))
+	target := enode.PubkeyEncoded(req.Target).ID()
 	closest := t.tab.findnodeByID(target, bucketSize, true).entries
 
 	// Send neighbors in chunks with at most maxNeighbors per packet

diff --git a/p2p/discover/v4_udp_test.go b/p2p/discover/v4_udp_test.go
@@ -264,7 +264,8 @@ func TestUDPv4_findnode(t *testing.T) {
 	// put a few nodes into the table. their exact
 	// distribution shouldn't matter much, although we need to
 	// take care not to overflow any bucket.
-	nodes := &nodesByDistance{target: testTarget.ID()}
+	testTargetID := enode.PubkeyEncoded(testTarget).ID()
+	nodes := &nodesByDistance{target: testTargetID}
 	live := make(map[enode.ID]bool)
 	numCandidates := 2 * bucketSize
 	for i := 0; i < numCandidates; i++ {
@@ -282,23 +283,24 @@ func TestUDPv4_findnode(t *testing.T) {
 
 	// ensure there's a bond with the test node,
 	// findnode won't be accepted otherwise.
-	remoteID := v4wire.EncodePubkey(&test.remotekey.PublicKey).ID()
+	remoteID := enode.PubkeyToIDV4(&test.remotekey.PublicKey)
 	test.table.db.UpdateLastPongReceived(remoteID, test.remoteaddr.IP, time.Now())
 
 	// check that closest neighbors are returned.
-	expected := test.table.findnodeByID(testTarget.ID(), bucketSize, true)
+	expected := test.table.findnodeByID(testTargetID, bucketSize, true)
 	test.packetIn(nil, &v4wire.Findnode{Target: testTarget, Expiration: futureExp})
 	waitNeighbors := func(want []*node) {
 		test.waitPacketOut(func(p *v4wire.Neighbors, to *net.UDPAddr, hash []byte) {
 			if len(p.Nodes) != len(want) {
 				t.Errorf("wrong number of results: got %d, want %d", len(p.Nodes), bucketSize)
 			}
 			for i, n := range p.Nodes {
-				if n.ID.ID() != want[i].ID() {
+				nodeID := enode.PubkeyEncoded(n.ID).ID()
+				if nodeID != want[i].ID() {
 					t.Errorf("result mismatch at %d:\n  got:  %v\n  want: %v", i, n, expected.entries[i])
 				}
-				if !live[n.ID.ID()] {
-					t.Errorf("result includes dead node %v", n.ID.ID())
+				if !live[nodeID] {
+					t.Errorf("result includes dead node %v", nodeID)
 				}
 			}
 		})
@@ -322,7 +324,7 @@ func TestUDPv4_findnodeMultiReply(t *testing.T) {
 	// queue a pending findnode request
 	resultc, errc := make(chan []*node), make(chan error)
 	go func() {
-		rid := encodePubkey(&test.remotekey.PublicKey).id()
+		rid := enode.PubkeyToIDV4(&test.remotekey.PublicKey)
 		ns, err := test.udp.findnode(rid, test.remoteaddr, testTarget)
 		if err != nil && len(ns) == 0 {
 			errc <- err
@@ -446,7 +448,7 @@ func TestUDPv4_successfulPing(t *testing.T) {
 	// pong packet.
 	select {
 	case n := <-added:
-		rid := encodePubkey(&test.remotekey.PublicKey).id()
+		rid := enode.PubkeyToIDV4(&test.remotekey.PublicKey)
 		if n.ID() != rid {
 			t.Errorf("node has wrong ID: got %v, want %v", n.ID(), rid)
 		}

diff --git a/p2p/discover/v4wire/v4wire.go b/p2p/discover/v4wire/v4wire.go
@@ -29,7 +29,6 @@ import (
 
 	"github.com/ledgerwatch/erigon/common/math"
 	"github.com/ledgerwatch/erigon/crypto"
-	"github.com/ledgerwatch/erigon/p2p/enode"
 	"github.com/ledgerwatch/erigon/p2p/enr"
 	"github.com/ledgerwatch/erigon/rlp"
 )
@@ -125,11 +124,6 @@ const MaxNeighbors = 12
 // Pubkey represents an encoded 64-byte secp256k1 public key.
 type Pubkey [64]byte
 
-// ID returns the node ID corresponding to the public key.
-func (e Pubkey) ID() enode.ID {
-	return enode.ID(crypto.Keccak256Hash(e[:]))
-}
-
 // Node represents information about a node.
 type Node struct {
 	IP  net.IP // len 4 for IPv4 or 16 for IPv6
@@ -192,6 +186,7 @@ func seqFromTail(tail []rlp.RawValue) uint64 {
 		return 0
 	}
 	var seq uint64
+	//goland:noinspection GoUnhandledErrorResult
 	rlp.DecodeBytes(tail[0], &seq) //nolint:errcheck
 	return seq
 }
@@ -279,7 +274,8 @@ func recoverNodeKey(hash, sig []byte) (key Pubkey, err error) {
 	return key, nil
 }
 
-// EncodePubkey encodes a secp256k1 public key.
+// EncodePubkey converts a public key into a binary format.
+// The logic matches DecodePubkey.
 func EncodePubkey(key *ecdsa.PublicKey) Pubkey {
 	var e Pubkey
 	math.ReadBits(key.X, e[:len(e)/2])

diff --git a/p2p/discover/v5_udp_test.go b/p2p/discover/v5_udp_test.go
@@ -565,7 +565,7 @@ func TestUDPv5_lookup(t *testing.T) {
 	test := newUDPV5Test(t)
 
 	// Lookup on empty table returns no nodes.
-	if results := test.udp.Lookup(lookupTestnet.target.id()); len(results) > 0 {
+	if results := test.udp.Lookup(lookupTestnet.target.ID()); len(results) > 0 {
 		t.Fatalf("lookup on empty table returned %d results: %#v", len(results), results)
 	}
 
@@ -584,7 +584,7 @@ func TestUDPv5_lookup(t *testing.T) {
 	// Start the lookup.
 	resultC := make(chan []*enode.Node, 1)
 	go func() {
-		resultC <- test.udp.Lookup(lookupTestnet.target.id())
+		resultC <- test.udp.Lookup(lookupTestnet.target.ID())
 		test.close()
 	}()
 
@@ -768,7 +768,7 @@ func (test *udpV5Test) packetInFrom(key *ecdsa.PrivateKey, addr *net.UDPAddr, pa
 
 // getNode ensures the test knows about a node at the given endpoint.
 func (test *udpV5Test) getNode(key *ecdsa.PrivateKey, addr *net.UDPAddr) *enode.LocalNode {
-	id := encodePubkey(&key.PublicKey).id()
+	id := enode.PubkeyToIDV4(&key.PublicKey)
 	ln := test.nodesByID[id]
 	if ln == nil {
 		db, err := enode.OpenDB(test.t.TempDir())

diff --git a/p2p/enode/idscheme.go b/p2p/enode/idscheme.go
@@ -21,8 +21,8 @@ import (
 	"fmt"
 	"io"
 
-	"github.com/ledgerwatch/erigon/common/math"
 	"github.com/ledgerwatch/erigon/crypto"
+	"github.com/ledgerwatch/erigon/p2p/discover/v4wire"
 	"github.com/ledgerwatch/erigon/p2p/enr"
 	"github.com/ledgerwatch/erigon/rlp"
 	"golang.org/x/crypto/sha3"
@@ -83,10 +83,20 @@ func (V4ID) NodeAddr(r *enr.Record) []byte {
 	if err != nil {
 		return nil
 	}
-	buf := make([]byte, 64)
-	math.ReadBits(pubkey.X, buf[:32])
-	math.ReadBits(pubkey.Y, buf[32:])
-	return crypto.Keccak256(buf)
+	id := PubkeyToIDV4((*ecdsa.PublicKey)(&pubkey))
+	return id[:]
+}
+
+// PubkeyToIDV4 derives the v4 node address from the given public key.
+func PubkeyToIDV4(key *ecdsa.PublicKey) ID {
+	return PubkeyEncoded(v4wire.EncodePubkey(key)).ID()
+}
+
+type PubkeyEncoded v4wire.Pubkey
+
+// ID returns the node ID corresponding to the public key.
+func (e PubkeyEncoded) ID() ID {
+	return ID(crypto.Keccak256Hash(e[:]))
 }
 
 // Secp256k1 is the "secp256k1" key, which holds a public key.

diff --git a/p2p/enode/urlv4.go b/p2p/enode/urlv4.go
@@ -26,7 +26,6 @@ import (
 	"regexp"
 	"strconv"
 
-	"github.com/ledgerwatch/erigon/common/math"
 	"github.com/ledgerwatch/erigon/crypto"
 	"github.com/ledgerwatch/erigon/p2p/enr"
 )
@@ -193,11 +192,3 @@ func (n *Node) URLv4() string {
 	}
 	return u.String()
 }
-
-// PubkeyToIDV4 derives the v4 node address from the given public key.
-func PubkeyToIDV4(key *ecdsa.PublicKey) ID {
-	e := make([]byte, 64)
-	math.ReadBits(key.X, e[:len(e)/2])
-	math.ReadBits(key.Y, e[len(e)/2:])
-	return ID(crypto.Keccak256Hash(e))
-}