diff --git a/addrbook.go b/addrbook.go index 2bbae64f..ea438120 100644 --- a/addrbook.go +++ b/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() { diff --git a/addrbook_test.go b/addrbook_test.go index 5eb7c0b6..7e8cb8d7 100644 --- a/addrbook_test.go +++ b/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,20 +22,114 @@ 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()) + + // 100 addresses + randAddrs := randNetAddressPairs(t, 100) + + for _, addrSrc := range randAddrs { + book.AddAddress(addrSrc.addr, addrSrc.src) } + + assert.Equal(t, 100, book.Size()) + book.saveToFile(fname) + + book = NewAddrBook(fname, true) + book.loadFromFile(fname) + + assert.Equal(t, 100, 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()] + 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 TestAddrBookPromoteToOld(t *testing.T) { + fname := createTempFileName("addrbook_test") + + 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) + } + } + + // TODO: do more testing :) + + selection := book.GetSelection() + t.Logf("selection: %v", selection) + + if len(selection) > book.Size() { + t.Errorf("selection could not be bigger than the book") + } +} + +func TestAddrBookHandlesDuplicates(t *testing.T) { + fname := createTempFileName("addrbook_test") + + 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 { @@ -56,111 +148,3 @@ func randIPv4Address(t *testing.T) *NetAddress { } } } - -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()) - } - book.saveToFile(fname) - - // Reload the book - book = NewAddrBook(fname, addrBookStrict) - book.loadFromFile(fname) - - // Test ... - - if book.Size() != 100 { - t.Errorf("Expected 100 addresses, found %v", book.Size()) - } - - for _, addrSrc := range randAddrs { - addr := addrSrc.addr - src := addrSrc.src - ka := book.addrLookup[addr.String()] - if ka == nil { - t.Fatalf("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) { - 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) - } - // 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) - - // Test ... - - if book.Size() != 100 { - t.Errorf("Expected 100 addresses, found %v", book.Size()) - } - - // TODO: do more testing :) - - selection := book.GetSelection() - t.Logf("selection: %v", selection) -} diff --git a/pex_reactor.go b/pex_reactor.go index c02d9a67..b8cb43f8 100644 --- a/pex_reactor.go +++ b/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)))