Changeset View
Changeset View
Standalone View
Standalone View
test/functional/test_framework/avatools.py
| #!/usr/bin/env python3 | #!/usr/bin/env python3 | ||||
| # Copyright (c) 2021 The Bitcoin ABC developers | # Copyright (c) 2021 The Bitcoin ABC developers | ||||
| # Distributed under the MIT software license, see the accompanying | # Distributed under the MIT software license, see the accompanying | ||||
| # file COPYING or http://www.opensource.org/licenses/mit-license.php. | # file COPYING or http://www.opensource.org/licenses/mit-license.php. | ||||
| """Utilities for avalanche tests.""" | """Utilities for avalanche tests.""" | ||||
| import random | import random | ||||
| import struct | import struct | ||||
| from typing import Any, Dict, List, Optional | from typing import Any, Dict, List, Optional | ||||
| from .authproxy import JSONRPCException | from .authproxy import JSONRPCException | ||||
| from .key import ECKey | from .key import ECKey | ||||
| from .messages import ( | from .messages import ( | ||||
| MSG_AVA_PROOF, | |||||
| MSG_BLOCK, | MSG_BLOCK, | ||||
| NODE_AVALANCHE, | NODE_AVALANCHE, | ||||
| NODE_NETWORK, | NODE_NETWORK, | ||||
| AvalancheDelegation, | AvalancheDelegation, | ||||
| AvalancheProof, | AvalancheProof, | ||||
| AvalancheResponse, | AvalancheResponse, | ||||
| CInv, | CInv, | ||||
| CTransaction, | CTransaction, | ||||
| FromHex, | FromHex, | ||||
| LegacyAvalancheProof, | LegacyAvalancheProof, | ||||
| TCPAvalancheResponse, | TCPAvalancheResponse, | ||||
| ToHex, | ToHex, | ||||
| calculate_shortid, | calculate_shortid, | ||||
| hash256, | hash256, | ||||
| msg_avahello, | msg_avahello, | ||||
| msg_avapoll, | msg_avapoll, | ||||
| msg_avaproof, | msg_avaproof, | ||||
| msg_avaproofs, | msg_avaproofs, | ||||
| msg_notfound, | |||||
| msg_tcpavaresponse, | msg_tcpavaresponse, | ||||
| ) | ) | ||||
| from .p2p import P2PInterface, p2p_lock | from .p2p import P2PInterface, p2p_lock | ||||
| from .test_node import TestNode | from .test_node import TestNode | ||||
| from .util import assert_equal, satoshi_round, wait_until_helper | from .util import assert_equal, satoshi_round, wait_until_helper | ||||
| from .wallet_util import bytes_to_wif | from .wallet_util import bytes_to_wif | ||||
| ▲ Show 20 Lines • Show All 125 Lines • ▼ Show 20 Lines | def proof_found(): | ||||
| except JSONRPCException: | except JSONRPCException: | ||||
| return False | return False | ||||
| wait_until_helper(proof_found, timeout=timeout) | wait_until_helper(proof_found, timeout=timeout) | ||||
| if expect_orphan is not None: | if expect_orphan is not None: | ||||
| assert_equal(expect_orphan, wait_for_proof.is_orphan) | assert_equal(expect_orphan, wait_for_proof.is_orphan) | ||||
| class AvaP2PInterface(P2PInterface): | class NoHandshakeAvaP2PInterface(P2PInterface): | ||||
| """P2PInterface with avalanche capabilities""" | """P2PInterface with avalanche capabilities""" | ||||
| def __init__(self): | def __init__(self): | ||||
| self.round = 0 | self.round = 0 | ||||
| self.avahello = None | self.avahello = None | ||||
| self.avaresponses = [] | self.avaresponses = [] | ||||
| self.avapolls = [] | self.avapolls = [] | ||||
| self.nodeid: Optional[int] = None | self.nodeid: Optional[int] = None | ||||
| ▲ Show 20 Lines • Show All 69 Lines • ▼ Show 20 Lines | class NoHandshakeAvaP2PInterface(P2PInterface): | ||||
| def wait_for_avahello(self, timeout=5): | def wait_for_avahello(self, timeout=5): | ||||
| self.wait_until( | self.wait_until( | ||||
| lambda: self.avahello is not None, | lambda: self.avahello is not None, | ||||
| timeout=timeout) | timeout=timeout) | ||||
| with p2p_lock: | with p2p_lock: | ||||
| return self.avahello | return self.avahello | ||||
| def send_avahello(self, delegation_hex: str, delegated_privkey: ECKey): | def build_avahello(self, delegation: AvalancheDelegation, | ||||
| delegation = FromHex(AvalancheDelegation(), delegation_hex) | delegated_privkey: ECKey) -> msg_avahello: | ||||
| local_sighash = hash256( | local_sighash = hash256( | ||||
| delegation.getid() + | delegation.getid() + | ||||
| struct.pack("<QQQQ", self.local_nonce, self.remote_nonce, | struct.pack("<QQQQ", self.local_nonce, self.remote_nonce, | ||||
| self.local_extra_entropy, self.remote_extra_entropy)) | self.local_extra_entropy, self.remote_extra_entropy)) | ||||
| msg = msg_avahello() | msg = msg_avahello() | ||||
| msg.hello.delegation = delegation | msg.hello.delegation = delegation | ||||
| msg.hello.sig = delegated_privkey.sign_schnorr(local_sighash) | msg.hello.sig = delegated_privkey.sign_schnorr(local_sighash) | ||||
| return msg | |||||
| def send_avahello(self, delegation_hex: str, delegated_privkey: ECKey): | |||||
| delegation = FromHex(AvalancheDelegation(), delegation_hex) | |||||
| msg = self.build_avahello(delegation, delegated_privkey) | |||||
| self.send_message(msg) | self.send_message(msg) | ||||
| return delegation.proofid | return msg.hello.delegation.proofid | ||||
| def send_avaproof(self, proof: AvalancheProof): | def send_avaproof(self, proof: AvalancheProof): | ||||
| msg = msg_avaproof() | msg = msg_avaproof() | ||||
| msg.proof = proof | msg.proof = proof | ||||
| self.send_message(msg) | self.send_message(msg) | ||||
| class AvaP2PInterface(NoHandshakeAvaP2PInterface): | |||||
| def __init__(self, node=None): | |||||
| super().__init__() | |||||
| self.master_privkey = None | |||||
| self.proof = None | |||||
| self.delegated_privkey = ECKey() | |||||
| self.delegated_privkey.generate() | |||||
| self.delegation = None | |||||
| if node is not None: | |||||
| self.master_privkey, self.proof = gen_proof(node) | |||||
| delegation_hex = node.delegateavalancheproof( | |||||
| f"{self.proof.limited_proofid:0{64}x}", | |||||
| bytes_to_wif(self.master_privkey.get_bytes()), | |||||
| self.delegated_privkey.get_pubkey().get_bytes().hex(), | |||||
| ) | |||||
| assert node.verifyavalanchedelegation(delegation_hex) | |||||
| self.delegation = FromHex(AvalancheDelegation(), delegation_hex) | |||||
| def on_version(self, message): | |||||
| super().on_version(message) | |||||
| avahello = msg_avahello() | |||||
| if self.delegation is not None: | |||||
| avahello = self.build_avahello( | |||||
| self.delegation, self.delegated_privkey) | |||||
| elif self.proof is not None: | |||||
| avahello = self.build_avahello( | |||||
| AvalancheDelegation( | |||||
| self.proof.limited_proofid, | |||||
| self.master_privkey.get_pubkey().get_bytes()), | |||||
| self.master_privkey) | |||||
| self.send_message(avahello) | |||||
| def on_getdata(self, message): | |||||
| super().on_getdata(message) | |||||
| not_found = [] | |||||
| for inv in message.inv: | |||||
| if inv.type == MSG_AVA_PROOF and self.proof is not None and inv.hash == self.proof.proofid: | |||||
| self.send_avaproof(self.proof) | |||||
| else: | |||||
| not_found.append(inv) | |||||
| if len(not_found) > 0: | |||||
| self.send_message(msg_notfound(not_found)) | |||||
| def get_ava_p2p_interface( | def get_ava_p2p_interface( | ||||
| node: TestNode, | node: TestNode, | ||||
| services=NODE_NETWORK | NODE_AVALANCHE) -> AvaP2PInterface: | services=NODE_NETWORK | NODE_AVALANCHE) -> NoHandshakeAvaP2PInterface: | ||||
| """Build and return an AvaP2PInterface connected to the specified | """Build and return a NoHandshakeAvaP2PInterface connected to the specified | ||||
| TestNode. | TestNode. | ||||
| """ | """ | ||||
| n = AvaP2PInterface() | n = NoHandshakeAvaP2PInterface() | ||||
| node.add_p2p_connection( | node.add_p2p_connection( | ||||
| n, services=services) | n, services=services) | ||||
| n.wait_for_verack() | n.wait_for_verack() | ||||
| n.nodeid = node.getpeerinfo()[-1]['id'] | n.nodeid = node.getpeerinfo()[-1]['id'] | ||||
| return n | return n | ||||
| Show All 30 Lines | |||||