From 7ff61423dca494b5a8dcfabd4a80076de9206576 Mon Sep 17 00:00:00 2001 From: vakarisz Date: Mon, 12 Sep 2022 12:15:07 +0300 Subject: [PATCH] Agent: Log uncaught errors in tcp_connection_handler.py thread If there's an uncaught error in the thread, the error won't propagate to the caller and we wouldn't know what went wrong from the logs. This fixes things and all uncaught errors get logged --- .../network/relay/tcp_connection_handler.py | 32 +++++++++++-------- 1 file changed, 19 insertions(+), 13 deletions(-) diff --git a/monkey/infection_monkey/network/relay/tcp_connection_handler.py b/monkey/infection_monkey/network/relay/tcp_connection_handler.py index b70a133bf..12d029d0f 100644 --- a/monkey/infection_monkey/network/relay/tcp_connection_handler.py +++ b/monkey/infection_monkey/network/relay/tcp_connection_handler.py @@ -1,3 +1,4 @@ +import logging import socket from threading import Thread from typing import Callable, List @@ -6,6 +7,8 @@ from infection_monkey.utils.threading import InterruptableThreadMixin PROXY_TIMEOUT = 2.5 +logger = logging.getLogger(__name__) + class TCPConnectionHandler(Thread, InterruptableThreadMixin): """Accepts connections on a TCP socket.""" @@ -24,18 +27,21 @@ class TCPConnectionHandler(Thread, InterruptableThreadMixin): InterruptableThreadMixin.__init__(self) def run(self): - l_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - l_socket.bind((self.bind_host, self.bind_port)) - l_socket.settimeout(PROXY_TIMEOUT) - l_socket.listen(5) + try: + l_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + l_socket.bind((self.bind_host, self.bind_port)) + l_socket.settimeout(PROXY_TIMEOUT) + l_socket.listen(5) - while not self._interrupted.is_set(): - try: - source, _ = l_socket.accept() - except socket.timeout: - continue + while not self._interrupted.is_set(): + try: + source, _ = l_socket.accept() + except socket.timeout: + continue - for notify_client_connected in self._client_connected: - notify_client_connected(source) - - l_socket.close() + for notify_client_connected in self._client_connected: + notify_client_connected(source) + except: + logging.exception("Uncaught error in TCPConnectionHandler thread") + finally: + l_socket.close()