Agent: Refactor Log4Shell LDAP server to avoid race condition

A race condition existed between the time when the LDAP server was
instructed to start and the first exploit was sent to the victim.
Sometimes, the first exploit would be sent before the LDAP server
finished starting, resulting in failed exploitation.

To remedy this, the LDAPExploitServer.run() function now blocks until
the server has successfully started. Once the server has started,
LDAPExploitServer.run() returns. This allows the caller to have
confidence that the LDAP server is running after LDAPExploitServer.run()
returns and alleviates the need to sleep in order to avoid the race
condition.

monkey/infection_monkey/exploit/log4shell.py

@@ -34,7 +34,6 @@ class Log4ShellExploiter(WebRCE):
     _EXPLOITED_SERVICE = "Log4j"
     DOWNLOAD_TIMEOUT = 15
     REQUEST_TO_VICTIM_TIME = 5  # How long the request from victim to monkey might take. In seconds
-    LDAP_SERVER_INIT_DELAY = 5  # Time period that code halts waiting for ldap server to start
 
     def __init__(self, host: VictimHost):
         super().__init__(host)
@@ -45,7 +44,6 @@ class Log4ShellExploiter(WebRCE):
         self.class_http_server_port = get_free_tcp_port()
 
         self._ldap_server = None
-        self._ldap_server_thread = None
         self._exploit_class_http_server = None
         self._exploit_class_http_server_thread = None
         self._agent_http_server_thread = None
@@ -106,14 +104,7 @@ class Log4ShellExploiter(WebRCE):
             storage_dir=get_monkey_dir_path(),
         )
 
-        # Setting `daemon=True` to save ourselves some trouble when this is merged to the
+        self._ldap_server.run()
-        # agent-refactor branch.
-        # TODO: Make a call to `create_daemon_thread()` instead of calling the `Thread()`
-        # constructor directly after merging to the agent-refactor branch.
-        self._ldap_server_thread = Thread(target=self._ldap_server.run, daemon=True)
-        self._ldap_server_thread.start()
-        logger.debug(f"Sleeping {Log4ShellExploiter.LDAP_SERVER_INIT_DELAY} seconds for ldap process to start")
-        sleep(Log4ShellExploiter.LDAP_SERVER_INIT_DELAY)
 
     def _stop_servers(self):
         logger.debug("Stopping all LDAP and HTTP Servers")
@@ -123,8 +114,7 @@ class Log4ShellExploiter(WebRCE):
         self._exploit_class_http_server.stop()
         self._exploit_class_http_server_thread.join(Log4ShellExploiter.DOWNLOAD_TIMEOUT)
 
-        self._ldap_server.stop()
+        self._ldap_server.stop(Log4ShellExploiter.DOWNLOAD_TIMEOUT)
-        self._ldap_server_thread.join(Log4ShellExploiter.DOWNLOAD_TIMEOUT)
 
     def _build_ldap_payload(self):
         interface_ip = get_interface_to_target(self.host.ip_addr)

monkey/infection_monkey/exploit/log4shell_utils/ldap_server.py

@@ -1,6 +1,9 @@
 import logging
+import math
 import multiprocessing
 import tempfile
+import threading
+import time
 from pathlib import Path
 
 from ldaptor.interfaces import IConnectedLDAPEntry
@@ -15,6 +18,11 @@ from twisted.python.components import registerAdapter
 logger = logging.getLogger(__name__)
 
 EXPLOIT_RDN = "dn=Exploit"
+REACTOR_START_TIMEOUT_SEC = 30.0
+
+
+class LDAPServerStartError(Exception):
+    pass
 
 
 class Tree:
@@ -67,6 +75,22 @@ class LDAPExploitServer:
     def __init__(
         self, ldap_server_port: int, http_server_ip: str, http_server_port: int, storage_dir: Path
     ):
+        """
+        :param ldap_server_port: The port that the LDAP server will listen on.
+        :type ldap_server_port: int
+
+        :param http_server_ip: The IP address of the HTTP server that serves the malicious Log4Shell
+                               Java class.
+        :type http_server_ip: str
+
+        :param http_server_port: The port the HTTP server is listening on.
+        :type ldap_server_port: int
+
+        :param storage_dir: A directory where the LDAP server can safely store files it needs during
+                            runtime.
+        :type storage_dir: Path
+        """
+        self._reactor_startup_completed = multiprocessing.Event()
         self._ldap_server_port = ldap_server_port
         self._http_server_ip = http_server_ip
         self._http_server_port = http_server_port
@@ -81,14 +105,45 @@ class LDAPExploitServer:
         )
 
     def run(self):
+        """
+        Runs the Log4Shell LDAP exploit server in a subprocess. This method attempts to start the
+        server and blocks until either the server has successfully started or it times out.
+
+        :raises LDAPServerStartError: Indicates there was a problem starting the LDAP server.
+        """
+        logger.info("Starting LDAP exploit server")
         self._server_process.start()
-        self._server_process.join()
+        reactor_running = self._reactor_startup_completed.wait(REACTOR_START_TIMEOUT_SEC)
+
+        if not reactor_running:
+            raise LDAPServerStartError("An unknown error prevented the LDAP server from starting")
+
+        logger.debug("The LDAP exploit server has successfully started")
 
     def _run_twisted_reactor(self):
-        self._configure_twisted_reactor()
         logger.debug(f"Starting log4shell LDAP server on port {self._ldap_server_port}")
+        self._configure_twisted_reactor()
+
+        # Since the call to reactor.run() blocks, a separate thread is started to poll the value
+        # of `reactor.running` and set the self._reactor_startup_complete Event when the reactor
+        # is running. This allows the self.run() function to block until the reactor has
+        # successfully started.
+        threading.Thread(target=self._check_if_reactor_startup_completed, daemon=True).start()
         reactor.run()
 
+    def _check_if_reactor_startup_completed(self):
+        check_interval_sec = 0.25
+        num_checks = math.ceil(REACTOR_START_TIMEOUT_SEC / check_interval_sec)
+
+        for _ in range(0, num_checks):
+            if reactor.running:
+                logger.debug("Twisted reactor startup completed")
+                self._reactor_startup_completed.set()
+                break
+
+            logger.debug("Twisted reactor has not yet started")
+            time.sleep(check_interval_sec)
+
     def _configure_twisted_reactor(self):
         LDAPExploitServer._output_twisted_logs_to_python_logger()
 
@@ -110,7 +165,25 @@ class LDAPExploitServer:
         log_observer = log.PythonLoggingObserver()
         log_observer.start()
 
-    def stop(self):
+    def stop(self, timeout: float = None):
-        # The Twisted reactor registers signal handlers so it can catch SIGTERM and gracefully
+        """
-        # shutdown.
+        Stops the LDAP server.
-        self._server_process.terminate()
+
+        :param timeout: A floating point number of seconds to wait for the server to stop. If this
+                        argument is None (the default), the method blocks until the LDAP server
+                        terminates. If `timeout` is a positive floating point number, this method
+                        blocks for at most `timeout` seconds.
+        :type timeout: float
+        """
+        if self._server_process.is_alive():
+            logger.debug("Stopping LDAP exploit server")
+
+            # The Twisted reactor registers signal handlers so it can catch SIGTERM and gracefully
+            # shutdown.
+            self._server_process.terminate()
+            self._server_process.join(timeout)
+
+            if self._server_process.is_alive():
+                logger.warning("Timed out while waiting for the LDAP exploit server to stop.")
+            else:
+                logger.debug("Successfully stopped the LDAP exploit server")