Address Book POC implementation

libp2p/peerstore.nim

@@ -0,0 +1,142 @@
+## Nim-LibP2P
+## Copyright (c) 2021 Status Research & Development GmbH
+## Licensed under either of
+##  * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE))
+##  * MIT license ([LICENSE-MIT](LICENSE-MIT))
+## at your option.
+## This file may not be copied, modified, or distributed except according to
+## those terms.
+
+import
+  std/[tables, sets, sequtils],
+  ./peerid,
+  ./multiaddress
+
+type
+  ##############################
+  # Listener and handler types #
+  ##############################
+
+  AddrChangeHandler* = proc(peerId: PeerID, multiaddrs: HashSet[MultiAddress])
+
+  EventListener* = object
+    # Listener object with client-defined handlers for peer store events
+    addrChange*: AddrChangeHandler
+    # @TODO add handlers for other event types
+
+  #########
+  # Books #
+  #########
+
+  # Each book contains a book (map) and event handler(s)
+  AddressBook* = object
+    book*: Table[PeerID, HashSet[MultiAddress]]
+    addrChange: AddrChangeHandler
+
+  ####################
+  # Peer store types #
+  ####################
+
+  PeerStore* = ref object of RootObj
+    addressBook*: AddressBook
+    listeners: seq[EventListener]
+
+  StoredInfo* = object
+    # Collates stored info about a peer
+    ## @TODO include data from other stores once added
+    peerId*: PeerID
+    addrs*: HashSet[MultiAddress]
+
+proc init(T: type AddressBook, addrChange: AddrChangeHandler): AddressBook =
+  T(book: initTable[PeerId, HashSet[MultiAddress]](),
+    addrChange: addrChange)
+
+proc init*(p: PeerStore) =
+  p.listeners = newSeq[EventListener]()
+
+  proc addrChange(peerId: PeerID, multiaddrs: HashSet[MultiAddress]) =
+    # Notify all listeners of change in multiaddr
+    for listener in p.listeners:
+      listener.addrChange(peerId, multiaddrs)
+
+  p.addressBook = AddressBook.init(addrChange)
+
+proc init*(T: type PeerStore): PeerStore =
+  var p: PeerStore
+  new(p)
+  p.init()
+  return p
+
+####################
+# Address Book API #
+####################
+
+proc get*(addressBook: AddressBook,
+          peerId: PeerID): HashSet[MultiAddress] =
+  ## Get the known addresses of a provided peer.
+
+  addressBook.book.getOrDefault(peerId,
+                                initHashSet[MultiAddress]())
+
+proc add*(addressBook: var AddressBook,
+          peerId: PeerID,
+          multiaddr: MultiAddress) = 
+  ## Add known multiaddr of a given peer. If the peer is not known, 
+  ## it will be set with the provided multiaddr.
+
+  addressBook.book.mgetOrPut(peerId,
+                             initHashSet[MultiAddress]()).incl(multiaddr)
+  addressBook.addrChange(peerId, addressBook.get(peerId)) # Notify clients
+
+proc delete*(addressBook: var AddressBook,
+             peerId: PeerID): bool =
+  ## Delete the provided peer from the book.
+
+  if not addressBook.book.hasKey(peerId):
+    return false
+  else:
+    addressBook.book.del(peerId)
+    return true
+
+proc set*(addressBook: var AddressBook,
+          peerId: PeerID,
+          addrs: HashSet[MultiAddress]) =
+  ## Set known multiaddresses for a given peer. This will replace previously
+  ## stored addresses. Replacing stored multiaddresses might
+  ## result in losing obtained certified addresses, which is not desirable.
+  ## Consider using addressBook.add() as alternative.
+
+  addressBook.book[peerId] = addrs
+  addressBook.addrChange(peerId, addressBook.get(peerId)) # Notify clients
+
+##################  
+# Peer Store API #
+##################
+
+proc addListener*(peerStore: PeerStore,
+                  listener: EventListener) =
+  ## Register event listener to notify clients of changes in the peer store
+
+  peerStore.listeners.add(listener)
+
+proc delete*(peerStore: PeerStore,
+             peerId: PeerID): bool =
+  ## Delete the provided peer from every book.
+
+  peerStore.addressBook.delete(peerId)
+
+proc get*(peerStore: PeerStore,
+          peerId: PeerID): StoredInfo =
+  ## Get the stored information of a given peer.
+
+  StoredInfo(
+    peerId: peerId,
+    addrs: peerStore.addressBook.get(peerId)
+  )
+
+proc peers*(peerStore: PeerStore): seq[StoredInfo] =
+  ## Get all the stored information of every peer.
+
+  let allKeys = toSeq(keys(peerStore.addressBook.book)) # @TODO concat keys from other books
+
+  return allKeys.mapIt(peerStore.get(it))

