From 608ac82f53fe22495dde4645873caa68202a730b Mon Sep 17 00:00:00 2001
From: Mark Vujevits <mark.vujevits@prezi.com>
Date: Fri, 19 Oct 2018 11:20:26 +0200
Subject: [PATCH] lightnodes should be omitted from the addressbook

---
 swarm/network/kademlia.go      | 20 ++++++++++++--------
 swarm/network/kademlia_test.go | 32 ++++++++++++++++++++++++++++----
 2 files changed, 40 insertions(+), 12 deletions(-)

diff --git a/swarm/network/kademlia.go b/swarm/network/kademlia.go
index 55a0c6f13d..cd94741be3 100644
--- a/swarm/network/kademlia.go
+++ b/swarm/network/kademlia.go
@@ -261,7 +261,7 @@ func (k *Kademlia) On(p *Peer) (uint8, bool) {
 		// found among live peers, do nothing
 		return v
 	})
-	if ins {
+	if ins && !p.BzzPeer.LightNode {
 		a := newEntry(p.BzzAddr)
 		a.conn = p
 		// insert new online peer into addrs
@@ -329,14 +329,18 @@ func (k *Kademlia) Off(p *Peer) {
 	k.lock.Lock()
 	defer k.lock.Unlock()
 	var del bool
-	k.addrs, _, _, _ = pot.Swap(k.addrs, p, pof, func(v pot.Val) pot.Val {
-		// v cannot be nil, must check otherwise we overwrite entry
-		if v == nil {
-			panic(fmt.Sprintf("connected peer not found %v", p))
-		}
+	if !p.BzzPeer.LightNode {
+		k.addrs, _, _, _ = pot.Swap(k.addrs, p, pof, func(v pot.Val) pot.Val {
+			// v cannot be nil, must check otherwise we overwrite entry
+			if v == nil {
+				panic(fmt.Sprintf("connected peer not found %v", p))
+			}
+			del = true
+			return newEntry(p.BzzAddr)
+		})
+	} else {
 		del = true
-		return newEntry(p.BzzAddr)
-	})
+	}
 
 	if del {
 		k.conns, _, _, _ = pot.Swap(k.conns, p, pof, func(_ pot.Val) pot.Val {
diff --git a/swarm/network/kademlia_test.go b/swarm/network/kademlia_test.go
index 903c8dbdac..8af584846c 100644
--- a/swarm/network/kademlia_test.go
+++ b/swarm/network/kademlia_test.go
@@ -46,19 +46,19 @@ func newTestKademlia(b string) *Kademlia {
 	return NewKademlia(base, params)
 }
 
-func newTestKadPeer(k *Kademlia, s string) *Peer {
-	return NewPeer(&BzzPeer{BzzAddr: testKadPeerAddr(s)}, k)
+func newTestKadPeer(k *Kademlia, s string, lightNode bool) *Peer {
+	return NewPeer(&BzzPeer{BzzAddr: testKadPeerAddr(s), LightNode: lightNode}, k)
 }
 
 func On(k *Kademlia, ons ...string) {
 	for _, s := range ons {
-		k.On(newTestKadPeer(k, s))
+		k.On(newTestKadPeer(k, s, false))
 	}
 }
 
 func Off(k *Kademlia, offs ...string) {
 	for _, s := range offs {
-		k.Off(newTestKadPeer(k, s))
+		k.Off(newTestKadPeer(k, s, false))
 	}
 }
 
@@ -254,6 +254,30 @@ func TestSuggestPeerFindPeers(t *testing.T) {
 
 }
 
+func TestOffEffectingAddressBookNormalNode(t *testing.T) {
+	k := newTestKademlia("00000000")
+	k.On(newTestKadPeer(k, "01000000", false))
+	if k.addrs.Size() != 1 {
+		t.Fatal("known peer addresses should contain 1 entry")
+	}
+	k.Off(newTestKadPeer(k, "01000000", false))
+	if k.addrs.Size() != 1 {
+		t.Fatal("known peer addresses should contain 1 entry")
+	}
+}
+
+func TestOffEffectingAddressBookLightNode(t *testing.T) {
+	k := newTestKademlia("00000000")
+	k.On(newTestKadPeer(k, "01000000", true))
+	if k.addrs.Size() != 0 {
+		t.Fatal("known peer addresses should contain 0 entry")
+	}
+	k.Off(newTestKadPeer(k, "01000000", true))
+	if k.addrs.Size() != 0 {
+		t.Fatal("known peer addresses should contain 0 entry")
+	}
+}
+
 func TestSuggestPeerRetries(t *testing.T) {
 	k := newTestKademlia("00000000")
 	k.RetryInterval = int64(300 * time.Millisecond) // cycle