diff --git a/libp2p/peerstore.nim b/libp2p/peerstore.nim new file mode 100644 index 0000000..85787e8 --- /dev/null +++ b/libp2p/peerstore.nim @@ -0,0 +1,170 @@ +## 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. + +{.push raises: [Defect].} + +import + std/[tables, sets, sequtils], + ./crypto/crypto, + ./peerid, + ./multiaddress + +type + ################# + # Handler types # + ################# + + PeerBookChangeHandler*[T] = proc(peerId: PeerID, entry: T) + + AddrChangeHandler* = PeerBookChangeHandler[HashSet[MultiAddress]] + ProtoChangeHandler* = PeerBookChangeHandler[HashSet[string]] + KeyChangeHandler* = PeerBookChangeHandler[PublicKey] + + ######### + # Books # + ######### + + # Each book contains a book (map) and event handler(s) + PeerBook*[T] = object of RootObj + book*: Table[PeerID, T] + changeHandlers: seq[PeerBookChangeHandler[T]] + + AddressBook* = object of PeerBook[HashSet[MultiAddress]] + ProtoBook* = object of PeerBook[HashSet[string]] + KeyBook* = object of PeerBook[PublicKey] + + #################### + # Peer store types # + #################### + + PeerStore* = ref object of RootObj + addressBook*: AddressBook + protoBook*: ProtoBook + keyBook*: KeyBook + + StoredInfo* = object + # Collates stored info about a peer + peerId*: PeerID + addrs*: HashSet[MultiAddress] + protos*: HashSet[string] + publicKey*: PublicKey + +## Constructs a new PeerStore with metadata of type M +proc new*(T: type PeerStore): PeerStore = + var p: PeerStore + new(p) + return p + +######################### +# Generic Peer Book API # +######################### + +proc get*[T](peerBook: PeerBook[T], + peerId: PeerID): T = + ## Get all the known metadata of a provided peer. + + peerBook.book.getOrDefault(peerId) + +proc set*[T](peerBook: var PeerBook[T], + peerId: PeerID, + entry: T) = + ## Set metadata for a given peerId. This will replace any + ## previously stored metadata. + + peerBook.book[peerId] = entry + + # Notify clients + for handler in peerBook.changeHandlers: + handler(peerId, peerBook.get(peerId)) + +proc delete*[T](peerBook: var PeerBook[T], + peerId: PeerID): bool = + ## Delete the provided peer from the book. + + if not peerBook.book.hasKey(peerId): + return false + else: + peerBook.book.del(peerId) + return true + +#################### +# Address Book API # +#################### + +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) + + # Notify clients + for handler in addressBook.changeHandlers: + handler(peerId, addressBook.get(peerId)) + +##################### +# Protocol Book API # +##################### + +proc add*(protoBook: var ProtoBook, + peerId: PeerID, + protocol: string) = + ## Adds known protocol codec for a given peer. If the peer is not known, + ## it will be set with the provided protocol. + + protoBook.book.mgetOrPut(peerId, + initHashSet[string]()).incl(protocol) + + # Notify clients + for handler in protoBook.changeHandlers: + handler(peerId, protoBook.get(peerId)) + +################## +# Peer Store API # +################## + +proc addHandlers*(peerStore: PeerStore, + addrChangeHandler: AddrChangeHandler, + protoChangeHandler: ProtoChangeHandler, + keyChangeHandler: KeyChangeHandler) = + ## Register event handlers to notify clients of changes in the peer store + + peerStore.addressBook.changeHandlers.add(addrChangeHandler) + peerStore.protoBook.changeHandlers.add(protoChangeHandler) + peerStore.keyBook.changeHandlers.add(keyChangeHandler) + +proc delete*(peerStore: PeerStore, + peerId: PeerID): bool = + ## Delete the provided peer from every book. + + peerStore.addressBook.delete(peerId) and + peerStore.protoBook.delete(peerId) and + peerStore.keyBook.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), + protos: peerStore.protoBook.get(peerId), + publicKey: peerStore.keyBook.get(peerId) + ) + +proc peers*(peerStore: PeerStore): seq[StoredInfo] = + ## Get all the stored information of every peer. + + let allKeys = concat(toSeq(keys(peerStore.addressBook.book)), + toSeq(keys(peerStore.protoBook.book)), + toSeq(keys(peerStore.keyBook.book))).toHashSet() + + return allKeys.mapIt(peerStore.get(it)) diff --git a/tests/testpeerstore.nim b/tests/testpeerstore.nim new file mode 100644 index 0000000..7eba389 --- /dev/null +++ b/tests/testpeerstore.nim @@ -0,0 +1,223 @@ +import + std/[unittest, tables, sequtils, sets], + ../libp2p/crypto/crypto, + ../libp2p/multiaddress, + ../libp2p/peerid, + ../libp2p/peerstore, + ./helpers + +suite "PeerStore": + # Testvars + let + # Peer 1 + keyPair1 = KeyPair.random(ECDSA, rng[]).get() + peerId1 = PeerID.init(keyPair1.secKey).get() + multiaddrStr1 = "/ip4/127.0.0.1/udp/1234/p2p/QmcgpsyWgH8Y8ajJz1Cu72KnS5uo2Aa2LpzU7kinSupNKC" + multiaddr1 = MultiAddress.init(multiaddrStr1).get() + testcodec1 = "/nim/libp2p/test/0.0.1-beta1" + # Peer 2 + keyPair2 = KeyPair.random(ECDSA, rng[]).get() + peerId2 = PeerID.init(keyPair2.secKey).get() + multiaddrStr2 = "/ip4/0.0.0.0/tcp/1234/ipfs/QmcgpsyWgH8Y8ajJz1Cu72KnS5uo2Aa2LpzU7kinSupNKC" + multiaddr2 = MultiAddress.init(multiaddrStr2).get() + testcodec2 = "/nim/libp2p/test/0.0.2-beta1" + + test "PeerStore API": + # Set up peer store + var + peerStore = PeerStore.new() + + peerStore.addressBook.add(peerId1, multiaddr1) + peerStore.addressBook.add(peerId2, multiaddr2) + peerStore.protoBook.add(peerId1, testcodec1) + peerStore.protoBook.add(peerId2, testcodec2) + peerStore.keyBook.set(peerId1, keyPair1.pubKey) + peerStore.keyBook.set(peerId2, keyPair2.pubKey) + + # Test PeerStore::get + let + peer1Stored = peerStore.get(peerId1) + peer2Stored = peerStore.get(peerId2) + check: + peer1Stored.peerId == peerId1 + peer1Stored.addrs == toHashSet([multiaddr1]) + peer1Stored.protos == toHashSet([testcodec1]) + peer1Stored.publicKey == keyPair1.pubkey + peer2Stored.peerId == peerId2 + peer2Stored.addrs == toHashSet([multiaddr2]) + peer2Stored.protos == toHashSet([testcodec2]) + peer2Stored.publicKey == keyPair2.pubkey + + # Test PeerStore::peers + let peers = peerStore.peers() + check: + peers.len == 2 + peers.anyIt(it.peerId == peerId1 and + it.addrs == toHashSet([multiaddr1]) and + it.protos == toHashSet([testcodec1]) and + it.publicKey == keyPair1.pubkey) + peers.anyIt(it.peerId == peerId2 and + it.addrs == toHashSet([multiaddr2]) and + it.protos == toHashSet([testcodec2]) and + it.publicKey == keyPair2.pubkey) + + # 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.new() + addrChanged = false + protoChanged = false + keyChanged = false + + proc addrChange(peerId: PeerID, addrs: HashSet[MultiAddress]) = + addrChanged = true + + proc protoChange(peerId: PeerID, protos: HashSet[string]) = + protoChanged = true + + proc keyChange(peerId: PeerID, publicKey: PublicKey) = + keyChanged = true + + peerStore.addHandlers(addrChangeHandler = addrChange, + protoChangeHandler = protoChange, + keyChangeHandler = keyChange) + + # 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 listener triggered on adding proto + peerStore.protoBook.add(peerId1, testcodec1) + check: + protoChanged == true + + # Test listener triggered on setting protos + protoChanged = false + peerStore.protoBook.set(peerId2, + toHashSet([testcodec1, testcodec2])) + check: + protoChanged == true + + # Test listener triggered on setting public key + peerStore.keyBook.set(peerId1, + keyPair1.pubkey) + check: + keyChanged == true + + # Test listener triggered on changing public key + keyChanged = false + peerStore.keyBook.set(peerId1, + keyPair2.pubkey) + check: + keyChanged == true + + test "AddressBook API": + # Set up address book + var + addressBook = PeerStore.new().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]) + + test "ProtoBook API": + # Set up protocol book + var + protoBook = PeerStore.new().protoBook + + # Test ProtoBook::add + protoBook.add(peerId1, testcodec1) + + check: + toSeq(keys(protoBook.book))[0] == peerId1 + toSeq(values(protoBook.book))[0] == toHashSet([testcodec1]) + + # Test ProtoBook::get + check: + protoBook.get(peerId1) == toHashSet([testcodec1]) + + # Test ProtoBook::delete + check: + # Try to delete peerId that doesn't exist + protoBook.delete(peerId2) == false + + # Delete existing peerId + protoBook.book.len == 1 # sanity + protoBook.delete(peerId1) == true + protoBook.book.len == 0 + + # Test ProtoBook::set + # Set peerId2 with multiple protocols + protoBook.set(peerId2, + toHashSet([testcodec1, testcodec2])) + check: + toSeq(keys(protoBook.book))[0] == peerId2 + toSeq(values(protoBook.book))[0] == toHashSet([testcodec1, testcodec2]) + + test "KeyBook API": + # Set up key book + var + keyBook = PeerStore.new().keyBook + + # Test KeyBook::set + keyBook.set(peerId1, + keyPair1.pubkey) + check: + toSeq(keys(keyBook.book))[0] == peerId1 + toSeq(values(keyBook.book))[0] == keyPair1.pubkey + + # Test KeyBook::get + check: + keyBook.get(peerId1) == keyPair1.pubkey + + # Test KeyBook::delete + check: + # Try to delete peerId that doesn't exist + keyBook.delete(peerId2) == false + + # Delete existing peerId + keyBook.book.len == 1 # sanity + keyBook.delete(peerId1) == true + keyBook.book.len == 0