From 94500ae52b1663458b640b0a01e98a7d9f8a187f Mon Sep 17 00:00:00 2001
From: guillaumemichel <guillaume@michel.id>
Date: Fri, 9 Aug 2024 13:14:27 +0200
Subject: [PATCH] adding quic field to enr

---
 p2p/enode/localnode.go |  2 +-
 p2p/enode/node.go      | 24 +++++++++++++++++++++---
 p2p/enode/node_test.go | 22 +++++++++++++++++-----
 p2p/enr/entries.go     | 10 ++++++++++
 4 files changed, 49 insertions(+), 9 deletions(-)

diff --git a/p2p/enode/localnode.go b/p2p/enode/localnode.go
index 6e79c9cbdc..255b673594 100644
--- a/p2p/enode/localnode.go
+++ b/p2p/enode/localnode.go
@@ -298,7 +298,7 @@ func (ln *LocalNode) sign() {
 		panic(fmt.Errorf("enode: can't verify local record: %v", err))
 	}
 	ln.cur.Store(n)
-	log.Info("New local node record", "seq", ln.seq, "id", n.ID(), "ip", n.IPAddr(), "udp", n.UDP(), "tcp", n.TCP())
+	log.Info("New local node record", "seq", ln.seq, "id", n.ID(), "ip", n.IPAddr(), "udp", n.UDP(), "tcp", n.TCP(), "quic", n.QUIC())
 }
 
 func (ln *LocalNode) bumpSeq() {
diff --git a/p2p/enode/node.go b/p2p/enode/node.go
index cb4ac8d172..c4dcfb8e20 100644
--- a/p2p/enode/node.go
+++ b/p2p/enode/node.go
@@ -38,9 +38,10 @@ type Node struct {
 	r  enr.Record
 	id ID
 	// endpoint information
-	ip  netip.Addr
-	udp uint16
-	tcp uint16
+	ip   netip.Addr
+	udp  uint16
+	tcp  uint16
+	quic uint16
 }
 
 // New wraps a node record. The record must be valid according to the given
@@ -105,6 +106,7 @@ func (n *Node) setIP4(ip netip.Addr) {
 	n.ip = ip
 	n.Load((*enr.UDP)(&n.udp))
 	n.Load((*enr.TCP)(&n.tcp))
+	n.Load((*enr.QUIC)(&n.quic))
 }
 
 func (n *Node) setIP6(ip netip.Addr) {
@@ -119,6 +121,9 @@ func (n *Node) setIP6(ip netip.Addr) {
 	if err := n.Load((*enr.TCP6)(&n.tcp)); err != nil {
 		n.Load((*enr.TCP)(&n.tcp))
 	}
+	if err := n.Load((*enr.QUIC6)(&n.quic)); err != nil {
+		n.Load((*enr.QUIC)(&n.quic))
+	}
 }
 
 // MustParse parses a node record or enode:// URL. It panics if the input is invalid.
@@ -184,6 +189,11 @@ func (n *Node) TCP() int {
 	return int(n.tcp)
 }
 
+// QUIC returns the QUIC port of the node.
+func (n *Node) QUIC() int {
+	return int(n.quic)
+}
+
 // UDPEndpoint returns the announced UDP endpoint.
 func (n *Node) UDPEndpoint() (netip.AddrPort, bool) {
 	if !n.ip.IsValid() || n.ip.IsUnspecified() || n.udp == 0 {
@@ -200,6 +210,14 @@ 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) {
+	if !n.ip.IsValid() || n.ip.IsUnspecified() || n.quic == 0 {
+		return netip.AddrPort{}, false
+	}
+	return netip.AddrPortFrom(n.ip, n.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..05946b50b6 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,14 @@ 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: "ipv4-only-loopback",
 			node: func() *Node {
@@ -222,6 +231,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 test.wantQUIC != test.node.QUIC() {
+				t.Errorf("node has wrong QUIC port %d, want %d", test.node.QUIC(), 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