diff --git a/src/avalanche/peermanager.cpp b/src/avalanche/peermanager.cpp --- a/src/avalanche/peermanager.cpp +++ b/src/avalanche/peermanager.cpp @@ -17,6 +17,11 @@ return it->peerid; } + // Reject invalid proof. + if (!proof.verify()) { + return NO_PEER; + } + // We have no peer for this proof, time to create it. const PeerId peerid = nextPeerId++; auto inserted = peers.emplace(peerid, uint32_t(slots.size()), proof); @@ -61,9 +66,7 @@ bool PeerManager::addNode(NodeId nodeid, const Proof &proof, const CPubKey &pubkey) { const PeerId peerid = getPeer(proof); - - auto pit = peers.find(peerid); - if (pit == peers.end()) { + if (peerid == NO_PEER) { return false; } diff --git a/src/avalanche/proofbuilder.h b/src/avalanche/proofbuilder.h --- a/src/avalanche/proofbuilder.h +++ b/src/avalanche/proofbuilder.h @@ -38,6 +38,12 @@ Proof build(); + /** + * Builds a randomized (and therefore invalid) Proof. + * Useful for tests. + */ + static Proof buildRandom(uint32_t score); + private: ProofId getProofId() const; }; diff --git a/src/avalanche/proofbuilder.cpp b/src/avalanche/proofbuilder.cpp --- a/src/avalanche/proofbuilder.cpp +++ b/src/avalanche/proofbuilder.cpp @@ -4,6 +4,8 @@ #include +#include + namespace avalanche { SignedStake ProofBuilder::StakeSigner::sign(const ProofId &proofid) { @@ -62,4 +64,14 @@ return ProofId(ss.GetHash()); } +Proof ProofBuilder::buildRandom(uint32_t score) { + CKey key; + key.MakeNewKey(true); + + ProofBuilder pb(0, std::numeric_limits::max(), CPubKey()); + pb.addUTXO(COutPoint(TxId(GetRandHash()), 0), (int64_t(score) * COIN) / 100, + 0, std::move(key)); + return pb.build(); +} + } // namespace avalanche 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 @@ -291,7 +291,7 @@ PeerManager pm; // Create one peer. - Proof proof = buildRandomProof(100); + Proof proof = buildRandomProof(100000000); BOOST_CHECK_EQUAL(pm.selectNode(), NO_NODE); // Add 4 nodes. @@ -327,13 +327,20 @@ pm.updateNextRequestTime(n, std::chrono::steady_clock::now())); } - // Move a node from a peer to another. - Proof altproof = buildRandomProof(0); + // Move a node from a peer to another. This peer has a very low score such + // as chances of being picked are 1 in a billion. + Proof altproof = buildRandomProof(1); BOOST_CHECK(pm.addNode(3, altproof, CPubKey())); + int node3selected = 0; for (int i = 0; i < 100; i++) { NodeId n = pm.selectNode(); - BOOST_CHECK(n == 0); + if (n == 3) { + // Selecting this node should be exceedingly unlikely. + BOOST_CHECK(node3selected++ < 1); + } else { + BOOST_CHECK_EQUAL(n, 0); + } BOOST_CHECK( pm.updateNextRequestTime(n, std::chrono::steady_clock::now())); }