From f544fc3b4659aeca24a6de83f820dd61ea9b39db Mon Sep 17 00:00:00 2001
From: Guillaume Michel <guillaumemichel@users.noreply.github.com>
Date: Fri, 13 Sep 2024 23:47:18 +0200
Subject: [PATCH] p2p/enode: add quic ENR entry (#30283)

Add `quic` entry to the ENR as proposed in
https://github.com/ethereum/consensus-specs/pull/3644

---------

Co-authored-by: lightclient <lightclient@protonmail.com>
---
 p2p/enode/node.go      | 14 ++++++++
 p2p/enode/node_test.go | 72 +++++++++++++++++++++++++++++++++++++++---
 p2p/enr/entries.go     | 10 ++++++
 3 files changed, 91 insertions(+), 5 deletions(-)

diff --git a/p2p/enode/node.go b/p2p/enode/node.go
index cb4ac8d172..4d93d3f6be 100644
--- a/p2p/enode/node.go
+++ b/p2p/enode/node.go
@@ -200,6 +200,20 @@ func (n *Node) TCPEndpoint() (netip.AddrPort, bool) {
 	return netip.AddrPortFrom(n.ip, n.tcp), true
 }
 
+// QUICEndpoint returns the announced QUIC endpoint.
+func (n *Node) QUICEndpoint() (netip.AddrPort, bool) {
+	var quic uint16
+	if n.ip.Is4() || n.ip.Is4In6() {
+		n.Load((*enr.QUIC)(&quic))
+	} else if n.ip.Is6() {
+		n.Load((*enr.QUIC6)(&quic))
+	}
+	if !n.ip.IsValid() || n.ip.IsUnspecified() || quic == 0 {
+		return netip.AddrPort{}, false
+	}
+	return netip.AddrPortFrom(n.ip, quic), true
+}
+
 // Pubkey returns the secp256k1 public key of the node, if present.
 func (n *Node) Pubkey() *ecdsa.PublicKey {
 	var key ecdsa.PublicKey
diff --git a/p2p/enode/node_test.go b/p2p/enode/node_test.go
index 56e196e82e..f38c77415e 100644
--- a/p2p/enode/node_test.go
+++ b/p2p/enode/node_test.go
@@ -68,11 +68,12 @@ func TestPythonInterop(t *testing.T) {
 func TestNodeEndpoints(t *testing.T) {
 	id := HexID("00000000000000806ad9b61fa5ae014307ebdc964253adcd9f2c0a392aa11abc")
 	type endpointTest struct {
-		name    string
-		node    *Node
-		wantIP  netip.Addr
-		wantUDP int
-		wantTCP int
+		name     string
+		node     *Node
+		wantIP   netip.Addr
+		wantUDP  int
+		wantTCP  int
+		wantQUIC int
 	}
 	tests := []endpointTest{
 		{
@@ -98,6 +99,22 @@ func TestNodeEndpoints(t *testing.T) {
 				return SignNull(&r, id)
 			}(),
 		},
+		{
+			name: "quic-only",
+			node: func() *Node {
+				var r enr.Record
+				r.Set(enr.QUIC(9000))
+				return SignNull(&r, id)
+			}(),
+		},
+		{
+			name: "quic6-only",
+			node: func() *Node {
+				var r enr.Record
+				r.Set(enr.QUIC6(9000))
+				return SignNull(&r, id)
+			}(),
+		},
 		{
 			name: "ipv4-only-loopback",
 			node: func() *Node {
@@ -209,6 +226,48 @@ func TestNodeEndpoints(t *testing.T) {
 			wantIP:  netip.MustParseAddr("192.168.2.2"),
 			wantUDP: 30304,
 		},
+		{
+			name: "ipv4-quic",
+			node: func() *Node {
+				var r enr.Record
+				r.Set(enr.IPv4Addr(netip.MustParseAddr("99.22.33.1")))
+				r.Set(enr.QUIC(9001))
+				return SignNull(&r, id)
+			}(),
+			wantIP:   netip.MustParseAddr("99.22.33.1"),
+			wantQUIC: 9001,
+		},
+		{ // Because the node is IPv4, the quic6 entry won't be loaded.
+			name: "ipv4-quic6",
+			node: func() *Node {
+				var r enr.Record
+				r.Set(enr.IPv4Addr(netip.MustParseAddr("99.22.33.1")))
+				r.Set(enr.QUIC6(9001))
+				return SignNull(&r, id)
+			}(),
+			wantIP: netip.MustParseAddr("99.22.33.1"),
+		},
+		{
+			name: "ipv6-quic",
+			node: func() *Node {
+				var r enr.Record
+				r.Set(enr.IPv6Addr(netip.MustParseAddr("2001::ff00:0042:8329")))
+				r.Set(enr.QUIC(9001))
+				return SignNull(&r, id)
+			}(),
+			wantIP: netip.MustParseAddr("2001::ff00:0042:8329"),
+		},
+		{
+			name: "ipv6-quic6",
+			node: func() *Node {
+				var r enr.Record
+				r.Set(enr.IPv6Addr(netip.MustParseAddr("2001::ff00:0042:8329")))
+				r.Set(enr.QUIC6(9001))
+				return SignNull(&r, id)
+			}(),
+			wantIP:   netip.MustParseAddr("2001::ff00:0042:8329"),
+			wantQUIC: 9001,
+		},
 	}
 
 	for _, test := range tests {
@@ -222,6 +281,9 @@ func TestNodeEndpoints(t *testing.T) {
 			if test.wantTCP != test.node.TCP() {
 				t.Errorf("node has wrong TCP port %d, want %d", test.node.TCP(), test.wantTCP)
 			}
+			if quic, _ := test.node.QUICEndpoint(); test.wantQUIC != int(quic.Port()) {
+				t.Errorf("node has wrong QUIC port %d, want %d", quic.Port(), test.wantQUIC)
+			}
 		})
 	}
 }
diff --git a/p2p/enr/entries.go b/p2p/enr/entries.go
index 155ec4c023..58e660c154 100644
--- a/p2p/enr/entries.go
+++ b/p2p/enr/entries.go
@@ -77,6 +77,16 @@ type UDP6 uint16
 
 func (v UDP6) ENRKey() string { return "udp6" }
 
+// QUIC is the "quic" key, which holds the QUIC port of the node.
+type QUIC uint16
+
+func (v QUIC) ENRKey() string { return "quic" }
+
+// QUIC6 is the "quic6" key, which holds the IPv6-specific quic6 port of the node.
+type QUIC6 uint16
+
+func (v QUIC6) ENRKey() string { return "quic6" }
+
 // ID is the "id" key, which holds the name of the identity scheme.
 type ID string