diff --git a/src/avalanche/peermanager.h b/src/avalanche/peermanager.h --- a/src/avalanche/peermanager.h +++ b/src/avalanche/peermanager.h @@ -134,6 +134,23 @@ bool forNode(NodeId nodeid, std::function func) const; bool updateNextRequestTime(NodeId nodeid, TimePoint timeout); + /** + * Add a proof, return true if it was successfully added as a new peer + * (good proof) or if it was added as an orphan proof. Return false if + * the proof is permanently bad or if we already have it. + * + * This method should be used to add proofs when we don't already know + * any owner's node for them, e.g. for proofs relayed through the + * inventory network message. For proofs that already belong to known + * nodes, use addNode instead. + */ + bool addProof(Proof proof); + + /** + * Get a proof by ID if we have it, else nullptr. + */ + const Proof *getProof(const ProofId proofId) const; + /** * Randomly select a node to poll. */ @@ -183,6 +200,8 @@ private: PeerSet::iterator fetchOrCreatePeer(const Proof &proof); + PeerSet::iterator fetchOrCreatePeer(const Proof &proof, + ProofValidationState &state); bool addNodeToPeer(const PeerSet::iterator &it); bool removeNodeFromPeer(const PeerSet::iterator &it, uint32_t count = 1); }; diff --git a/src/avalanche/peermanager.cpp b/src/avalanche/peermanager.cpp --- a/src/avalanche/peermanager.cpp +++ b/src/avalanche/peermanager.cpp @@ -180,6 +180,25 @@ } } +bool PeerManager::addProof(Proof proof) { + if (getProof(proof.getId())) { + return false; + } + ProofValidationState state; + auto it = fetchOrCreatePeer(proof, state); + // TODO: check validation state and add orphan proof if needed + return it != peers.end(); +} + +const Proof *PeerManager::getProof(const ProofId proofId) const { + auto &peerView = peers.get(); + auto peerIt = peerView.find(proofId); + if (peerIt != peerView.end()) { + return &peerIt->proof; + } + return nullptr; +} + PeerId PeerManager::getPeerId(const Proof &proof) { auto it = fetchOrCreatePeer(proof); return it == peers.end() ? NO_PEER : it->peerid; @@ -187,6 +206,13 @@ PeerManager::PeerSet::iterator PeerManager::fetchOrCreatePeer(const Proof &proof) { + ProofValidationState state; + return fetchOrCreatePeer(proof, state); +} + +PeerManager::PeerSet::iterator +PeerManager::fetchOrCreatePeer(const Proof &proof, + ProofValidationState &state) { { // Check if we already know of that peer. auto &pview = peers.get(); @@ -201,7 +227,6 @@ LOCK(cs_main); const CCoinsViewCache &coins = ::ChainstateActive().CoinsTip(); - ProofValidationState state; if (!proof.verify(state, coins)) { return peers.end(); } diff --git a/src/avalanche/test/peermanager_tests.cpp b/src/avalanche/test/peermanager_tests.cpp --- a/src/avalanche/test/peermanager_tests.cpp +++ b/src/avalanche/test/peermanager_tests.cpp @@ -438,4 +438,42 @@ NO_PEER); } +BOOST_AUTO_TEST_CASE(add_get_proofs) { + avalanche::PeerManager pm; + BOOST_CHECK_EQUAL(pm.selectPeer(), NO_PEER); + + // Add 5 good proofs + std::vector proofIds; + Proof p1 = buildRandomProof(100); + proofIds.push_back(p1.getId()); + BOOST_CHECK(pm.addProof(p1)); + + for (int i = 0; i < 4; i++) { + Proof p = buildRandomProof(100); + proofIds.push_back(p.getId()); + BOOST_CHECK(pm.addProof(std::move(p))); + } + + // Fail to add a proof already added previously + BOOST_CHECK(!pm.addProof(p1)); + const Proof *pp2 = pm.getProof(p1.getId()); + BOOST_CHECK(pp2); + BOOST_CHECK_EQUAL(pp2->getId(), p1.getId()); + + // Fail to add a bad proof (below dust threshold) + Proof p3 = buildRandomProof(1); + BOOST_CHECK(!pm.addProof(p3)); + + // Fail to get a proof that was never added. + Proof p2 = buildRandomProof(100); + BOOST_CHECK(!pm.getProof(p2.getId())); + + // Check that we still have all the good proofs + for (auto id : proofIds) { + const Proof *pp = pm.getProof(id); + BOOST_CHECK(pp); + BOOST_CHECK_EQUAL(pp->getId(), id); + } +} + BOOST_AUTO_TEST_SUITE_END()