diff --git a/src/avalanche.h b/src/avalanche.h --- a/src/avalanche.h +++ b/src/avalanche.h @@ -275,6 +275,7 @@ int64_t score; TimePoint nextRequestTime; + CPubKey pubkey; }; typedef boost::multi_index_container< @@ -351,7 +352,8 @@ bool registerVotes(NodeId nodeid, const AvalancheResponse &response, std::vector &updates); - bool addPeer(NodeId nodeid, int64_t score); + bool addPeer(NodeId nodeid, int64_t score, CPubKey pubkey); + CPubKey getPubKey(NodeId nodeid) const; CPubKey getSessionPubKey() const { return sessionKey.GetPubKey(); } diff --git a/src/avalanche.cpp b/src/avalanche.cpp --- a/src/avalanche.cpp +++ b/src/avalanche.cpp @@ -340,12 +340,23 @@ return true; } -bool AvalancheProcessor::addPeer(NodeId nodeid, int64_t score) { +bool AvalancheProcessor::addPeer(NodeId nodeid, int64_t score, CPubKey pubkey) { return peerSet.getWriteView() - ->insert({nodeid, score, std::chrono::steady_clock::now()}) + ->insert({nodeid, score, std::chrono::steady_clock::now(), + std::move(pubkey)}) .second; } +CPubKey AvalancheProcessor::getPubKey(NodeId nodeid) const { + auto r = peerSet.getReadView(); + auto it = r->find(nodeid); + if (it == r->end()) { + return CPubKey(); + } + + return it->pubkey; +} + bool AvalancheProcessor::startEventLoop(CScheduler &scheduler) { LOCK(cs_running); if (running) { diff --git a/src/rpc/avalanche.cpp b/src/rpc/avalanche.cpp --- a/src/rpc/avalanche.cpp +++ b/src/rpc/avalanche.cpp @@ -31,17 +31,19 @@ static UniValue addavalanchepeer(const Config &config, const JSONRPCRequest &request) { - if (request.fHelp || request.params.size() != 1) { + if (request.fHelp || request.params.size() != 2) { throw std::runtime_error( RPCHelpMan{ "addavalanchepeer", "\nAdd a peer to the set of peer to poll for avalanche.\n", { {"nodeid", RPCArg::Type::NUM, false}, + {"publickey", RPCArg::Type::STR_HEX, false}, }} .ToString() + "\nArguments\n" + "1. nodeid (number, required) Node to be added to avalanche.\n" + + "2. publickey (sring, required) The public key of the node.\n" + "\nExamples:\n" + HelpExampleRpc("addavalanchepeer", "5")); } @@ -49,6 +51,7 @@ throw JSONRPCError(RPC_INTERNAL_ERROR, "Avalanche is not initialized"); } + // Parse nodeid if (!request.params[0].isNum()) { throw JSONRPCError( RPC_INVALID_PARAMETER, @@ -56,7 +59,19 @@ } NodeId nodeid = request.params[0].get_int64(); - g_avalanche->addPeer(nodeid, 0); + + // Parse the pubkey + const std::string keyHex = request.params[1].get_str(); + if ((keyHex.length() != 2 * CPubKey::COMPRESSED_PUBLIC_KEY_SIZE && + keyHex.length() != 2 * CPubKey::PUBLIC_KEY_SIZE) || + !IsHex(keyHex)) { + throw JSONRPCError(RPC_INVALID_ADDRESS_OR_KEY, + strprintf("Invalid public key: %s\n", keyHex)); + } + + CPubKey pubkey{HexToPubKey(keyHex)}; + + g_avalanche->addPeer(nodeid, 0, pubkey); return {}; } diff --git a/src/test/avalanche_tests.cpp b/src/test/avalanche_tests.cpp --- a/src/test/avalanche_tests.cpp +++ b/src/test/avalanche_tests.cpp @@ -197,7 +197,7 @@ std::array nodes; for (CNode *&n : nodes) { n = ConnectNode(config, nServices, peerLogic, connman); - BOOST_CHECK(p.addPeer(n->GetId(), 0)); + BOOST_CHECK(p.addPeer(n->GetId(), 0, CPubKey())); } return nodes; @@ -521,7 +521,7 @@ auto avanode = ConnectNode(config, NODE_AVALANCHE, *peerLogic, connman.get()); NodeId avanodeid = avanode->GetId(); - BOOST_CHECK(p.addPeer(avanodeid, 0)); + BOOST_CHECK(p.addPeer(avanodeid, 0, CPubKey())); // It returns the avalanche peer. BOOST_CHECK_EQUAL(AvalancheTest::getSuitableNodeToQuery(p), avanodeid); @@ -675,7 +675,7 @@ auto avanode = ConnectNode(config, NODE_AVALANCHE, *peerLogic, connman.get()); NodeId avanodeid = avanode->GetId(); - BOOST_CHECK(p.addPeer(avanodeid, 0)); + BOOST_CHECK(p.addPeer(avanodeid, 0, CPubKey())); // Expire requests after some time. auto queryTimeDuration = std::chrono::milliseconds(10); @@ -730,7 +730,7 @@ std::array nodes; for (auto &n : nodes) { n = ConnectNode(config, NODE_AVALANCHE, *peerLogic, connman.get()); - BOOST_CHECK(p.addPeer(n->GetId(), 0)); + BOOST_CHECK(p.addPeer(n->GetId(), 0, CPubKey())); } // Add a block to poll @@ -891,7 +891,7 @@ auto avanode = ConnectNode(config, NODE_AVALANCHE, *peerLogic, connman.get()); NodeId nodeid = avanode->GetId(); - BOOST_CHECK(p.addPeer(nodeid, 0)); + BOOST_CHECK(p.addPeer(nodeid, 0, CPubKey())); // There is no query in flight at the moment. BOOST_CHECK_EQUAL(AvalancheTest::getSuitableNodeToQuery(p), nodeid); diff --git a/test/functional/abc-p2p-avalanche.py b/test/functional/abc-p2p-avalanche.py --- a/test/functional/abc-p2p-avalanche.py +++ b/test/functional/abc-p2p-avalanche.py @@ -6,7 +6,11 @@ import random from test_framework.mininode import P2PInterface, mininode_lock -from test_framework.messages import AvalancheVote, CInv, msg_avapoll +from test_framework.messages import ( + AvalancheVote, + CInv, + msg_avapoll, +) from test_framework.test_framework import BitcoinTestFramework from test_framework.util import assert_equal, wait_until from test_framework import schnorr @@ -156,7 +160,12 @@ [AvalancheVote(BLOCK_UNKNOWN, h) for h in various_block_hashes[-3:]]) self.log.info("Trigger polling from the node...") - node.addavalanchepeer(nodeid) + # duplicate the deterministic sig test from src/test/key_tests.cpp + privkey = bytes.fromhex( + "12b004fff7f4b69ef8650e767f18f11ede158148b425660723b9f9a66e61f747") + pubkey = schnorr.getpubkey(privkey, compressed=True) + + node.addavalanchepeer(nodeid, pubkey.hex()) # Sanity check fork_node = self.nodes[1]