1 files changed, 20 insertions, 36 deletions

diff --git a/p2p/enr/entries.go b/p2p/enr/entries.go
index 7591e6eff..71c7653a2 100644
--- a/p2p/enr/entries.go
+++ b/p2p/enr/entries.go
@@ -57,59 +57,43 @@ func WithEntry(k string, v interface{}) Entry {
 	return &generic{key: k, value: v}
 }
 
-// DiscPort is the "discv5" key, which holds the UDP port for discovery v5.
-type DiscPort uint16
+// TCP is the "tcp" key, which holds the TCP port of the node.
+type TCP uint16
 
-func (v DiscPort) ENRKey() string { return "discv5" }
+func (v TCP) ENRKey() string { return "tcp" }
+
+// UDP is the "udp" key, which holds the UDP port of the node.
+type UDP uint16
+
+func (v UDP) ENRKey() string { return "udp" }
 
 // ID is the "id" key, which holds the name of the identity scheme.
 type ID string
 
+const IDv4 = ID("v4") // the default identity scheme
+
 func (v ID) ENRKey() string { return "id" }
 
-// IP4 is the "ip4" key, which holds a 4-byte IPv4 address.
-type IP4 net.IP
+// IP is the "ip" key, which holds the IP address of the node.
+type IP net.IP
 
-func (v IP4) ENRKey() string { return "ip4" }
+func (v IP) ENRKey() string { return "ip" }
 
 // EncodeRLP implements rlp.Encoder.
-func (v IP4) EncodeRLP(w io.Writer) error {
-	ip4 := net.IP(v).To4()
-	if ip4 == nil {
-		return fmt.Errorf("invalid IPv4 address: %v", v)
-	}
-	return rlp.Encode(w, ip4)
-}
-
-// DecodeRLP implements rlp.Decoder.
-func (v *IP4) 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)
+func (v IP) EncodeRLP(w io.Writer) error {
+	if ip4 := net.IP(v).To4(); ip4 != nil {
+		return rlp.Encode(w, ip4)
 	}
-	return nil
-}
-
-// IP6 is the "ip6" key, which holds a 16-byte IPv6 address.
-type IP6 net.IP
-
-func (v IP6) ENRKey() string { return "ip6" }
-
-// EncodeRLP implements rlp.Encoder.
-func (v IP6) EncodeRLP(w io.Writer) error {
-	ip6 := net.IP(v)
-	return rlp.Encode(w, ip6)
+	return rlp.Encode(w, net.IP(v))
 }
 
 // DecodeRLP implements rlp.Decoder.
-func (v *IP6) DecodeRLP(s *rlp.Stream) error {
+func (v *IP) 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)
+	if len(*v) != 4 && len(*v) != 16 {
+		return fmt.Errorf("invalid IP address, want 4 or 16 bytes: %v", *v)
 	}
 	return nil
 }