Changeset View
Changeset View
Standalone View
Standalone View
test/functional/test_framework/mininode.py
Show All 17 Lines | |||||
import asyncio | import asyncio | ||||
from collections import defaultdict | from collections import defaultdict | ||||
from io import BytesIO | from io import BytesIO | ||||
import logging | import logging | ||||
import struct | import struct | ||||
import sys | import sys | ||||
import threading | import threading | ||||
from typing import List | |||||
from test_framework.messages import ( | from test_framework.messages import ( | ||||
AvalancheResponse, | |||||
CBlockHeader, | CBlockHeader, | ||||
CInv, | |||||
MIN_VERSION_SUPPORTED, | MIN_VERSION_SUPPORTED, | ||||
msg_addr, | msg_addr, | ||||
msg_addrv2, | msg_addrv2, | ||||
msg_avahello, | |||||
msg_avapoll, | msg_avapoll, | ||||
msg_avaproof, | msg_avaproof, | ||||
msg_tcpavaresponse, | msg_tcpavaresponse, | ||||
msg_avahello, | |||||
msg_block, | msg_block, | ||||
MSG_BLOCK, | MSG_BLOCK, | ||||
msg_blocktxn, | msg_blocktxn, | ||||
msg_cfcheckpt, | msg_cfcheckpt, | ||||
msg_cfheaders, | msg_cfheaders, | ||||
msg_cfilter, | msg_cfilter, | ||||
msg_cmpctblock, | msg_cmpctblock, | ||||
msg_feefilter, | msg_feefilter, | ||||
Show All 17 Lines | from test_framework.messages import ( | ||||
msg_sendheaders, | msg_sendheaders, | ||||
msg_tx, | msg_tx, | ||||
MSG_TX, | MSG_TX, | ||||
MSG_TYPE_MASK, | MSG_TYPE_MASK, | ||||
msg_verack, | msg_verack, | ||||
msg_version, | msg_version, | ||||
NODE_NETWORK, | NODE_NETWORK, | ||||
sha256, | sha256, | ||||
TCPAvalancheResponse, | |||||
) | ) | ||||
from test_framework.util import wait_until | from test_framework.util import wait_until | ||||
logger = logging.getLogger("TestFramework.mininode") | logger = logging.getLogger("TestFramework.mininode") | ||||
MESSAGEMAP = { | MESSAGEMAP = { | ||||
b"addr": msg_addr, | b"addr": msg_addr, | ||||
b"addrv2": msg_addrv2, | b"addrv2": msg_addrv2, | ||||
▲ Show 20 Lines • Show All 709 Lines • ▼ Show 20 Lines | def wait_for_broadcast(self, txns, timeout=60): | ||||
"""Waits for the txns (list of txids) to complete initial broadcast. | """Waits for the txns (list of txids) to complete initial broadcast. | ||||
The mempool should mark unbroadcast=False for these transactions. | The mempool should mark unbroadcast=False for these transactions. | ||||
""" | """ | ||||
# Wait until invs have been received (and getdatas sent) for each txid. | # Wait until invs have been received (and getdatas sent) for each txid. | ||||
self.wait_until(lambda: set(self.get_invs()) == set( | self.wait_until(lambda: set(self.get_invs()) == set( | ||||
[int(tx, 16) for tx in txns]), timeout) | [int(tx, 16) for tx in txns]), timeout) | ||||
# Flush messages and wait for the getdatas to be processed | # Flush messages and wait for the getdatas to be processed | ||||
self.sync_with_ping() | self.sync_with_ping() | ||||
class P2PAvalancheTestNode(P2PInterface): | |||||
"""A P2PInterface supporting Avalanche messages""" | |||||
DUMMY_PROOFID = 1337 | |||||
def __init__(self): | |||||
self.round = 0 | |||||
self.avahello = None | |||||
self.avaresponses = [] | |||||
self.avapolls = [] | |||||
self.avaproof = None | |||||
super().__init__() | |||||
def peer_connect(self, *args, **kwargs): | |||||
create_conn = super().peer_connect(*args, **kwargs) | |||||
# Save the nonce and extra entropy so they can be reused later. | |||||
self.local_nonce = self.on_connection_send_msg.nNonce | |||||
self.local_extra_entropy = self.on_connection_send_msg.nExtraEntropy | |||||
return create_conn | |||||
def on_version(self, message): | |||||
super().on_version(message) | |||||
# Save the nonce and extra entropy so they can be reused later. | |||||
self.remote_nonce = message.nNonce | |||||
self.remote_extra_entropy = message.nExtraEntropy | |||||
def on_avaresponse(self, message): | |||||
with mininode_lock: | |||||
self.avaresponses.append(message.response) | |||||
def on_avapoll(self, message): | |||||
with mininode_lock: | |||||
self.avapolls.append(message.poll) | |||||
def on_avahello(self, message): | |||||
with mininode_lock: | |||||
assert(self.avahello is None) | |||||
self.avahello = message | |||||
def send_avaresponse(self, round, votes, privkey): | |||||
response = AvalancheResponse(round, 0, votes) | |||||
sig = privkey.sign_schnorr(response.get_hash()) | |||||
msg = msg_tcpavaresponse() | |||||
msg.response = TCPAvalancheResponse(response, sig) | |||||
self.send_message(msg) | |||||
def wait_for_avaresponse(self, timeout=5): | |||||
wait_until( | |||||
lambda: len(self.avaresponses) > 0, | |||||
timeout=timeout, | |||||
lock=mininode_lock) | |||||
with mininode_lock: | |||||
return self.avaresponses.pop(0) | |||||
def send_poll(self, hashes): | |||||
msg = msg_avapoll() | |||||
msg.poll.round = self.round | |||||
self.round += 1 | |||||
for h in hashes: | |||||
msg.poll.invs.append(CInv(2, h)) | |||||
self.send_message(msg) | |||||
def get_avapoll_if_available(self): | |||||
with mininode_lock: | |||||
return self.avapolls.pop(0) if len(self.avapolls) > 0 else None | |||||
def wait_for_avahello(self, timeout=5): | |||||
wait_until( | |||||
lambda: self.avahello is not None, | |||||
timeout=timeout, | |||||
lock=mininode_lock) | |||||
with mininode_lock: | |||||
return self.avahello | |||||
def send_avahello(self): | |||||
msg = msg_avahello() | |||||
msg.hello.delegation.proofid = self.DUMMY_PROOFID | |||||
self.send_message(msg) | |||||
def send_getdata(self, inv: List[CInv]): | |||||
msg = msg_getdata() | |||||
msg.inv = inv | |||||
self.send_message(msg) | |||||
def on_avaproof(self, message): | |||||
with mininode_lock: | |||||
assert(self.avaproof is None) | |||||
self.avaproof = message | |||||
def wait_for_avaproof(self, timeout=10): | |||||
wait_until( | |||||
lambda: self.avaproof is not None, | |||||
timeout=timeout, | |||||
lock=mininode_lock) | |||||
with mininode_lock: | |||||
return self.avaproof |