def test_ping_handler(self):
"""
Test the PingHandler allows a ping from a connection who has not
finished authorization and returns a NetworkAck. Also test that
if that connection sends another ping in a short time period, the
connection is deemed malicious, an AuthorizationViolation is returned,
and the connection is closed.
"""
ping = PingRequest()
network = MockNetwork({},
connection_status={"connection_id":
None})
handler = PingHandler(network)
handler_status = handler.handle("connection_id",
ping.SerializeToString())
self.assertEqual(handler_status.status, HandlerStatus.RETURN)
self.assertEqual(
handler_status.message_type,
validator_pb2.Message.PING_RESPONSE)
handler_status = handler.handle("connection_id",
ping.SerializeToString())
self.assertEqual(handler_status.status, HandlerStatus.RETURN_AND_CLOSE)
self.assertEqual(
handler_status.message_type,
validator_pb2.Message.AUTHORIZATION_VIOLATION)
def _do_heartbeat(self):
ping = PingRequest()
while True:
try:
if self._socket.getsockopt(zmq.TYPE) == zmq.ROUTER:
expired = \
[ident for ident in self._last_message_times
if time.time() - self._last_message_times[ident] >
self._heartbeat_interval]
for zmq_identity in expired:
if self._is_connection_lost(
self._last_message_times[zmq_identity]):
LOGGER.debug(
"No response from %s in %s seconds"
" - removing connection.", zmq_identity,
self._connection_timeout)
self.remove_connected_identity(zmq_identity)
else:
message = validator_pb2.Message(
correlation_id=_generate_id(),
content=ping.SerializeToString(),
message_type=validator_pb2.Message.PING_REQUEST
)
fut = future.Future(
message.correlation_id,
message.content,
)
self._futures.put(fut)
message_frame = [
bytes(zmq_identity),
message.SerializeToString()
]
yield from self._send_message_frame(message_frame)
elif self._socket.getsockopt(zmq.TYPE) == zmq.DEALER:
if self._last_message_time and \
self._is_connection_lost(self._last_message_time):
LOGGER.debug(
"No response from %s in %s seconds"
" - removing connection.", self._connection,
self._connection_timeout)
connection_id = hashlib.sha512(
self.connection.encode()).hexdigest()
if connection_id in self._connections:
del self._connections[connection_id]
yield from self._stop()
yield from asyncio.sleep(self._heartbeat_interval)
except CancelledError:
# The concurrent.futures.CancelledError is caught by asyncio
# when the Task associated with the coroutine is cancelled.
# The raise is required to stop this component.
raise
except Exception as e: # pylint: disable=broad-except
LOGGER.exception(
"An error occurred while sending heartbeat: %s", e)
def do_ping_handler():
"""
Test the PingHandler returns a NetworkAck if the connection has
has finished authorization.
"""
ping = PingRequest()
network = MockNetwork(
{}, connection_status={"connection_id": ConnectionStatus.CONNECTED})
handler = PingHandler(network)
handler_status = handler.handle("connection_id", ping.SerializeToString())
return handler_status
def test_ping_handler(self):
"""
Test the PingHandler returns a NetworkAck if the connection has
has finished authorization.
"""
ping = PingRequest()
network = MockNetwork(
{},
connection_status={"connection_id": ConnectionStatus.CONNECTED})
handler = PingHandler(network)
handler_status = handler.handle("connection_id",
ping.SerializeToString())
self.assertEqual(handler_status.status, HandlerStatus.RETURN)
self.assertEqual(handler_status.message_type,
validator_pb2.Message.PING_RESPONSE)
def _do_heartbeat(self):
with self._condition:
self._condition.wait_for(lambda: self._socket is not None)
ping = PingRequest()
while True:
if self._socket.getsockopt(zmq.TYPE) == zmq.ROUTER:
expired = [
ident for ident in self._last_message_times
if time.time() -
self._last_message_times[ident] > self._heartbeat_interval
]
for zmq_identity in expired:
if self._is_connection_lost(
self._last_message_times[zmq_identity]):
LOGGER.debug(
"No response from %s in %s seconds"
" - removing connection.", zmq_identity,
self._connection_timeout)
self._remove_connected_identity(zmq_identity)
else:
message = validator_pb2.Message(
correlation_id=_generate_id(),
content=ping.SerializeToString(),
message_type=validator_pb2.Message.NETWORK_PING)
fut = future.Future(message.correlation_id,
message.content,
has_callback=False)
self._futures.put(fut)
yield from self._send_message(zmq_identity, message)
elif self._socket.getsockopt(zmq.TYPE) == zmq.DEALER:
if self._last_message_time:
if self._is_connection_lost(self._last_message_time):
LOGGER.debug(
"No response from %s in %s seconds"
" - removing connection.", self._connection,
self._connection_timeout)
yield from self._stop()
yield from asyncio.sleep(self._heartbeat_interval)
def _send_heartbeat(self):
with self._condition:
self._condition.wait_for(lambda: self._socket is not None)
ping = PingRequest()
while True:
if self._socket.getsockopt(zmq.TYPE) == zmq.ROUTER:
expired = [
ident for ident in self._connected_identities
if time.time() - self._connected_identities[ident] >
self._heartbeat_interval
]
for zmq_identity in expired:
message = validator_pb2.Message(
correlation_id=_generate_id(),
content=ping.SerializeToString(),
message_type=validator_pb2.Message.NETWORK_PING)
fut = future.Future(message.correlation_id,
message.content,
has_callback=False)
self._futures.put(fut)
yield from self._send_message(zmq_identity, message)
yield from asyncio.sleep(self._heartbeat_interval)
