From 4f0155e626a9f7b0c8a71d66328618eb6a780453 Mon Sep 17 00:00:00 2001
From: jangko <jangko128@gmail.com>
Date: Tue, 23 Aug 2022 15:29:34 +0700
Subject: [PATCH] add some peer convenience funcs to PeerPool and EthNode

---
 eth/p2p.nim           | 18 ++++++++++++++++++
 eth/p2p/peer_pool.nim | 24 ++++++++++++++++++++----
 2 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/eth/p2p.nim b/eth/p2p.nim
index a0b6bcb..ec14ae1 100644
--- a/eth/p2p.nim
+++ b/eth/p2p.nim
@@ -176,3 +176,21 @@ proc getPeer*(node: EthereumNode, peerId: NodeId, Protocol: type): Option[Peer]
   for peer in node.peers(Protocol):
     if peer.remote.id == peerId:
       return some(peer)
+
+proc connectToNode*(node: EthereumNode, n: Node) {.async.} =
+  await node.peerPool.connectToNode(n)
+
+proc connectToNode*(node: EthereumNode, n: ENode) {.async.} =
+  await node.peerPool.connectToNode(n)
+
+func numPeers*(node: EthereumNode): int =
+  node.peerPool.numPeers
+
+func hasPeer*(node: EthereumNode, n: ENode): bool =
+  n in node.peerPool
+
+func hasPeer*(node: EthereumNode, n: Node): bool =
+  n in node.peerPool
+
+func hasPeer*(node: EthereumNode, n: Peer): bool =
+  n in node.peerPool
diff --git a/eth/p2p/peer_pool.nim b/eth/p2p/peer_pool.nim
index 3a66518..9a5fb18 100644
--- a/eth/p2p/peer_pool.nim
+++ b/eth/p2p/peer_pool.nim
@@ -14,7 +14,7 @@ import
   std/[os, tables, times, random, sequtils, options],
   chronos, chronicles,
   ".."/[rlp, keys, common],
-  ./private/p2p_types, "."/[discovery, kademlia, rlpx]
+  ./private/p2p_types, "."/[discovery, kademlia, rlpx, enode]
 
 const
   lookupInterval = 5
@@ -36,7 +36,7 @@ proc newPeerPool*(
 proc nodesToConnect(p: PeerPool): seq[Node] =
   p.discovery.randomNodes(p.minPeers).filterIt(it notin p.discovery.bootstrapNodes)
 
-proc addObserver(p: PeerPool, observerId: int, observer: PeerObserver) =
+proc addObserver*(p: PeerPool, observerId: int, observer: PeerObserver) =
   doAssert(observerId notin p.observers)
   p.observers[observerId] = observer
   if not observer.onPeerConnected.isNil:
@@ -44,7 +44,7 @@ proc addObserver(p: PeerPool, observerId: int, observer: PeerObserver) =
       if observer.protocol.isNil or peer.supports(observer.protocol):
         observer.onPeerConnected(peer)
 
-proc delObserver(p: PeerPool, observerId: int) =
+proc delObserver*(p: PeerPool, observerId: int) =
   p.observers.del(observerId)
 
 proc addObserver*(p: PeerPool, observerId: ref, observer: PeerObserver) =
@@ -123,9 +123,12 @@ proc connectToNode*(p: PeerPool, n: Node) {.async.} =
     trace "Connection established (outgoing)", peer
     p.addPeer(peer)
 
+proc connectToNode*(p: PeerPool, n: ENode) {.async.} =
+  await p.connectToNode(newNode(n))
+
 proc connectToNodes(p: PeerPool, nodes: seq[Node]) {.async.} =
   for node in nodes:
-    discard p.connectToNode(node)
+    await p.connectToNode(node)
 
     # # TODO: Consider changing connect() to raise an exception instead of
     # # returning None, as discussed in
@@ -213,3 +216,16 @@ iterator peers*(p: PeerPool, Protocol: type): Peer =
     if peer.supports(Protocol):
       yield peer
 
+func numPeers*(p: PeerPool): int =
+  p.connectedNodes.len
+
+func contains*(p: PeerPool, n: ENode): bool =
+  for remote, _ in p.connectedNodes:
+    if remote.node == n:
+      return true
+
+func contains*(p: PeerPool, n: Node): bool =
+  n in p.connectedNodes
+
+func contains*(p: PeerPool, n: Peer): bool =
+  n.remote in p.connectedNodes