Skip to content

Commit

Permalink
[qa] mininode: Expose connection state through is_connected
Browse files Browse the repository at this point in the history
  • Loading branch information
MarcoFalke authored and furszy committed Jun 28, 2021
1 parent 0e264e2 commit 688190c
Show file tree
Hide file tree
Showing 3 changed files with 48 additions and 44 deletions.
6 changes: 3 additions & 3 deletions test/functional/p2p_leak.py
Expand Up @@ -118,11 +118,11 @@ def run_test(self):
time.sleep(5)

#This node should have been banned
assert no_version_bannode.state != "connected"
assert not no_version_bannode.is_connected

# These nodes should have been disconnected
assert unsupported_service_bit5_node.state != "connected"
assert unsupported_service_bit7_node.state != "connected"
assert not unsupported_service_bit5_node.is_connected
assert not unsupported_service_bit7_node.is_connected

self.nodes[0].disconnect_p2ps()

Expand Down
18 changes: 9 additions & 9 deletions test/functional/p2p_timeouts.py
Expand Up @@ -47,9 +47,9 @@ def run_test(self):

sleep(1)

assert no_verack_node.connected
assert no_version_node.connected
assert no_send_node.connected
assert no_verack_node.is_connected
assert no_version_node.is_connected
assert no_send_node.is_connected

no_verack_node.send_message(msg_ping())
no_version_node.send_message(msg_ping())
Expand All @@ -58,18 +58,18 @@ def run_test(self):

assert "version" in no_verack_node.last_message

assert no_verack_node.connected
assert no_version_node.connected
assert no_send_node.connected
assert no_verack_node.is_connected
assert no_version_node.is_connected
assert no_send_node.is_connected

no_verack_node.send_message(msg_ping())
no_version_node.send_message(msg_ping())

sleep(31)

assert not no_verack_node.connected
assert not no_version_node.connected
assert not no_send_node.connected
assert not no_verack_node.is_connected
assert not no_version_node.is_connected
assert not no_send_node.is_connected

if __name__ == '__main__':
TimeoutsTest().main()
68 changes: 36 additions & 32 deletions test/functional/test_framework/mininode.py
Expand Up @@ -78,14 +78,20 @@ def __init__(self):

super().__init__(map=mininode_socket_map)

self._conn_open = False

@property
def is_connected(self):
return self._conn_open

def peer_connect(self, dstaddr, dstport, net="regtest"):
self.dstaddr = dstaddr
self.dstport = dstport
self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
self.socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
self.sendbuf = b""
self.recvbuf = b""
self.state = "connecting"
self._asyncore_pre_connection = True
self.network = net
self.disconnect = False

Expand All @@ -98,22 +104,23 @@ def peer_connect(self, dstaddr, dstport, net="regtest"):

def peer_disconnect(self):
# Connection could have already been closed by other end.
if self.state == "connected":
self.disconnect_node()
if self.is_connected:
self.disconnect = True # Signal asyncore to disconnect

# Connection and disconnection methods

def handle_connect(self):
"""asyncore callback when a connection is opened."""
if self.state != "connected":
if not self.is_connected:
logger.debug("Connected & Listening: %s:%d" % (self.dstaddr, self.dstport))
self.state = "connected"
self._conn_open = True
self._asyncore_pre_connection = False
self.on_open()

def handle_close(self):
"""asyncore callback when a connection is closed."""
logger.debug("Closing connection to: %s:%d" % (self.dstaddr, self.dstport))
self.state = "closed"
self._conn_open = False
self.recvbuf = b""
self.sendbuf = b""
try:
Expand All @@ -122,13 +129,6 @@ def handle_close(self):
pass
self.on_close()

def disconnect_node(self):
"""Disconnect the p2p connection.
Called by the test logic thread. Causes the p2p connection
to be disconnected on the next iteration of the asyncore loop."""
self.disconnect = True

# Socket read methods

def handle_read(self):
Expand Down Expand Up @@ -184,17 +184,16 @@ def on_message(self, message):
def writable(self):
"""asyncore method to determine whether the handle_write() callback should be called on the next loop."""
with mininode_lock:
pre_connection = self.state == "connecting"
length = len(self.sendbuf)
return (length > 0 or pre_connection)
return length > 0 or self._asyncore_pre_connection

def handle_write(self):
"""asyncore callback when data should be written to the socket."""
with mininode_lock:
# asyncore does not expose socket connection, only the first read/write
# event, thus we must check connection manually here to know when we
# actually connect
if self.state == "connecting":
if self._asyncore_pre_connection:
self.handle_connect()
if not self.writable():
return
Expand All @@ -206,26 +205,17 @@ def handle_write(self):
return
self.sendbuf = self.sendbuf[sent:]

def send_message(self, message, pushbuf=False):
def send_message(self, message):
"""Send a P2P message over the socket.
This method takes a P2P payload, builds the P2P header and adds
the message to the send buffer to be sent over the socket."""
if self.state != "connected" and not pushbuf:
raise IOError('Not connected, no pushbuf')
if not self.is_connected:
raise IOError('Not connected')
self._log_message("send", message)
command = message.command
data = message.serialize()
tmsg = MAGIC_BYTES[self.network]
tmsg += command
tmsg += b"\x00" * (12 - len(command))
tmsg += struct.pack("<I", len(data))
th = sha256(data)
h = sha256(th)
tmsg += h[:4]
tmsg += data
tmsg = self._build_message(message)
with mininode_lock:
if (len(self.sendbuf) == 0 and not pushbuf):
if len(self.sendbuf) == 0:
try:
sent = self.send(tmsg)
self.sendbuf = tmsg[sent:]
Expand All @@ -236,6 +226,20 @@ def send_message(self, message, pushbuf=False):

# Class utility methods

def _build_message(self, message):
"""Build a serialized P2P message"""
command = message.command
data = message.serialize()
tmsg = MAGIC_BYTES[self.network]
tmsg += command
tmsg += b"\x00" * (12 - len(command))
tmsg += struct.pack("<I", len(data))
th = sha256(data)
h = sha256(th)
tmsg += h[:4]
tmsg += data
return tmsg

def _log_message(self, direction, msg):
"""Logs a message being sent or received over the connection."""
if direction == "send":
Expand Down Expand Up @@ -282,7 +286,7 @@ def peer_connect(self, *args, services=NODE_NETWORK, send_version=True, **kwargs
vt.addrTo.port = self.dstport
vt.addrFrom.ip = "0.0.0.0"
vt.addrFrom.port = 0
self.send_message(vt, True)
self.sendbuf = self._build_message(vt) # Will be sent right after handle_connect

# Message receiving methods

Expand Down Expand Up @@ -350,7 +354,7 @@ def on_version(self, message):
# Connection helper methods

def wait_for_disconnect(self, timeout=60):
test_function = lambda: self.state != "connected"
test_function = lambda: not self.is_connected
wait_until(test_function, timeout=timeout, lock=mininode_lock)

# Message receiving helper methods
Expand Down

0 comments on commit 688190c

Please sign in to comment.