tests/testpeerstore.nim

@@ -0,0 +1,114 @@
+import
+  std/[unittest, tables, sequtils, sets],
+  ../libp2p/crypto/crypto,
+  ../libp2p/multiaddress,
+  ../libp2p/peerid,
+  ../libp2p/peerstore,
+  ./helpers
+
+suite "PeerStore":
+  # Testvars
+  let
+    # Peer 1
+    seckey1 = PrivateKey.random(ECDSA, rng[]).get()
+    peerId1 = PeerID.init(seckey1).get()
+    multiaddrStr1 = "/ip4/127.0.0.1/udp/1234/p2p/QmcgpsyWgH8Y8ajJz1Cu72KnS5uo2Aa2LpzU7kinSupNKC"
+    multiaddr1 = MultiAddress.init(multiaddrStr1).get()
+    # Peer 2
+    seckey2 = PrivateKey.random(ECDSA, rng[]).get()
+    peerId2 = PeerID.init(seckey2).get()
+    multiaddrStr2 = "/ip4/0.0.0.0/tcp/1234/ipfs/QmcgpsyWgH8Y8ajJz1Cu72KnS5uo2Aa2LpzU7kinSupNKC"
+    multiaddr2 = MultiAddress.init(multiaddrStr2).get()
+
+  test "PeerStore API":
+    # Set up peer store
+    var
+      peerStore = PeerStore.init()
+
+    peerStore.addressBook.add(peerId1, multiaddr1)
+    peerStore.addressBook.add(peerId2, multiaddr2)
+
+    # Test PeerStore::get
+    let
+      peer1Stored = peerStore.get(peerId1)
+      peer2Stored = peerStore.get(peerId2)
+    check:
+      peer1Stored.peerId == peerId1
+      peer1Stored.addrs == toHashSet([multiaddr1])
+      peer2Stored.peerId == peerId2
+      peer2Stored.addrs == toHashSet([multiaddr2])
+
+    # Test PeerStore::peers
+    let peers = peerStore.peers()
+    check:
+      peers.len == 2
+      peers.anyIt(it.peerId == peerId1 and it.addrs == toHashSet([multiaddr1]))
+      peers.anyIt(it.peerId == peerId2 and it.addrs == toHashSet([multiaddr2]))
+
+    # Test PeerStore::delete
+    check:
+      # Delete existing peerId
+      peerStore.delete(peerId1) == true
+      peerStore.peers().anyIt(it.peerId == peerId1) == false
+
+      # Now try and delete it again
+      peerStore.delete(peerId1) == false
+
+  test "PeerStore listeners":
+    # Set up peer store with listener
+    var
+      peerStore = PeerStore.init()
+      addrChanged = false
+
+    proc addrChange(peerId: PeerID, addrs: HashSet[MultiAddress]) =
+      addrChanged = true
+
+    let listener = EventListener(addrChange: addrChange)
+
+    peerStore.addListener(listener)
+
+    # Test listener triggered on adding multiaddr
+    peerStore.addressBook.add(peerId1, multiaddr1)
+    check:
+      addrChanged == true
+
+    # Test listener triggered on setting addresses
+    addrChanged = false
+    peerStore.addressBook.set(peerId2,
+                              toHashSet([multiaddr1, multiaddr2]))
+    check:
+      addrChanged == true    
+
+  test "AddressBook API":
+    # Set up address book
+    var
+      addressBook = PeerStore.init().addressBook
+
+    # Test AddressBook::add
+    addressBook.add(peerId1, multiaddr1)
+
+    check:
+      toSeq(keys(addressBook.book))[0] == peerId1
+      toSeq(values(addressBook.book))[0] == toHashSet([multiaddr1])
+
+    # Test AddressBook::get
+    check:
+      addressBook.get(peerId1) == toHashSet([multiaddr1])
+
+    # Test AddressBook::delete
+    check:
+      # Try to delete peerId that doesn't exist
+      addressBook.delete(peerId2) == false
+
+      # Delete existing peerId
+      addressBook.book.len == 1 # sanity
+      addressBook.delete(peerId1) == true
+      addressBook.book.len == 0
+
+    # Test AddressBook::set
+    # Set peerId2 with multiple multiaddrs
+    addressBook.set(peerId2,
+                toHashSet([multiaddr1, multiaddr2]))
+    check:
+      toSeq(keys(addressBook.book))[0] == peerId2
+      toSeq(values(addressBook.book))[0] == toHashSet([multiaddr1, multiaddr2])