diff --git a/pstoremem/addr_book.go b/pstoremem/addr_book.go index 7446bd7..0826a83 100644 --- a/pstoremem/addr_book.go +++ b/pstoremem/addr_book.go @@ -25,7 +25,7 @@ type expiringAddr struct { } func (e *expiringAddr) ExpiredBy(t time.Time) bool { - return t.After(e.Expires) + return !t.Before(e.Expires) } type peerRecordState struct { @@ -315,9 +315,15 @@ func (mab *memoryAddrBook) UpdateAddrs(p peer.ID, oldTTL time.Duration, newTTL t defer s.Unlock() exp := time.Now().Add(newTTL) amap, found := s.addrs[p] - if found { - for k, a := range amap { - if oldTTL == a.TTL { + if !found { + return + } + + for k, a := range amap { + if oldTTL == a.TTL { + if newTTL == 0 { + delete(amap, k) + } else { a.TTL = newTTL a.Expires = exp amap[k] = a diff --git a/test/addr_book_suite.go b/test/addr_book_suite.go index 6bc9dd8..2c09262 100644 --- a/test/addr_book_suite.go +++ b/test/addr_book_suite.go @@ -210,6 +210,18 @@ func testUpdateTTLs(m pstore.AddrBook) func(t *testing.T) { m.UpdateAddrs(id, time.Hour, time.Minute) }) + t.Run("update to 0 clears addrs", func(t *testing.T) { + id := GeneratePeerIDs(1)[0] + addrs := GenerateAddrs(1) + + // Shouldn't panic. + m.SetAddrs(id, addrs, time.Hour) + m.UpdateAddrs(id, time.Hour, 0) + if len(m.Addrs(id)) != 0 { + t.Error("expected no addresses") + } + }) + t.Run("update ttls successfully", func(t *testing.T) { ids := GeneratePeerIDs(2) addrs1, addrs2 := GenerateAddrs(2), GenerateAddrs(2)