Merge pull request #12 from tendermint/bugfix/pex-issues-335

PEX issues #335

addrbook.go

@@ -153,6 +153,7 @@ func (a *AddrBook) OurAddresses() []*NetAddress {
 	return addrs
 }
 
+// NOTE: addr must not be nil
 func (a *AddrBook) AddAddress(addr *NetAddress, src *NetAddress) {
 	a.mtx.Lock()
 	defer a.mtx.Unlock()
@@ -368,6 +369,12 @@ func (a *AddrBook) loadFromFile(filePath string) bool {
 	return true
 }
 
+// Save saves the book.
+func (a *AddrBook) Save() {
+	log.Info("Saving AddrBook to file", "size", a.Size())
+	a.saveToFile(a.filePath)
+}
+
 /* Private methods */
 
 func (a *AddrBook) saveRoutine() {

addrbook_test.go

@@ -9,8 +9,6 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-const addrBookStrict = true
-
 func createTempFileName(prefix string) string {
 	f, err := ioutil.TempFile("", prefix)
 	if err != nil {
@@ -24,143 +22,129 @@ func createTempFileName(prefix string) string {
 	return fname
 }
 
-func TestEmpty(t *testing.T) {
+func TestAddrBookSaveLoad(t *testing.T) {
 	fname := createTempFileName("addrbook_test")
-	// t.Logf("New tempfile name: %v", fname)
 
-	// Save an empty book & load it
-	book := NewAddrBook(fname, addrBookStrict)
+	// 0 addresses
+	book := NewAddrBook(fname, true)
 	book.saveToFile(fname)
 
-	book = NewAddrBook(fname, addrBookStrict)
+	book = NewAddrBook(fname, true)
 	book.loadFromFile(fname)
 
-	if book.Size() != 0 {
-		t.Errorf("Expected 0 addresses, found %v", book.Size())
-	}
-}
+	assert.Zero(t, book.Size())
 
-func randIPv4Address(t *testing.T) *NetAddress {
-	for {
-		ip := fmt.Sprintf("%v.%v.%v.%v",
-			rand.Intn(254)+1,
-			rand.Intn(255),
-			rand.Intn(255),
-			rand.Intn(255),
-		)
-		port := rand.Intn(65535-1) + 1
-		addr, err := NewNetAddressString(fmt.Sprintf("%v:%v", ip, port))
-		assert.Nil(t, err, "error generating rand network address")
-		if addr.Routable() {
-			return addr
-		}
-	}
-}
+	// 100 addresses
+	randAddrs := randNetAddressPairs(t, 100)
 
-func TestSaveAddresses(t *testing.T) {
-	fname := createTempFileName("addrbook_test")
-	//t.Logf("New tempfile name: %v", fname)
-
-	// Create some random addresses
-	randAddrs := []struct {
-		addr *NetAddress
-		src  *NetAddress
-	}{}
-	for i := 0; i < 100; i++ {
-		addr := randIPv4Address(t)
-		src := randIPv4Address(t)
-		randAddrs = append(randAddrs, struct {
-			addr *NetAddress
-			src  *NetAddress
-		}{
-			addr: addr,
-			src:  src,
-		})
-	}
-
-	// Create the book & populate & save
-	book := NewAddrBook(fname, addrBookStrict)
 	for _, addrSrc := range randAddrs {
 		book.AddAddress(addrSrc.addr, addrSrc.src)
 	}
-	if book.Size() != 100 {
-		t.Errorf("Expected 100 addresses, found %v", book.Size())
-	}
+
+	assert.Equal(t, 100, book.Size())
 	book.saveToFile(fname)
 
-	// Reload the book
-	book = NewAddrBook(fname, addrBookStrict)
+	book = NewAddrBook(fname, true)
 	book.loadFromFile(fname)
 
-	// Test ...
+	assert.Equal(t, 100, book.Size())
+}
 
-	if book.Size() != 100 {
-		t.Errorf("Expected 100 addresses, found %v", book.Size())
-	}
+func TestAddrBookLookup(t *testing.T) {
+	fname := createTempFileName("addrbook_test")
+
+	randAddrs := randNetAddressPairs(t, 100)
 
+	book := NewAddrBook(fname, true)
 	for _, addrSrc := range randAddrs {
 		addr := addrSrc.addr
 		src := addrSrc.src
+		book.AddAddress(addr, src)
+
 		ka := book.addrLookup[addr.String()]
-		if ka == nil {
-			t.Fatalf("Expected to find KnownAddress %v but wasn't there.", addr)
-		}
+		assert.NotNil(t, ka, "Expected to find KnownAddress %v but wasn't there.", addr)
+
 		if !(ka.Addr.Equals(addr) && ka.Src.Equals(src)) {
 			t.Fatalf("KnownAddress doesn't match addr & src")
 		}
 	}
 }
 
-func TestPromoteToOld(t *testing.T) {
+func TestAddrBookPromoteToOld(t *testing.T) {
 	fname := createTempFileName("addrbook_test")
-	t.Logf("New tempfile name: %v", fname)
-
-	// Create some random addresses
-	randAddrs := []struct {
-		addr *NetAddress
-		src  *NetAddress
-	}{}
-	for i := 0; i < 100; i++ {
-		addr := randIPv4Address(t)
-		src := randIPv4Address(t)
-		randAddrs = append(randAddrs, struct {
-			addr *NetAddress
-			src  *NetAddress
-		}{
-			addr: addr,
-			src:  src,
-		})
-	}
 
-	// Create the book & populate & save
-	book := NewAddrBook(fname, addrBookStrict)
+	randAddrs := randNetAddressPairs(t, 100)
+
+	book := NewAddrBook(fname, true)
 	for _, addrSrc := range randAddrs {
 		book.AddAddress(addrSrc.addr, addrSrc.src)
 	}
+
 	// Attempt all addresses.
 	for _, addrSrc := range randAddrs {
 		book.MarkAttempt(addrSrc.addr)
 	}
+
 	// Promote half of them
 	for i, addrSrc := range randAddrs {
 		if i%2 == 0 {
 			book.MarkGood(addrSrc.addr)
 		}
 	}
-	book.saveToFile(fname)
 
-	// Reload the book
-	book = NewAddrBook(fname, addrBookStrict)
-	book.loadFromFile(fname)
+	// TODO: do more testing :)
 
-	// Test ...
+	selection := book.GetSelection()
+	t.Logf("selection: %v", selection)
 
-	if book.Size() != 100 {
-		t.Errorf("Expected 100 addresses, found %v", book.Size())
+	if len(selection) > book.Size() {
+		t.Errorf("selection could not be bigger than the book")
 	}
+}
 
-	// TODO: do more testing :)
+func TestAddrBookHandlesDuplicates(t *testing.T) {
+	fname := createTempFileName("addrbook_test")
 
-	selection := book.GetSelection()
-	t.Logf("selection: %v", selection)
+	book := NewAddrBook(fname, true)
+
+	randAddrs := randNetAddressPairs(t, 100)
+
+	differentSrc := randIPv4Address(t)
+	for _, addrSrc := range randAddrs {
+		book.AddAddress(addrSrc.addr, addrSrc.src)
+		book.AddAddress(addrSrc.addr, addrSrc.src)  // duplicate
+		book.AddAddress(addrSrc.addr, differentSrc) // different src
+	}
+
+	assert.Equal(t, 100, book.Size())
+}
+
+type netAddressPair struct {
+	addr *NetAddress
+	src  *NetAddress
+}
+
+func randNetAddressPairs(t *testing.T, n int) []netAddressPair {
+	randAddrs := make([]netAddressPair, n)
+	for i := 0; i < n; i++ {
+		randAddrs[i] = netAddressPair{addr: randIPv4Address(t), src: randIPv4Address(t)}
+	}
+	return randAddrs
+}
+
+func randIPv4Address(t *testing.T) *NetAddress {
+	for {
+		ip := fmt.Sprintf("%v.%v.%v.%v",
+			rand.Intn(254)+1,
+			rand.Intn(255),
+			rand.Intn(255),
+			rand.Intn(255),
+		)
+		port := rand.Intn(65535-1) + 1
+		addr, err := NewNetAddressString(fmt.Sprintf("%v:%v", ip, port))
+		assert.Nil(t, err, "error generating rand network address")
+		if addr.Routable() {
+			return addr
+		}
+	}
 }

pex_reactor.go

@@ -42,12 +42,14 @@ func NewPEXReactor(book *AddrBook) *PEXReactor {
 
 func (pexR *PEXReactor) OnStart() error {
 	pexR.BaseReactor.OnStart()
+	pexR.book.OnStart()
 	go pexR.ensurePeersRoutine()
 	return nil
 }
 
 func (pexR *PEXReactor) OnStop() {
 	pexR.BaseReactor.OnStop()
+	pexR.book.OnStop()
 }
 
 // Implements Reactor
@@ -110,7 +112,9 @@ func (pexR *PEXReactor) Receive(chID byte, src *Peer, msgBytes []byte) {
 		// (We don't want to get spammed with bad peers)
 		srcAddr := src.Connection().RemoteAddress
 		for _, addr := range msg.Addrs {
-			pexR.book.AddAddress(addr, srcAddr)
+			if addr != nil {
+				pexR.book.AddAddress(addr, srcAddr)
+			}
 		}
 	default:
 		log.Warn(Fmt("Unknown message type %v", reflect.TypeOf(msg)))