From 85110224b47108b5c8659a3f49b7159d6844f08b Mon Sep 17 00:00:00 2001
From: Jae Kwon <jkwon.work@gmail.com>
Date: Wed, 21 Jan 2015 13:40:26 -0800
Subject: [PATCH] Fix PeerSet.Remove() bug.

---
 p2p/addrbook.go      |  2 +-
 p2p/peer_set.go      |  5 ++--
 p2p/peer_set_test.go | 63 ++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 67 insertions(+), 3 deletions(-)
 create mode 100644 p2p/peer_set_test.go

diff --git a/p2p/addrbook.go b/p2p/addrbook.go
index dcb41d82f..0fc61b89c 100644
--- a/p2p/addrbook.go
+++ b/p2p/addrbook.go
@@ -334,7 +334,7 @@ func (a *AddrBook) loadFromFile(filePath string) {
 	// If doesn't exist, do nothing.
 	_, err := os.Stat(filePath)
 	if os.IsNotExist(err) {
-		return
+		panic(Fmt("File does not exist: %v", filePath))
 	}
 
 	// Load addrBookJSON{}
diff --git a/p2p/peer_set.go b/p2p/peer_set.go
index 158ffe542..b4230ffa3 100644
--- a/p2p/peer_set.go
+++ b/p2p/peer_set.go
@@ -70,12 +70,13 @@ func (ps *PeerSet) Remove(peer *Peer) {
 	// If it's the last peer, that's an easy special case.
 	if index == len(ps.list)-1 {
 		ps.list = newList
+		delete(ps.lookup, peer.Key)
 		return
 	}
 	// Move the last item from ps.list to "index" in list.
 	lastPeer := ps.list[len(ps.list)-1]
-	lastPeerAddr := lastPeer.mconn.RemoteAddress.String()
-	lastPeerItem := ps.lookup[lastPeerAddr]
+	lastPeerKey := lastPeer.Key
+	lastPeerItem := ps.lookup[lastPeerKey]
 	newList[index] = lastPeer
 	lastPeerItem.index = index
 	ps.list = newList
diff --git a/p2p/peer_set_test.go b/p2p/peer_set_test.go
new file mode 100644
index 000000000..6a0c38ae2
--- /dev/null
+++ b/p2p/peer_set_test.go
@@ -0,0 +1,63 @@
+package p2p
+
+import (
+	"math/rand"
+	"testing"
+
+	. "github.com/tendermint/tendermint/common"
+)
+
+// Returns an empty dummy peer
+func randPeer() *Peer {
+	return &Peer{
+		Key: Fmt("%v.%v.%v.%v:%v", rand.Int()%256, rand.Int()%256, rand.Int()%256, rand.Int()%256, rand.Int()%10000+80),
+	}
+}
+
+func TestAddRemoveOne(t *testing.T) {
+	peerSet := NewPeerSet()
+
+	peer := randPeer()
+	added := peerSet.Add(peer)
+	if !added {
+		t.Errorf("Failed to add new peer")
+	}
+	if peerSet.Size() != 1 {
+		t.Errorf("Failed to add new peer and increment size")
+	}
+
+	peerSet.Remove(peer)
+	if peerSet.Has(peer.Key) {
+		t.Errorf("Failed to remove peer")
+	}
+	if peerSet.Size() != 0 {
+		t.Errorf("Failed to remove peer and decrement size")
+	}
+}
+
+func TestAddRemoveMany(t *testing.T) {
+	peerSet := NewPeerSet()
+
+	peers := []*Peer{}
+	for i := 0; i < 100; i++ {
+		peer := randPeer()
+		added := peerSet.Add(peer)
+		if !added {
+			t.Errorf("Failed to add new peer")
+		}
+		if peerSet.Size() != i+1 {
+			t.Errorf("Failed to add new peer and increment size")
+		}
+		peers = append(peers, peer)
+	}
+
+	for i, peer := range peers {
+		peerSet.Remove(peer)
+		if peerSet.Has(peer.Key) {
+			t.Errorf("Failed to remove peer")
+		}
+		if peerSet.Size() != len(peers)-i-1 {
+			t.Errorf("Failed to remove peer and decrement size")
+		}
+	}
+}