diff --git a/addrmgr/addrmanager.go b/addrmgr/addrmanager.go index fc693068..89569aa0 100644 --- a/addrmgr/addrmanager.go +++ b/addrmgr/addrmanager.go @@ -635,21 +635,9 @@ func (a *AddrManager) NeedMoreAddresses() bool { // AddressCache returns the current address cache. It must be treated as // read-only (but since it is a copy now, this is not as dangerous). func (a *AddrManager) AddressCache() []*wire.NetAddress { - a.mtx.Lock() - defer a.mtx.Unlock() + allAddr := a.getAddresses() - addrIndexLen := len(a.addrIndex) - if addrIndexLen == 0 { - return nil - } - - allAddr := make([]*wire.NetAddress, 0, addrIndexLen) - // Iteration order is undefined here, but we randomise it anyway. - for _, v := range a.addrIndex { - allAddr = append(allAddr, v.na) - } - - numAddresses := addrIndexLen * getAddrPercent / 100 + numAddresses := len(allAddr) * getAddrPercent / 100 if numAddresses > getAddrMax { numAddresses = getAddrMax } @@ -658,7 +646,7 @@ func (a *AddrManager) AddressCache() []*wire.NetAddress { // `numAddresses' since we are throwing the rest. for i := 0; i < numAddresses; i++ { // pick a number between current index and the end - j := rand.Intn(addrIndexLen-i) + i + j := rand.Intn(len(allAddr)-i) + i allAddr[i], allAddr[j] = allAddr[j], allAddr[i] } @@ -666,6 +654,25 @@ func (a *AddrManager) AddressCache() []*wire.NetAddress { return allAddr[0:numAddresses] } +// getAddresses returns all of the addresses currently found within the +// manager's address cache. +func (a *AddrManager) getAddresses() []*wire.NetAddress { + a.mtx.Lock() + defer a.mtx.Unlock() + + addrIndexLen := len(a.addrIndex) + if addrIndexLen == 0 { + return nil + } + + addrs := make([]*wire.NetAddress, 0, addrIndexLen) + for _, v := range a.addrIndex { + addrs = append(addrs, v.na) + } + + return addrs +} + // reset resets the address manager by reinitialising the random source // and allocating fresh empty bucket storage. func (a *AddrManager) reset() { diff --git a/addrmgr/addrmanager_internal_test.go b/addrmgr/addrmanager_internal_test.go new file mode 100644 index 00000000..82036d37 --- /dev/null +++ b/addrmgr/addrmanager_internal_test.go @@ -0,0 +1,119 @@ +package addrmgr + +import ( + "io/ioutil" + "math/rand" + "net" + "os" + "testing" + + "github.com/btcsuite/btcd/wire" +) + +// randAddr generates a *wire.NetAddress backed by a random IPv4/IPv6 address. +func randAddr(t *testing.T) *wire.NetAddress { + t.Helper() + + ipv4 := rand.Intn(2) == 0 + var ip net.IP + if ipv4 { + var b [4]byte + if _, err := rand.Read(b[:]); err != nil { + t.Fatal(err) + } + ip = b[:] + } else { + var b [16]byte + if _, err := rand.Read(b[:]); err != nil { + t.Fatal(err) + } + ip = b[:] + } + + return &wire.NetAddress{ + Services: wire.ServiceFlag(rand.Uint64()), + IP: ip, + Port: uint16(rand.Uint32()), + } +} + +// assertAddr ensures that the two addresses match. The timestamp is not +// checked as it does not affect uniquely identifying a specific address. +func assertAddr(t *testing.T, got, expected *wire.NetAddress) { + if got.Services != expected.Services { + t.Fatalf("expected address services %v, got %v", + expected.Services, got.Services) + } + if !got.IP.Equal(expected.IP) { + t.Fatalf("expected address IP %v, got %v", expected.IP, got.IP) + } + if got.Port != expected.Port { + t.Fatalf("expected address port %d, got %d", expected.Port, + got.Port) + } +} + +// assertAddrs ensures that the manager's address cache matches the given +// expected addresses. +func assertAddrs(t *testing.T, addrMgr *AddrManager, + expectedAddrs map[string]*wire.NetAddress) { + + t.Helper() + + addrs := addrMgr.getAddresses() + + if len(addrs) != len(expectedAddrs) { + t.Fatalf("expected to find %d addresses, found %d", + len(expectedAddrs), len(addrs)) + } + + for _, addr := range addrs { + addrStr := NetAddressKey(addr) + expectedAddr, ok := expectedAddrs[addrStr] + if !ok { + t.Fatalf("expected to find address %v", addrStr) + } + + assertAddr(t, addr, expectedAddr) + } +} + +// TestAddrManagerSerialization ensures that we can properly serialize and +// deserialize the manager's current address cache. +func TestAddrManagerSerialization(t *testing.T) { + t.Parallel() + + // We'll start by creating our address manager backed by a temporary + // directory. + tempDir, err := ioutil.TempDir("", "addrmgr") + if err != nil { + t.Fatalf("unable to create temp dir: %v", err) + } + defer os.RemoveAll(tempDir) + + addrMgr := New(tempDir, nil) + + // We'll be adding 5 random addresses to the manager. + const numAddrs = 5 + + expectedAddrs := make(map[string]*wire.NetAddress, numAddrs) + for i := 0; i < numAddrs; i++ { + addr := randAddr(t) + expectedAddrs[NetAddressKey(addr)] = addr + addrMgr.AddAddress(addr, randAddr(t)) + } + + // Now that the addresses have been added, we should be able to retrieve + // them. + assertAddrs(t, addrMgr, expectedAddrs) + + // Then, we'll persist these addresses to disk and restart the address + // manager. + addrMgr.savePeers() + addrMgr = New(tempDir, nil) + + // Finally, we'll read all of the addresses from disk and ensure they + // match as expected. + addrMgr.loadPeers() + assertAddrs(t, addrMgr, expectedAddrs) +}