diff --git a/src/avalanche/processor.h b/src/avalanche/processor.h --- a/src/avalanche/processor.h +++ b/src/avalanche/processor.h @@ -179,6 +179,17 @@ struct query_timeout {}; class AvalancheProcessor { +public: + using TimePoint = std::chrono::time_point; + + struct Node { + NodeId nodeid; + int64_t score; + + TimePoint nextRequestTime; + CPubKey pubkey; + }; + private: CConnman *connman; std::chrono::milliseconds queryTimeoutDuration; @@ -193,26 +204,16 @@ */ std::atomic round; - using TimePoint = std::chrono::time_point; - - struct Peer { - NodeId nodeid; - int64_t score; - - TimePoint nextRequestTime; - CPubKey pubkey; - }; - using PeerSet = boost::multi_index_container< - Peer, boost::multi_index::indexed_by< + Node, boost::multi_index::indexed_by< // index by nodeid boost::multi_index::hashed_unique< - boost::multi_index::member>, + boost::multi_index::member>, // sorted by nextRequestTime boost::multi_index::ordered_non_unique< boost::multi_index::tag, - boost::multi_index::member>>>; + boost::multi_index::member>>>; RWCollection peerSet; @@ -270,7 +271,7 @@ std::vector &updates); bool addPeer(NodeId nodeid, int64_t score, CPubKey pubkey); - CPubKey getPubKey(NodeId nodeid) const; + bool forNode(NodeId nodeid, std::function func) const; CPubKey getSessionPubKey() const { return sessionKey.GetPubKey(); } diff --git a/src/avalanche/processor.cpp b/src/avalanche/processor.cpp --- a/src/avalanche/processor.cpp +++ b/src/avalanche/processor.cpp @@ -233,11 +233,11 @@ auto w = peerSet.getWriteView(); auto it = w->find(nodeid); if (it != w->end()) { - w->modify(it, [&response](Peer &p) { + w->modify(it, [&response](Node &n) { // FIXME: This will override the time even when we received an // old stale message. This should check that the message is // indeed the most up to date one before updating the time. - p.nextRequestTime = + n.nextRequestTime = std::chrono::steady_clock::now() + std::chrono::milliseconds(response.getCooldown()); }); @@ -346,14 +346,11 @@ .second; } -CPubKey AvalancheProcessor::getPubKey(NodeId nodeid) const { +bool AvalancheProcessor::forNode( + NodeId nodeid, std::function func) const { auto r = peerSet.getReadView(); auto it = r->find(nodeid); - if (it == r->end()) { - return CPubKey(); - } - - return it->pubkey; + return it != r->end() && func(*it); } bool AvalancheProcessor::startEventLoop(CScheduler &scheduler) { @@ -502,8 +499,8 @@ auto w = peerSet.getWriteView(); auto it = w->find(pnode->GetId()); if (it != w->end()) { - w->modify(it, [&timeout](Peer &p) { - p.nextRequestTime = timeout; + w->modify(it, [&timeout](Node &n) { + n.nextRequestTime = timeout; }); } } diff --git a/src/net_processing.cpp b/src/net_processing.cpp --- a/src/net_processing.cpp +++ b/src/net_processing.cpp @@ -3630,18 +3630,18 @@ AvalancheResponse response; verifier >> response; - { - std::array sig; - vRecv >> sig; - - // Unfortunately, the verify API require a vector. - std::vector vchSig{sig.begin(), sig.end()}; - if (!g_avalanche->getPubKey(pfrom->GetId()) - .VerifySchnorr(verifier.GetHash(), vchSig)) { - LOCK(cs_main); - Misbehaving(pfrom, 100, "invalid-ava-response-signature"); - return true; - } + if (!g_avalanche->forNode( + pfrom->GetId(), [&](const AvalancheProcessor::Node &n) { + std::array sig; + vRecv >> sig; + + // Unfortunately, the verify API require a vector. + std::vector vchSig{sig.begin(), sig.end()}; + return n.pubkey.VerifySchnorr(verifier.GetHash(), vchSig); + })) { + LOCK(cs_main); + Misbehaving(pfrom, 100, "invalid-ava-response-signature"); + return true; } std::vector updates;