diff --git a/Kademlia.go b/Kademlia.go index 1b2101a..f0189be 100644 --- a/Kademlia.go +++ b/Kademlia.go @@ -181,6 +181,14 @@ func IsNodeContact(nodeID []byte) (node *dht.Node, peer *PeerInfo) { // FindNode finds a node via the DHT func FindNode(nodeID []byte, Timeout time.Duration) (node *dht.Node, peer *PeerInfo, err error) { + // first check if in mirrored node list + var nodeID2 [protocol.HashSize]byte + copy(nodeID2[:], nodeID) + if peer, ok := nodeList[nodeID2]; ok { + return nil, peer, nil + } + + // Search the node via DHT. node, err = nodesDHT.FindNode(nodeID) if node == nil { return nil, nil, err diff --git a/Peer ID.go b/Peer ID.go index 1f7f83f..b4786da 100644 --- a/Peer ID.go +++ b/Peer ID.go @@ -28,6 +28,7 @@ var nodeID []byte func initPeerID() { peerList = make(map[[btcec.PubKeyBytesLenCompressed]byte]*PeerInfo) + nodeList = make(map[[protocol.HashSize]byte]*PeerInfo) // load existing key from config, if available if len(config.PrivateKey) > 0 { @@ -120,6 +121,9 @@ type peerAddress struct { var peerList map[[btcec.PubKeyBytesLenCompressed]byte]*PeerInfo var peerlistMutex sync.RWMutex +// nodeList is a mirror of peerList but using the node ID +var nodeList map[[protocol.HashSize]byte]*PeerInfo + // PeerlistAdd adds a new peer to the peer list. It does not validate the peer info. If the peer is already added, it does nothing. Connections must be live. func PeerlistAdd(PublicKey *btcec.PublicKey, connections ...*Connection) (peer *PeerInfo, added bool) { if len(connections) == 0 { @@ -140,6 +144,11 @@ func PeerlistAdd(PublicKey *btcec.PublicKey, connections ...*Connection) (peer * peerList[publicKeyCompressed] = peer + // also add to mirrored nodeList + var nodeID [protocol.HashSize]byte + copy(nodeID[:], peer.NodeID) + nodeList[nodeID] = peer + // add to Kademlia nodesDHT.AddNode(&dht.Node{ID: peer.NodeID, Info: peer}) @@ -168,6 +177,11 @@ func PeerlistRemove(peer *PeerInfo) { nodesDHT.RemoveNode(peer.NodeID) delete(peerList, publicKey2Compressed(peer.PublicKey)) + + var nodeID [protocol.HashSize]byte + copy(nodeID[:], peer.NodeID) + + delete(nodeList, nodeID) } // PeerlistGet returns the full peer list