From e489d77d4bba591789db31099b2008a1c1c644c2 Mon Sep 17 00:00:00 2001
From: lash <dev@holbrook.no>
Date: Fri, 1 Mar 2019 13:51:59 +0100
Subject: [PATCH] swarm/network: Use enode package for enode rewrite instead of
 regex

---
 swarm/network/protocol.go      | 23 ++++++-----------------
 swarm/network/protocol_test.go | 13 +++----------
 2 files changed, 9 insertions(+), 27 deletions(-)

diff --git a/swarm/network/protocol.go b/swarm/network/protocol.go
index 4161907d74..f014ea0d1b 100644
--- a/swarm/network/protocol.go
+++ b/swarm/network/protocol.go
@@ -21,7 +21,6 @@ import (
 	"errors"
 	"fmt"
 	"net"
-	"regexp"
 	"sync"
 	"time"
 
@@ -40,8 +39,6 @@ const (
 	bzzHandshakeTimeout = 3000 * time.Millisecond
 )
 
-var regexpEnodeIP = regexp.MustCompile("@(.+):([0-9]+)")
-
 // BzzSpec is the spec of the generic swarm handshake
 var BzzSpec = &protocols.Spec{
 	Name:       "bzz",
@@ -217,7 +214,7 @@ func (b *Bzz) performHandshake(p *protocols.Peer, handshake *HandshakeMsg) error
 		return err
 	}
 	handshake.peerAddr = rsh.(*HandshakeMsg).Addr
-	sanitizeEnodeRemote(p.RemoteAddr(), handshake.peerAddr)
+	sanitizeEnodeRemote(p.Node(), handshake.peerAddr)
 	handshake.LightNode = rsh.(*HandshakeMsg).LightNode
 	return nil
 }
@@ -226,20 +223,12 @@ func (b *Bzz) performHandshake(p *protocols.Peer, handshake *HandshakeMsg) error
 // this method ensures that if this default is used in a networked environment, we replace
 // the ip with the one applicable on the interface the connection came in on
 // it modifies the passed bzzaddr in place, and returns the same pointer
-func sanitizeEnodeRemote(paddr net.Addr, baddr *BzzAddr) {
-	hsSubmatch := regexpEnodeIP.FindSubmatch(baddr.UAddr)
-	ip, _, err := net.SplitHostPort(paddr.String())
-	if len(hsSubmatch) < 2 {
-		log.Warn("sanitize found non ipv4 string", "remotestring", paddr.String(), "handshakeaddr", baddr)
-	} else if err == nil {
-		hsip := net.ParseIP(string(hsSubmatch[1]))
-		if hsip != nil && hsip.IsLoopback() {
-			remoteStr := fmt.Sprintf("@%s:%s", ip, string(hsSubmatch[2]))
-			log.Debug("rewrote peer uaddr host/port", "addr", baddr)
-			baddr.UAddr = regexpEnodeIP.ReplaceAll(baddr.UAddr, []byte(remoteStr))
+func sanitizeEnodeRemote(paddr *enode.Node, baddr *BzzAddr) {
+	enod, err := enode.ParseV4(string(baddr.UAddr))
+	if err == nil {
+		if enod.IP().IsLoopback() {
+			baddr.UAddr = []byte(paddr.String())
 		}
-	} else {
-		log.Trace("passthrough handshake addr rewrite", "submatch", hsSubmatch[1])
 	}
 }
 
diff --git a/swarm/network/protocol_test.go b/swarm/network/protocol_test.go
index 9516f69801..105168ad90 100644
--- a/swarm/network/protocol_test.go
+++ b/swarm/network/protocol_test.go
@@ -263,26 +263,19 @@ func TestSanitizeEnodeRemote(t *testing.T) {
 		t.Fatal(err)
 	}
 	remoteIP := net.IPv4(0x80, 0x40, 0x20, 0x10)
-	remoteAddr := net.TCPAddr{
-		IP:   remoteIP,
-		Port: 30399,
-	}
 	nodLocal := enode.NewV4(&pk.PublicKey, net.IPv4(0x7f, 0x00, 0x00, 0x01), 30341, 30341)
 	nodRemote := enode.NewV4(&pk.PublicKey, remoteIP, 30341, 30341)
 	baddr := RandomAddr()
 	oldUAddr := []byte(nodLocal.String())
 	baddr.UAddr = oldUAddr
-	sanitizeEnodeRemote(&remoteAddr, baddr)
+	sanitizeEnodeRemote(nodRemote, baddr)
 	if !bytes.Equal(baddr.UAddr, []byte(nodRemote.String())) {
 		t.Fatalf("insane address. expected %v, got %v", nodRemote.String(), string(baddr.UAddr))
 	}
 
 	remoteIP = net.IPv4(0x04, 0x04, 0x04, 0x04)
-	remoteAddr = net.TCPAddr{
-		IP:   remoteIP,
-		Port: 30399,
-	}
-	sanitizeEnodeRemote(&remoteAddr, baddr)
+	nodRemoteTwo := enode.NewV4(&pk.PublicKey, remoteIP, 30341, 30341)
+	sanitizeEnodeRemote(nodRemoteTwo, baddr)
 	if !bytes.Equal(baddr.UAddr, []byte(nodRemote.String())) {
 		t.Fatalf("Should not have rewritten non-localhost string. expected %v, got %v", nodRemote.String(), string(baddr.UAddr))
 	}