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)