From 3414632d94ea9a865b00ab2cdd0303446716dcd2 Mon Sep 17 00:00:00 2001
From: lash <dev@holbrook.no>
Date: Fri, 1 Mar 2019 16:08:52 +0100
Subject: [PATCH] swarm/network: Make sanitize method of bzzaddr

---
 swarm/network/protocol.go      | 32 ++++++++++++++++----------------
 swarm/network/protocol_test.go |  4 ++--
 2 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/swarm/network/protocol.go b/swarm/network/protocol.go
index 38ed5d37e2..175e93e1ac 100644
--- a/swarm/network/protocol.go
+++ b/swarm/network/protocol.go
@@ -214,26 +214,11 @@ func (b *Bzz) performHandshake(p *protocols.Peer, handshake *HandshakeMsg) error
 		return err
 	}
 	handshake.peerAddr = rsh.(*HandshakeMsg).Addr
-	sanitizeEnodeRemote(p.Node(), handshake.peerAddr)
+	handshake.peerAddr.sanitizeEnodeRemote(p.Node())
 	handshake.LightNode = rsh.(*HandshakeMsg).LightNode
 	return nil
 }
 
-// if started without the natip argument, the enode string will be localhost
-// 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 *enode.Node, baddr *BzzAddr) {
-	enod, err := enode.ParseV4(string(baddr.UAddr))
-	if err != nil {
-		log.Error("invalid bzz address", "addr", baddr)
-		return
-	}
-	if enod.IP().IsLoopback() {
-		baddr.UAddr = []byte(paddr.String())
-	}
-}
-
 // runBzz is the p2p protocol run function for the bzz base protocol
 // that negotiates the bzz handshake
 func (b *Bzz) runBzz(p *p2p.Peer, rw p2p.MsgReadWriter) error {
@@ -389,6 +374,21 @@ func (a *BzzAddr) String() string {
 	return fmt.Sprintf("%x <%s>", a.OAddr, a.UAddr)
 }
 
+// if started without the natip argument, the enode string will be localhost
+// 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 (b *BzzAddr) sanitizeEnodeRemote(paddr *enode.Node) {
+	enod, err := enode.ParseV4(string(b.UAddr))
+	if err != nil {
+		log.Error("invalid bzz address", "addr", b)
+		return
+	}
+	if enod.IP().IsLoopback() {
+		b.UAddr = []byte(paddr.String())
+	}
+}
+
 // RandomAddr is a utility method generating an address from a public key
 func RandomAddr() *BzzAddr {
 	key, err := crypto.GenerateKey()
diff --git a/swarm/network/protocol_test.go b/swarm/network/protocol_test.go
index 105168ad90..867bf5820f 100644
--- a/swarm/network/protocol_test.go
+++ b/swarm/network/protocol_test.go
@@ -268,14 +268,14 @@ func TestSanitizeEnodeRemote(t *testing.T) {
 	baddr := RandomAddr()
 	oldUAddr := []byte(nodLocal.String())
 	baddr.UAddr = oldUAddr
-	sanitizeEnodeRemote(nodRemote, baddr)
+	baddr.sanitizeEnodeRemote(nodRemote)
 	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)
 	nodRemoteTwo := enode.NewV4(&pk.PublicKey, remoteIP, 30341, 30341)
-	sanitizeEnodeRemote(nodRemoteTwo, baddr)
+	baddr.sanitizeEnodeRemote(nodRemoteTwo)
 	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))
 	}