3 files changed, 96 insertions, 16 deletions
diff --git a/p2p/enr/enr.go b/p2p/enr/enr.go
index 444820c15..c36ae9e3e 100644
--- a/p2p/enr/enr.go
+++ b/p2p/enr/enr.go
@@ -163,6 +163,16 @@ func (r *Record) invalidate() {
 	r.raw = nil
 }
 
+// Signature returns the signature of the record.
+func (r *Record) Signature() []byte {
+	if r.signature == nil {
+		return nil
+	}
+	cpy := make([]byte, len(r.signature))
+	copy(cpy, r.signature)
+	return cpy
+}
+
 // EncodeRLP implements rlp.Encoder. Encoding fails if
 // the record is unsigned.
 func (r Record) EncodeRLP(w io.Writer) error {
@@ -173,7 +183,7 @@ func (r Record) EncodeRLP(w io.Writer) error {
 	return err
 }
 
-// DecodeRLP implements rlp.Decoder. Decoding verifies the signature.
+// DecodeRLP implements rlp.Decoder. Decoding doesn't verify the signature.
 func (r *Record) DecodeRLP(s *rlp.Stream) error {
 	dec, raw, err := decodeRecord(s)
 	if err != nil {
diff --git a/p2p/enr/enr_test.go b/p2p/enr/enr_test.go
index 449c898a8..434685e0b 100644
--- a/p2p/enr/enr_test.go
+++ b/p2p/enr/enr_test.go
@@ -49,23 +49,23 @@ func TestGetSetID(t *testing.T) {
 }
 
 // TestGetSetIP4 tests encoding/decoding and setting/getting of the IP key.
-func TestGetSetIP4(t *testing.T) {
-	ip := IP{192, 168, 0, 3}
+func TestGetSetIPv4(t *testing.T) {
+	ip := IPv4{192, 168, 0, 3}
 	var r Record
 	r.Set(ip)
 
-	var ip2 IP
+	var ip2 IPv4
 	require.NoError(t, r.Load(&ip2))
 	assert.Equal(t, ip, ip2)
 }
 
-// TestGetSetIP6 tests encoding/decoding and setting/getting of the IP key.
-func TestGetSetIP6(t *testing.T) {
-	ip := IP{0x20, 0x01, 0x48, 0x60, 0, 0, 0x20, 0x01, 0, 0, 0, 0, 0, 0, 0x00, 0x68}
+// TestGetSetIP6 tests encoding/decoding and setting/getting of the IP6 key.
+func TestGetSetIPv6(t *testing.T) {
+	ip := IPv6{0x20, 0x01, 0x48, 0x60, 0, 0, 0x20, 0x01, 0, 0, 0, 0, 0, 0, 0x00, 0x68}
 	var r Record
 	r.Set(ip)
 
-	var ip2 IP
+	var ip2 IPv6
 	require.NoError(t, r.Load(&ip2))
 	assert.Equal(t, ip, ip2)
 }
@@ -83,7 +83,7 @@ func TestGetSetUDP(t *testing.T) {
 
 func TestLoadErrors(t *testing.T) {
 	var r Record
-	ip4 := IP{127, 0, 0, 1}
+	ip4 := IPv4{127, 0, 0, 1}
 	r.Set(ip4)
 
 	// Check error for missing keys.
@@ -185,13 +185,13 @@ func TestSeq(t *testing.T) {
 func TestGetSetOverwrite(t *testing.T) {
 	var r Record
 
-	ip := IP{192, 168, 0, 3}
+	ip := IPv4{192, 168, 0, 3}
 	r.Set(ip)
 
-	ip2 := IP{192, 168, 0, 4}
+	ip2 := IPv4{192, 168, 0, 4}
 	r.Set(ip2)
 
-	var ip3 IP
+	var ip3 IPv4
 	require.NoError(t, r.Load(&ip3))
 	assert.Equal(t, ip2, ip3)
 }
@@ -200,7 +200,7 @@ func TestGetSetOverwrite(t *testing.T) {
 func TestSignEncodeAndDecode(t *testing.T) {
 	var r Record
 	r.Set(UDP(30303))
-	r.Set(IP{127, 0, 0, 1})
+	r.Set(IPv4{127, 0, 0, 1})
 	require.NoError(t, signTest([]byte{5}, &r))
 
 	blob, err := rlp.EncodeToBytes(r)
diff --git a/p2p/enr/entries.go b/p2p/enr/entries.go
index 347990ab6..f2118401a 100644
--- a/p2p/enr/entries.go
+++ b/p2p/enr/entries.go
@@ -60,11 +60,21 @@ type TCP uint16
 
 func (v TCP) ENRKey() string { return "tcp" }
 
+// UDP is the "udp" key, which holds the IPv6-specific UDP port of the node.
+type TCP6 uint16
+
+func (v TCP6) ENRKey() string { return "tcp6" }
+
 // UDP is the "udp" key, which holds the UDP port of the node.
 type UDP uint16
 
 func (v UDP) ENRKey() string { return "udp" }
 
+// UDP is the "udp" key, which holds the IPv6-specific UDP port of the node.
+type UDP6 uint16
+
+func (v UDP6) ENRKey() string { return "udp6" }
+
 // ID is the "id" key, which holds the name of the identity scheme.
 type ID string
 
@@ -72,17 +82,27 @@ const IDv4 = ID("v4") // the default identity scheme
 
 func (v ID) ENRKey() string { return "id" }
 
-// IP is the "ip" key, which holds the IP address of the node.
+// IP is either the "ip" or "ip6" key, depending on the value.
+// Use this value to encode IP addresses that can be either v4 or v6.
+// To load an address from a record use the IPv4 or IPv6 types.
 type IP net.IP
 
-func (v IP) ENRKey() string { return "ip" }
+func (v IP) ENRKey() string {
+	if net.IP(v).To4() == nil {
+		return "ip6"
+	}
+	return "ip"
+}
 
 // EncodeRLP implements rlp.Encoder.
 func (v IP) EncodeRLP(w io.Writer) error {
 	if ip4 := net.IP(v).To4(); ip4 != nil {
 		return rlp.Encode(w, ip4)
 	}
-	return rlp.Encode(w, net.IP(v))
+	if ip6 := net.IP(v).To16(); ip6 != nil {
+		return rlp.Encode(w, ip6)
+	}
+	return fmt.Errorf("invalid IP address: %v", net.IP(v))
 }
 
 // DecodeRLP implements rlp.Decoder.
@@ -96,6 +116,56 @@ func (v *IP) DecodeRLP(s *rlp.Stream) error {
 	return nil
 }
 
+// IPv4 is the "ip" key, which holds the IP address of the node.
+type IPv4 net.IP
+
+func (v IPv4) ENRKey() string { return "ip" }
+
+// EncodeRLP implements rlp.Encoder.
+func (v IPv4) EncodeRLP(w io.Writer) error {
+	ip4 := net.IP(v).To4()
+	if ip4 == nil {
+		return fmt.Errorf("invalid IPv4 address: %v", net.IP(v))
+	}
+	return rlp.Encode(w, ip4)
+}
+
+// DecodeRLP implements rlp.Decoder.
+func (v *IPv4) DecodeRLP(s *rlp.Stream) error {
+	if err := s.Decode((*net.IP)(v)); err != nil {
+		return err
+	}
+	if len(*v) != 4 {
+		return fmt.Errorf("invalid IPv4 address, want 4 bytes: %v", *v)
+	}
+	return nil
+}
+
+// IPv6 is the "ip6" key, which holds the IP address of the node.
+type IPv6 net.IP
+
+func (v IPv6) ENRKey() string { return "ip6" }
+
+// EncodeRLP implements rlp.Encoder.
+func (v IPv6) EncodeRLP(w io.Writer) error {
+	ip6 := net.IP(v).To16()
+	if ip6 == nil {
+		return fmt.Errorf("invalid IPv6 address: %v", net.IP(v))
+	}
+	return rlp.Encode(w, ip6)
+}
+
+// DecodeRLP implements rlp.Decoder.
+func (v *IPv6) DecodeRLP(s *rlp.Stream) error {
+	if err := s.Decode((*net.IP)(v)); err != nil {
+		return err
+	}
+	if len(*v) != 16 {
+		return fmt.Errorf("invalid IPv6 address, want 16 bytes: %v", *v)
+	}
+	return nil
+}
+
 // KeyError is an error related to a key.
 type KeyError struct {
 	Key string