From b294ee6f6f1f2d7fbe4e2cd047cf81035db67dc6 Mon Sep 17 00:00:00 2001
From: Richard Ramos <info@richardramos.me>
Date: Thu, 31 Mar 2022 21:17:14 -0400
Subject: [PATCH] fix: port on discv5 (#224)

---
 waku/v2/discv5/discover.go | 32 +++++++++++++++++++++-----------
 waku/v2/node/wakunode2.go  | 25 +++++++++++++++++++++----
 2 files changed, 42 insertions(+), 15 deletions(-)

diff --git a/waku/v2/discv5/discover.go b/waku/v2/discv5/discover.go
index 18334f30..d7e42c95 100644
--- a/waku/v2/discv5/discover.go
+++ b/waku/v2/discv5/discover.go
@@ -43,6 +43,12 @@ type DiscoveryV5 struct {
 	wg *sync.WaitGroup
 
 	peerCache peerCache
+
+	// Used for those weird cases where updateAddress
+	// receives the same external address twice both with the original port
+	// and the nat port. Ideally this atribute should be removed by doing
+	// hole punching before starting waku
+	ogTCPPort int
 }
 
 type peerCache struct {
@@ -145,7 +151,8 @@ func NewDiscoveryV5(host host.Host, addresses []ma.Multiaddr, priv *ecdsa.Privat
 			IP:   net.IPv4zero,
 			Port: params.udpPort,
 		},
-		log: logger,
+		log:       logger,
+		ogTCPPort: ipAddr.Port,
 	}, nil
 }
 
@@ -241,7 +248,7 @@ func (d *DiscoveryV5) Stop() {
 	d.wg.Wait()
 }
 
-func (d *DiscoveryV5) UpdateAddr(addr net.IP) error {
+func (d *DiscoveryV5) UpdateAddr(addr *net.TCPAddr) error {
 	if !d.params.autoUpdate {
 		return nil
 	}
@@ -249,18 +256,21 @@ func (d *DiscoveryV5) UpdateAddr(addr net.IP) error {
 	d.Lock()
 	defer d.Unlock()
 
-	if addr.IsUnspecified() || d.localnode.Node().IP().Equal(addr) {
+	// TODO: This code is not elegant and should be improved
+
+	if !isExternal(addr) && !isExternal(&net.TCPAddr{IP: d.localnode.Node().IP()}) {
+		if !((d.localnode.Node().IP().IsLoopback() && isPrivate(addr)) || (isPrivate(&net.TCPAddr{IP: d.localnode.Node().IP()}) && isExternal(addr))) {
+			return nil
+		}
+	}
+
+	if addr.IP.IsUnspecified() || (d.localnode.Node().IP().Equal(addr.IP) && addr.Port == d.ogTCPPort) {
 		return nil
 	}
 
-	// TODO: improve this logic to determine if an address should be replaced or not
-	if !((d.localnode.Node().IP().IsLoopback() && isPrivate(&net.TCPAddr{IP: addr})) ||
-		(isPrivate(&net.TCPAddr{IP: d.localnode.Node().IP()}) && isExternal(&net.TCPAddr{IP: addr}))) {
-		return nil
-	}
-
-	d.localnode.SetStaticIP(addr)
-	d.log.Info(fmt.Sprintf("Updated Discovery V5 node IP: %s", d.localnode.Node().IP()))
+	d.localnode.SetStaticIP(addr.IP)
+	d.localnode.Set(enr.TCP(uint16(addr.Port))) // lgtm [go/incorrect-integer-conversion]
+	d.log.Info(fmt.Sprintf("Updated Discovery V5 node: %s:%d", d.localnode.Node().IP(), d.localnode.Node().TCP()))
 	d.log.Info("Discovery V5 ", d.localnode.Node())
 
 	return nil
diff --git a/waku/v2/node/wakunode2.go b/waku/v2/node/wakunode2.go
index 7b61f5b5..02f02d59 100644
--- a/waku/v2/node/wakunode2.go
+++ b/waku/v2/node/wakunode2.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"net"
+	"strconv"
 	"sync"
 	"time"
 
@@ -195,13 +196,29 @@ func (w *WakuNode) onAddrChange() {
 			w.log.Error(fmt.Sprintf("could not extract ip from ma %s: %s", m, err.Error()))
 			continue
 		}
-		ip := net.ParseIP(ipStr)
 
-		if !ip.IsLoopback() && !ip.IsUnspecified() {
+		portStr, err := m.ValueForProtocol(ma.P_TCP)
+		if err != nil {
+			w.log.Error(fmt.Sprintf("could not extract port from ma %s: %s", m, err.Error()))
+			continue
+		}
+
+		port, err := strconv.Atoi(portStr)
+		if err != nil {
+			w.log.Error(fmt.Sprintf("could not convert port to int: %s", err.Error()))
+			continue
+		}
+
+		addr := &net.TCPAddr{
+			IP:   net.ParseIP(ipStr),
+			Port: port,
+		}
+
+		if !addr.IP.IsLoopback() && !addr.IP.IsUnspecified() {
 			if w.opts.enableDiscV5 {
-				err := w.discoveryV5.UpdateAddr(ip)
+				err := w.discoveryV5.UpdateAddr(addr)
 				if err != nil {
-					w.log.Error(fmt.Sprintf("could not update DiscV5 address with IP %s: %s", ip, err.Error()))
+					w.log.Error(fmt.Sprintf("could not update DiscV5 address with IP %s:%d %s", addr.IP, addr.Port, err.Error()))
 					continue
 				}
 			}