From 68ad004e6d8c2faa1ea5d22afbabfde2ce47e825 Mon Sep 17 00:00:00 2001 From: Kamil Ziomek Date: Wed, 19 Aug 2026 17:49:09 +0200 Subject: [PATCH 1/2] fix: Report why rshell command timed out instead of hiding a stuck client The EFI client executes commands with a blocking os.system() call, so a tool that hangs on the DUT stops its polling loop for good and every later command times out. Track when each client last polled and which command it picked up, and return that diagnosis in the X-RShell-Timeout-Reason header. Signed-off-by: Ziomek, Kamil --- README.md | 68 +++- mfd_connect/rshell.py | 63 ++- mfd_connect/rshell_client.py | 77 +++- mfd_connect/rshell_server/rshell_server.py | 383 +++++++++++++++--- tests/unit/test_mfd_connect/test_rshell.py | 100 ++++- .../test_mfd_connect/test_rshell_server.py | 312 +++++++++++++- 6 files changed, 906 insertions(+), 97 deletions(-) diff --git a/README.md b/README.md index 6d129e1..d295c39 100644 --- a/README.md +++ b/README.md @@ -1173,6 +1173,14 @@ If code cannot establish connection, it will start deployment of python using [D * `timeout` - Timeout for command execution. * `command` - Command to be executed. * `ip` - IP address of the EFI Shell target system. + + Responses: + * `200` - command output in the body, `rc` header holds the return code. + * `400` - no command provided. + * `504` - the result did not arrive within `timeout` (+ client poll grace period). The body is empty + and `rc` is `-1`. The command is marked as *abandoned*, so it is never executed late and its result + is discarded on arrival - this keeps the caller and the EFI client in sync after a timeout. + The `X-RShell-Timeout-Reason` header explains *why* the wait failed (see below). * `/post_result` - Endpoint to post results back to the host. Headers fields: * `CommandID` - Unique identifier for the command. @@ -1184,8 +1192,56 @@ If code cannot establish connection, it will start deployment of python using [D * `CommandID` - Unique identifier for the command. Body: * Exception details. - * `/getCommandToExecute` - Endpoint to retrieve commands to be executed on the EFI Shell target system. Returns commandline with generated CommandID. - * `/health/` - Endpoint to check the health status of the connection. + * `/getCommandToExecute` - Endpoint to retrieve commands to be executed on the EFI Shell target system. Returns commandline with generated CommandID. Commands abandoned after a timeout are skipped. Every poll also refreshes the client liveness marker. + * `/health/` - Reports `200` only when the client is known **and** polled within `CLIENT_LIVENESS_TIMEOUT_SECONDS`. A client that was seen once but went silent is reported as `503`, so a dead client no longer looks healthy for the rest of the run. + * `/disconnect_client/` - Removes the client and drops everything still queued for it. + +### Diagnosing a timeout + +The EFI client runs every command through a blocking `os.system()` call and asks for new work only +after the previous command finished. A tool that hangs on the DUT - for example one that waits for a +key press because it was not started in batch/silent mode - therefore stops the polling loop for good. +From the caller side that is indistinguishable from a slow command: **every** later command times out. + +To make that visible the server tracks when each client last polled and which command it picked up, +and returns the conclusion in the `X-RShell-Timeout-Reason` header, which `RShellConnection` logs: + +| Situation | Reported reason | +| --- | --- | +| Client never contacted the server | `has never contacted the server - check that rshell_client.py is running on the DUT` | +| Client fetched this command and never came back | `picked this command up Ns ago and never reported back ... run it manually with output redirected to a file to check` | +| Client stopped polling while running an earlier command | `stopped polling Ns ago ... still blocked executing '' - the command most likely hangs on the DUT` | +| Client keeps polling, command is just slow | `is alive and polling ... consider passing a bigger timeout` | + +The second and third rows point at the DUT, not at the server - no server-side change can make a +hung EFI process return. + +### Timeouts + +`execute_command(timeout=...)` is forwarded to the server. When no timeout is given, the server falls +back to `600` seconds and the connection logs that the fallback was applied. Set `default_timeout` on +the connection for commands that legitimately run longer: + +```python +from mfd_connect import RShellConnection + +conn = RShellConnection(ip="10.10.10.10", server_ip="10.10.10.1", default_timeout=2500) +``` + +### Server resource limits + +The server keeps its in-memory state bounded, so a long test session cannot exhaust RAM: + +| Constant | Default | Meaning | +| --- | --- | --- | +| `STALE_OUTPUT_TTL_SECONDS` | `600` | Age after which an uncollected output or abandoned ID is evicted. | +| `MAX_STORED_OUTPUTS` | `512` | Hard cap on results waiting to be collected. | +| `MAX_ABANDONED_COMMAND_IDS` | `512` | Hard cap on remembered abandoned command IDs. | +| `MAX_PENDING_COMMANDS_PER_CLIENT` | `256` | Hard cap on commands queued for a single client. | +| `CLIENT_LIVENESS_TIMEOUT_SECONDS` | `300` | Silence after which `/health` reports the client as gone. | + +Waiting for a result is event driven - the caller is woken up as soon as the output is posted, +and all shared state is guarded by a lock because Werkzeug serves requests in threads. `rshell.py` is a Connection class that calls RESTful API endpoints provided by `rshell_server.py` to execute commands on the EFI Shell target system. If required, starts `rshell_server.py` on the host machine. @@ -1194,6 +1250,14 @@ RShell server can be started manually using the following command: python -m mfd_connect.rshell_server ``` +`rshell_client.py` runs on the DUT and accepts an optional source IP and source port: +```bash +rshell_client.py [source_ip] [source_port] +``` +The source port defaults to `80`; override it when that fixed local port collides with sockets left +in `TIME_WAIT` by earlier connections. A transient network error no longer terminates the client - +it retries instead, so a single glitch cannot silence the DUT for the rest of the run. + ## OS supported: * LNX * WINDOWS diff --git a/mfd_connect/rshell.py b/mfd_connect/rshell.py index 3b366b3..6793dfe 100644 --- a/mfd_connect/rshell.py +++ b/mfd_connect/rshell.py @@ -39,6 +39,12 @@ # can be adjusted based on requirements and observed behavior of platforms. PLATFORM_POWER_TRANSITION_DELAY_SECONDS = 10 +# Timeout applied by the RShell server when the caller does not provide one. +# Kept in sync with rshell_server.execute_command so the fallback can be logged explicitly. +SERVER_DEFAULT_TIMEOUT_SECONDS = 600 +# Header carrying the server side explanation of a 504. +TIMEOUT_REASON_HEADER = "X-RShell-Timeout-Reason" + class RShellConnection(Connection): """RShell Connection Class.""" @@ -50,6 +56,7 @@ def __init__( model: "BaseModel | None" = None, cache_system_data: bool = True, connection_timeout: int = 60, + default_timeout: int | None = None, ): """ Initialize RShellConnection. @@ -58,8 +65,12 @@ def __init__( :param server_ip: The IP address of the server to connect to (optional). :param model: The Pydantic model to use for the connection (optional). :param cache_system_data: Whether to cache system data (default: True). + :param connection_timeout: Time to wait for the RShell client to show up. + :param default_timeout: Timeout used by execute_command when the caller passes none. + When left as None the server applies its own default + (``SERVER_DEFAULT_TIMEOUT_SECONDS``). """ - super().__init__(model=model, cache_system_data=cache_system_data) + super().__init__(model=model, default_timeout=default_timeout, cache_system_data=cache_system_data) self._ip = ip self.server_ip = server_ip self.server_process: LocalProcess | None = None @@ -76,10 +87,14 @@ def wait_for_connection(self, connection_timeout: int) -> None: logger.log(level=log_levels.MODULE_DEBUG, msg="Checking RShell server health") try: status_code = requests.get( - f"http://{self.server_ip}/health/{self._ip}", proxies={"no_proxy": "*"} + f"http://{self.server_ip}/health/{self._ip}", + proxies={"no_proxy": "*"}, ).status_code except requests.RequestException as e: - logger.log(level=log_levels.MODULE_DEBUG, msg=f"RShell server health check failed with error: {e}") + logger.log( + level=log_levels.MODULE_DEBUG, + msg=f"RShell server health check failed with error: {e}", + ) status_code = None if status_code == 200: logger.log(level=log_levels.MODULE_DEBUG, msg="RShell server is healthy") @@ -96,7 +111,10 @@ def disconnect(self, stop_client: bool = False, stop_server: bool = False) -> No :param stop_client: Whether to stop the RShell client (default: False). """ - requests.post(f"http://{self.server_ip}/disconnect_client/{self._ip}", proxies={"no_proxy": "*"}) + requests.post( + f"http://{self.server_ip}/disconnect_client/{self._ip}", + proxies={"no_proxy": "*"}, + ) if stop_client: logger.log(level=log_levels.MODULE_DEBUG, msg="Stopping RShell client") self.execute_command("end") @@ -200,14 +218,40 @@ def execute_command( level=log_levels.MODULE_DEBUG, msg="Custom exceptions are not supported for RShellConnection and will be ignored.", ) - timeout_string = f" with timeout {timeout} seconds" if timeout is not None else "" - logger.log(level=log_levels.CMD, msg=f"Executing >{self._ip}> '{command}',{timeout_string}") + effective_timeout = timeout if timeout is not None else self.default_timeout + if effective_timeout is None: + logger.log( + level=log_levels.MODULE_DEBUG, + msg=f"No timeout given for '{command}'; the RShell server will apply its default of " + f"{SERVER_DEFAULT_TIMEOUT_SECONDS} seconds. Pass 'timeout' or set 'default_timeout' " + f"on the connection for commands that legitimately run longer.", + ) + timeout_string = f" with timeout {effective_timeout} seconds" if effective_timeout is not None else "" + logger.log( + level=log_levels.CMD, + msg=f"Executing >{self._ip}> '{command}',{timeout_string}", + ) response = requests.post( f"http://{self.server_ip}/execute_command", - data={"command": command, "timeout": timeout, "ip": self._ip}, + data={"command": command, "timeout": effective_timeout, "ip": self._ip}, proxies={"no_proxy": "*"}, ) + if response.status_code == 504: + reason = response.headers.get(TIMEOUT_REASON_HEADER) + waited = effective_timeout if effective_timeout is not None else SERVER_DEFAULT_TIMEOUT_SECONDS + message = ( + f"RShell server timed out after {waited}s waiting for the result of '{command}'. " + f"The command was dropped so it will not be executed later by the client." + ) + if reason: + message = f"{message} Reason: {reason}." + logger.log(level=log_levels.MODULE_DEBUG, msg=message) + elif response.status_code >= 500: + logger.log( + level=log_levels.MODULE_DEBUG, + msg=f"RShell server returned an internal error ({response.status_code}) for '{command}'.", + ) completed_process = ConnectionCompletedProcess( args=command, stdout=response.text, @@ -336,7 +380,10 @@ def stop_server(self) -> None: break time.sleep(1) else: - logger.log(level=log_levels.MODULE_DEBUG, msg="RShell server did not stop within timeout") + logger.log( + level=log_levels.MODULE_DEBUG, + msg="RShell server did not stop within timeout", + ) raise RuntimeError("RShell server did not stop within timeout") logger.log(level=log_levels.MODULE_DEBUG, msg="RShell server stopped") diff --git a/mfd_connect/rshell_client.py b/mfd_connect/rshell_client.py index a305c3f..5408e63 100644 --- a/mfd_connect/rshell_client.py +++ b/mfd_connect/rshell_client.py @@ -5,9 +5,18 @@ Make sure that the Python UEFI interpreter is compiled with Socket module support. + +Usage:: + + rshell_client.py [source_ip] [source_port] + +The client polls the server for work, runs one command at a time and posts the result back. +Commands are executed with a blocking ``os.system()`` call, so a tool that waits for input on +the DUT stops the whole loop - every later command then times out on the server side. Keep +that in mind when adding tools: always run them in a non interactive/batch mode. """ -__version__ = "1.1.0" +__version__ = "1.2.0" try: import httplib as client @@ -23,6 +32,16 @@ source_address = sys.argv[2] else: source_address = None +# Local port the outgoing connection is bound to. Only relevant together with source_address, +# which exists so that the server sees the expected client IP. Override it when the fixed +# port collides with sockets left in TIME_WAIT by previous connections. +if len(sys.argv) > 3: + source_port = int(sys.argv[3]) +else: + source_port = 80 + +# How long to wait before retrying after a failed server interaction. +RETRY_WAIT_SECONDS = 5 os_name = os.name @@ -41,6 +60,16 @@ def _sleep(interval): # noqa: ANN001, ANN202 time.sleep = _sleep +def _close(connection): # noqa: ANN001, ANN202 + """Close a connection without letting a broken socket kill the client.""" + if connection is None: + return + try: + connection.close() + except Exception as exp: # noqa: BLE001 + print("Ignoring error while closing the connection:", exp) + + def _get_command(): # noqa: ANN202 """Get the command from server to execute on client machine.""" # construct the list of tests by interacting with server @@ -61,27 +90,37 @@ def _get_command(): # noqa: ANN202 while True: # Connect to server - source_address_parameter = (source_address, 80) if source_address else None - conn = client.HTTPConnection(http_server, source_address=source_address_parameter) - # get the command from server - _command = _get_command() + source_address_parameter = (source_address, source_port) if source_address else None + conn = None + try: + conn = client.HTTPConnection(http_server, source_address=source_address_parameter) + # get the command from server + _command = _get_command() + except Exception as exp: # noqa: BLE001 + # A transient network error must not end the client. If it did, the DUT would stop + # asking for work and every following command would time out on the server. + print("Failed to get a command from the server:", exp) + _close(conn) + time.sleep(RETRY_WAIT_SECONDS) + continue + if not _command: - conn.close() - time.sleep(5) + _close(conn) + time.sleep(RETRY_WAIT_SECONDS) continue cmd_str, _id = _command cmd_str = cmd_str.decode("utf-8") cmd_name = cmd_str.split(" ")[0] if cmd_name == "end": print("No more commands available to run") - conn.close() + _close(conn) exit(0) print("Executing", cmd_str) if cmd_name.startswith("reset"): print("Reset command received, shutting down the platform") os.system(cmd_str) # execute reset command on machine - conn.close() + _close(conn) exit(0) non_echo = False @@ -120,15 +159,19 @@ def _get_command(): # noqa: ANN202 if non_echo and f: f.close() os.system("del " + out) - except Exception as exp: - conn.request( - "POST", - "exception", - body=cmd + str(exp), - headers={"Content-Type": "text/plain", "Connection": "keep-alive", "CommandID": _id}, - ) + except Exception as exp: # noqa: BLE001 + try: + conn.request( + "POST", + "exception", + body=cmd + str(exp), + headers={"Content-Type": "text/plain", "Connection": "keep-alive", "CommandID": _id}, + ) + except Exception as report_exp: # noqa: BLE001 + # Reporting failed too - stay alive so the next poll can still reach the server. + print("Failed to report the error to the server:", report_exp) print("output posted to server") - conn.close() + _close(conn) print("closed the connection") time.sleep(1) diff --git a/mfd_connect/rshell_server/rshell_server.py b/mfd_connect/rshell_server/rshell_server.py index 36e1646..5b0e23c 100644 --- a/mfd_connect/rshell_server/rshell_server.py +++ b/mfd_connect/rshell_server/rshell_server.py @@ -5,31 +5,235 @@ This script implements a RESTful server using Flask to manage command execution on connected RShell clients. + +Flow: + 1. ``/execute_command`` - caller queues a command and blocks until its result arrives. + 2. ``/getCommandToExecute`` - the EFI client polls for the next command to run. + 3. ``/post_result`` - the EFI client returns the command output. + 4. ``/exception`` - the EFI client reports a failure instead of an output. + +Because the EFI client polls in a slow loop (and some commands take minutes), a waiter may +give up before its result arrives. Such a command is marked as *abandoned* so that: + * it is skipped when the client asks for the next command (it is never executed late), and + * its result is dropped on arrival instead of being stored forever. + +This keeps the caller and the EFI client in sync after a timeout and keeps memory bounded. + +The server also tracks client liveness. The EFI client executes commands with a blocking +``os.system()`` call, so a command that hangs on the DUT (for example a tool waiting for a +key press) stops the polling loop for good. From the caller side this is indistinguishable +from a slow command, therefore every timeout is reported together with a diagnosis telling +whether the client kept polling or went silent while executing an earlier command. """ +import threading import time -from collections import namedtuple -from queue import Queue +from collections import OrderedDict +from queue import Empty, Queue +from typing import NamedTuple from uuid import uuid4 from flask import Flask, Response, request -__version__ = "1.1.0" +__version__ = "1.3.0" + +# How long the EFI client sleeps between polls - added to the caller timeout as a grace period. +CLIENT_LOOP_WAIT_SECONDS = 5 +# How long an output that nobody collected is kept before it is evicted. +STALE_OUTPUT_TTL_SECONDS = 600 +# Hard caps protecting the server against unbounded memory growth. +MAX_STORED_OUTPUTS = 512 +MAX_ABANDONED_COMMAND_IDS = 512 +MAX_PENDING_COMMANDS_PER_CLIENT = 256 +# Longest single wait inside get_output() - keeps the waiting loop responsive. +OUTPUT_WAIT_SLICE_SECONDS = 1.0 +# How long a client may stay silent before /health reports it as gone. The client does not poll +# while it executes a command, so this has to comfortably exceed a normal command duration. +CLIENT_LIVENESS_TIMEOUT_SECONDS = 300 +# Header carrying the human readable explanation of a 504 back to RShellConnection. +TIMEOUT_REASON_HEADER = "X-RShell-Timeout-Reason" + # Global command queue -output_object = namedtuple("OutputObject", ["output", "rc"]) -command_object = namedtuple("CommandObject", ["command_id", "str"]) +class OutputObject(NamedTuple): + """Store command output together with its return code.""" + + output: str + rc: int + + +class CommandObject(NamedTuple): + """Store queued command metadata sent to a specific client.""" + + command_id: str + str: str -output_queue: dict[str, output_object] = dict() + +class RunningCommand(NamedTuple): + """Store the command a client picked up and has not reported back yet.""" + + command_id: str + command: str + started_at: float + + +output_object = OutputObject +command_object = CommandObject + +# Results waiting to be collected by their /execute_command caller. +output_queue: "OrderedDict[str, OutputObject]" = OrderedDict() +output_queue_timestamps: dict[str, float] = dict() +# Commands whose caller already gave up - results must not be stored and they must not run. +abandoned_command_ids: "OrderedDict[str, float]" = OrderedDict() +# Per client (IP) queue of commands waiting to be picked up. command_dict_queue: dict[str, Queue] = dict() clients: list = [] +# Last time (monotonic) each client asked for work - used to tell a slow command from a dead client. +client_last_seen: dict[str, float] = dict() +# Command each client picked up last and has not reported back yet. +client_running_command: dict[str, RunningCommand] = dict() + +# Guards every structure above. Re-entrant so helpers can be called with the lock already held. +_state_lock = threading.RLock() +_output_available = threading.Condition(_state_lock) app = Flask(__name__) -def get_output(command_id: str, timeout: float = 600) -> output_object: +def _cleanup_stale_outputs(now: float | None = None, ttl: int = STALE_OUTPUT_TTL_SECONDS) -> None: + """ + Remove orphaned command outputs that have been kept longer than the configured TTL. + + Also enforces the hard caps on stored outputs and abandoned command IDs. + + :param now: Reference time (``time.monotonic()`` based). Defaults to the current time. + :param ttl: Maximum age, in seconds, of an uncollected output. + """ + with _state_lock: + current_time = time.monotonic() if now is None else now + + stale_outputs = [cid for cid, created in list(output_queue_timestamps.items()) if current_time - created >= ttl] + for command_id in stale_outputs: + output_queue.pop(command_id, None) + output_queue_timestamps.pop(command_id, None) + + stale_abandoned = [cid for cid, created in list(abandoned_command_ids.items()) if current_time - created >= ttl] + for command_id in stale_abandoned: + abandoned_command_ids.pop(command_id, None) + + while len(output_queue) > MAX_STORED_OUTPUTS: + oldest_id, _ = output_queue.popitem(last=False) + output_queue_timestamps.pop(oldest_id, None) + + while len(abandoned_command_ids) > MAX_ABANDONED_COMMAND_IDS: + abandoned_command_ids.popitem(last=False) + + +def _abandon_command(command_id: str) -> None: + """ + Mark a command as no longer awaited, so it is neither executed late nor stored on arrival. + + :param command_id: The ID of the command whose caller gave up. + """ + with _state_lock: + abandoned_command_ids[command_id] = time.monotonic() + output_queue.pop(command_id, None) + output_queue_timestamps.pop(command_id, None) + _cleanup_stale_outputs() + + +def _store_output(command_id: str, output: str, rc: int) -> bool: + """ + Persist command output together with its insertion timestamp and wake up the waiter. + + Results of abandoned commands are dropped instead of being kept forever. + + :param command_id: The ID of the command the output belongs to. + :param output: The command output. + :param rc: The return code of the command. + :return: True when the output was stored, False when it was dropped as abandoned. """ - Retrieve the output for a given command ID. + with _output_available: + if abandoned_command_ids.pop(command_id, None) is not None: + print(f"Dropping output of abandoned command {command_id} - caller already gave up") + return False + output_queue[command_id] = output_object(output=output, rc=rc) + output_queue_timestamps[command_id] = time.monotonic() + _cleanup_stale_outputs() + _output_available.notify_all() + return True + + +def _mark_client_finished(ip_address: str, command_id: str) -> None: + """ + Record that a client reported back, so it is no longer considered busy nor silent. + + :param ip_address: The IP address the result came from. + :param command_id: The ID of the command the client reported. + """ + with _state_lock: + client_last_seen[ip_address] = time.monotonic() + running = client_running_command.get(ip_address) + if running is not None and running.command_id == command_id: + client_running_command.pop(ip_address, None) + + +def diagnose_timeout(ip_address: str, command: str, waiting_since: float, command_id: str | None = None) -> str: + """ + Explain why a command timed out, based on whether the EFI client kept polling. + + The EFI client runs commands through a blocking ``os.system()`` call, so a tool that hangs + on the DUT stops the polling loop permanently. Distinguishing that from a genuinely slow + command is impossible for the caller, hence this explicit diagnosis. + + :param ip_address: The IP address of the client the command was queued for. + :param command: The command that timed out. + :param waiting_since: Monotonic timestamp of the moment the caller started waiting. + :param command_id: The ID of the command that timed out, when known. + :return: Human readable reason, safe to send as an HTTP header value. + """ + with _state_lock: + last_seen = client_last_seen.get(ip_address) + running = client_running_command.get(ip_address) + + if last_seen is None: + return ( + f"the RShell client {ip_address} has never contacted the server - " + f"check that rshell_client.py is running on the DUT" + ) + + now = time.monotonic() + + if running is not None and command_id is not None and running.command_id == command_id: + busy_for = now - running.started_at + return ( + f"the RShell client {ip_address} picked this command up {busy_for:.0f}s ago and never reported " + f"back - '{command}' is either still running or hangs on the DUT (for example waiting for " + f"input); run it manually with output redirected to a file to check" + ) + + if last_seen < waiting_since: + idle_for = now - last_seen + reason = ( + f"the RShell client {ip_address} stopped polling {idle_for:.0f}s ago, so it never asked " + f"for this command" + ) + if running is not None: + busy_for = now - running.started_at + reason += f"; it is still blocked executing '{running.command}' started {busy_for:.0f}s ago" + return reason + " - the command most likely hangs on the DUT (for example waiting for input)" + + return ( + f"the RShell client {ip_address} is alive and polling, but '{command}' did not finish in time - " + f"consider passing a bigger timeout" + ) + + +def get_output(command_id: str, timeout: float = 600) -> OutputObject: + """ + Retrieve the output for a given command ID, waiting until it arrives. + + The wait is event driven - the caller is woken up as soon as the result is posted. :param command_id: The ID of the command to retrieve output for. :param timeout: The maximum time to wait for output (in seconds). @@ -38,14 +242,18 @@ def get_output(command_id: str, timeout: float = 600) -> output_object: """ print("Getting output for command ID:", command_id) print(f"Waiting for output {timeout} seconds") - timeout = timeout + 5 # add time for client loop waiting - while timeout > 0: - result = output_queue.get(command_id, None) - if result is not None: - return result - time.sleep(1) - timeout -= 1 - raise TimeoutError("Command timed out") + deadline = time.monotonic() + timeout + CLIENT_LOOP_WAIT_SECONDS + with _output_available: + while True: + result = output_queue.pop(command_id, None) + if result is not None: + output_queue_timestamps.pop(command_id, None) + return result + remaining = deadline - time.monotonic() + if remaining <= 0: + _abandon_command(command_id) + raise TimeoutError("Command timed out") + _output_available.wait(min(remaining, OUTPUT_WAIT_SLICE_SECONDS)) def add_command_to_queue(command: str, ip_address: str) -> str: @@ -58,19 +266,39 @@ def add_command_to_queue(command: str, ip_address: str) -> str: """ print("Adding command to queue:", command) _id = str(uuid4().int) - if command_dict_queue.get(ip_address) is None: - command_dict_queue[ip_address] = Queue() - command_dict_queue[ip_address].put(command_object(command_id=_id, str=command)) + with _state_lock: + client_queue = command_dict_queue.get(ip_address) + if client_queue is None: + client_queue = Queue() + command_dict_queue[ip_address] = client_queue + while client_queue.qsize() >= MAX_PENDING_COMMANDS_PER_CLIENT: + try: + dropped = client_queue.get_nowait() + except Empty: + break + print(f"Dropping queued command {dropped.command_id} for {ip_address} - queue limit reached") + abandoned_command_ids[dropped.command_id] = time.monotonic() + client_queue.put(command_object(command_id=_id, str=command)) return _id @app.route("/health/", methods=["GET"]) def health_check(ip: str) -> Response: - """Health check endpoint.""" - if ip in clients: - return Response("OK", status=200) - else: + """ + Report whether the RShell client is connected and still polling. + + A client that was seen once but went silent is reported as gone, otherwise a dead client + would look healthy for the rest of the run. + """ + with _state_lock: + last_seen = client_last_seen.get(ip) + connected = ip in clients and last_seen is not None + if not connected or last_seen is None: return Response("Client not connected", status=503) + silent_for = time.monotonic() - last_seen + if silent_for <= CLIENT_LIVENESS_TIMEOUT_SECONDS: + return Response("OK", status=200) + return Response(f"Client stopped polling {silent_for:.0f}s ago", status=503) @app.route("/getCommandToExecute", methods=["GET"]) @@ -78,23 +306,37 @@ def get_command_to_execute() -> Response: """ Get the next command to execute for the connected client. + Commands whose caller already timed out are skipped, so the client never falls behind. + Every poll also refreshes the client liveness marker - the client only asks for work once + the previous command finished. + :return: The next command to execute. """ ip_address = str(request.remote_addr) - if ip_address not in clients: - print(f"Client connected: {ip_address}") - clients.append(ip_address) - client_queue = command_dict_queue.get(ip_address, Queue()) - if not client_queue.empty(): - command_object = client_queue.get() - return Response( - command_object.str, - status=200, - mimetype="text/plain", - headers={"CommandID": command_object.command_id}, - ) - else: - return Response("No more elements left in the queue", status=204) + with _state_lock: + if ip_address not in clients: + print(f"Client connected: {ip_address}") + clients.append(ip_address) + client_last_seen[ip_address] = time.monotonic() + client_running_command.pop(ip_address, None) + client_queue = command_dict_queue.get(ip_address) + while client_queue is not None and not client_queue.empty(): + queued_command = client_queue.get() + if abandoned_command_ids.pop(queued_command.command_id, None) is not None: + print(f"Skipping abandoned command {queued_command.command_id} for {ip_address}") + continue + client_running_command[ip_address] = RunningCommand( + command_id=queued_command.command_id, + command=queued_command.str, + started_at=time.monotonic(), + ) + return Response( + queued_command.str, + status=200, + mimetype="text/plain", + headers={"CommandID": queued_command.command_id}, + ) + return Response("No more elements left in the queue", status=204) @app.route("/exception", methods=["POST"]) @@ -110,7 +352,8 @@ def post_exception() -> Response: command_id = str(request.headers.get("CommandID")) print("CommandID: ", command_id) print(str(read_data, encoding="utf-8")) - output_queue[command_id] = output_object(output=str(read_data, encoding="utf-8"), rc=-1) + _mark_client_finished(str(request.remote_addr), command_id) + _store_output(command_id, str(read_data, encoding="utf-8"), rc=-1) return Response("Exception received", status=200) @@ -127,36 +370,67 @@ def execute_command() -> Response: timeout = int(request.form.get("timeout", 600)) command = request.form.get("command") ip_address = str(request.form.get("ip")) - if command: - _id = add_command_to_queue(command, ip_address) - if command == "end": - return Response("No more commands available to run", status=200) - if command.startswith("reset"): - return Response("Reset command sent", status=200) + if not command: + return Response("No command provided", status=400) + + _id = add_command_to_queue(command, ip_address) + if command == "end": + return Response("No more commands available to run", status=200) + if command.startswith("reset"): + return Response("Reset command sent", status=200) + + waiting_since = time.monotonic() + try: process = get_output(_id, timeout) + except TimeoutError: + # Return a clean gateway timeout instead of a Flask HTML 500 page, which the caller + # would otherwise store verbatim as the command stdout. The body stays empty on + # purpose - the explanation travels in a header so it never pollutes stdout. + reason = diagnose_timeout(ip_address, command, waiting_since, command_id=_id) + print(f"Command {_id} timed out after {timeout}s - marked as abandoned; {reason}") return Response( - process.output.encode("utf-8"), - status=200, + b"", + status=504, headers={ "Content-type": "text/plain", "CommandID": _id, - "rc": process.rc, + "rc": "-1", + TIMEOUT_REASON_HEADER: reason, }, ) - else: - return Response("No command provided", status=400) + + return Response( + process.output.encode("utf-8"), + status=200, + headers={ + "Content-type": "text/plain", + "CommandID": _id, + "rc": str(process.rc), + }, + ) @app.route("/disconnect_client/", methods=["POST"]) def disconnect_client(ip_address: str) -> Response: """ - Disconnect a client from the server. + Disconnect a client from the server and drop everything queued for it. :param ip_address: The IP address of the client to disconnect. """ - if ip_address in clients: - clients.remove(ip_address) - print(f"Client disconnected: {ip_address}") + with _state_lock: + if ip_address in clients: + clients.remove(ip_address) + client_last_seen.pop(ip_address, None) + client_running_command.pop(ip_address, None) + client_queue = command_dict_queue.pop(ip_address, None) + while client_queue is not None and not client_queue.empty(): + try: + pending = client_queue.get_nowait() + except Empty: + break + abandoned_command_ids[pending.command_id] = time.monotonic() + _cleanup_stale_outputs() + print(f"Client disconnected: {ip_address}") return Response("Client disconnected", status=200) @@ -168,7 +442,8 @@ def post_result() -> Response: rc = int(request.headers.get("rc", -1)) print("CommandID: ", command_id) print(str(read_data, encoding="utf-8")) - output_queue[command_id] = output_object(output=str(read_data, encoding="utf-8"), rc=rc) + _mark_client_finished(str(request.remote_addr), command_id) + _store_output(command_id, str(read_data, encoding="utf-8"), rc=rc) return Response("Results received", status=200) diff --git a/tests/unit/test_mfd_connect/test_rshell.py b/tests/unit/test_mfd_connect/test_rshell.py index 9a507d1..84e73ad 100644 --- a/tests/unit/test_mfd_connect/test_rshell.py +++ b/tests/unit/test_mfd_connect/test_rshell.py @@ -16,7 +16,12 @@ from mfd_connect.base import ConnectionCompletedProcess from mfd_connect.exceptions import ConnectionCalledProcessError, OsNotSupported -from mfd_connect.rshell import PLATFORM_POWER_TRANSITION_DELAY_SECONDS, RShellConnection +from mfd_connect.rshell import ( + PLATFORM_POWER_TRANSITION_DELAY_SECONDS, + SERVER_DEFAULT_TIMEOUT_SECONDS, + TIMEOUT_REASON_HEADER, + RShellConnection, +) class TestRShellConnection: @@ -30,6 +35,8 @@ def rshell(self): conn.server_ip = "127.0.0.1" conn.server_process = None conn.cache_system_data = False + # __init__ is mocked out, so the attribute normally set by Connection.__init__ is added here. + conn._default_timeout = None return conn def test_init_local_server_start(self, mocker): @@ -222,7 +229,7 @@ class _FakeBaseModel: def test_execute_command_with_all_unsupported_args_and_skip_logging(self, rshell, mocker): post = mocker.patch("mfd_connect.rshell.requests.post") - post.return_value = Mock(text="out", headers={"rc": "7"}) + post.return_value = Mock(status_code=200, text="out", headers={"rc": "7"}) result = rshell.execute_command( "echo hello", @@ -249,7 +256,10 @@ def test_execute_command_with_all_unsupported_args_and_skip_logging(self, rshell ) def test_execute_command_logs_stdout_and_default_rc(self, rshell, mocker): - mocker.patch("mfd_connect.rshell.requests.post", return_value=Mock(text="stdout", headers={})) + mocker.patch( + "mfd_connect.rshell.requests.post", + return_value=Mock(status_code=200, text="stdout", headers={}), + ) result = rshell.execute_command("echo hi") @@ -257,13 +267,95 @@ def test_execute_command_logs_stdout_and_default_rc(self, rshell, mocker): assert result.stdout == "stdout" def test_execute_command_no_stdout(self, rshell, mocker): - mocker.patch("mfd_connect.rshell.requests.post", return_value=Mock(text="", headers={"rc": "0"})) + mocker.patch( + "mfd_connect.rshell.requests.post", + return_value=Mock(status_code=200, text="", headers={"rc": "0"}), + ) result = rshell.execute_command("echo hi") assert result.return_code == 0 assert result.stdout == "" + def test_execute_command_logs_server_timeout(self, rshell, mocker, caplog): + """A 504 from the server means the command was dropped - it must be visible in the logs.""" + caplog.set_level(0) + mocker.patch( + "mfd_connect.rshell.requests.post", + return_value=Mock(status_code=504, text="", headers={"rc": "-1"}), + ) + + result = rshell.execute_command("FS0:\\Tools\\nvmupdate64e.efi /i /l") + + assert result.return_code == -1 + assert result.stdout == "" + assert "timed out" in caplog.text + + def test_execute_command_logs_server_internal_error(self, rshell, mocker, caplog): + caplog.set_level(0) + mocker.patch( + "mfd_connect.rshell.requests.post", + return_value=Mock(status_code=500, text="", headers={}), + ) + + result = rshell.execute_command("echo hi") + + assert result.return_code == -1 + assert "internal error" in caplog.text + + def test_execute_command_logs_server_timeout_reason(self, rshell, mocker, caplog): + """The server explains WHY it timed out - that diagnosis must reach the log.""" + caplog.set_level(0) + reason = "the RShell client 10.10.10.10 stopped polling 620s ago, so it never asked for this command" + mocker.patch( + "mfd_connect.rshell.requests.post", + return_value=Mock( + status_code=504, + text="", + headers={"rc": "-1", TIMEOUT_REASON_HEADER: reason}, + ), + ) + + rshell.execute_command("FS0:\\Tools\\nvmupdate64e.efi /i /l") + + assert reason in caplog.text + + def test_execute_command_without_timeout_warns_about_server_default(self, rshell, mocker, caplog): + """A command sent without a timeout is silently capped by the server - make that visible.""" + caplog.set_level(0) + post = mocker.patch( + "mfd_connect.rshell.requests.post", + return_value=Mock(status_code=200, text="out", headers={"rc": "0"}), + ) + + rshell.execute_command("FS0:\\Tools\\nvmupdate64e.efi /i /l") + + assert post.call_args.kwargs["data"]["timeout"] is None + assert str(SERVER_DEFAULT_TIMEOUT_SECONDS) in caplog.text + + def test_execute_command_uses_connection_default_timeout(self, rshell, mocker): + """default_timeout must be applied when the caller does not pass one.""" + rshell._default_timeout = 2500 + post = mocker.patch( + "mfd_connect.rshell.requests.post", + return_value=Mock(status_code=200, text="out", headers={"rc": "0"}), + ) + + rshell.execute_command("FS0:\\Tools\\nvmupdate64e.efi /i /l") + + assert post.call_args.kwargs["data"]["timeout"] == 2500 + + def test_execute_command_explicit_timeout_wins_over_default(self, rshell, mocker): + rshell._default_timeout = 2500 + post = mocker.patch( + "mfd_connect.rshell.requests.post", + return_value=Mock(status_code=200, text="out", headers={"rc": "0"}), + ) + + rshell.execute_command("ver", timeout=20) + + assert post.call_args.kwargs["data"]["timeout"] == 20 + def test_path_python_312_plus(self, rshell, monkeypatch, mocker): monkeypatch.setattr(sys, "version_info", (3, 12, 0)) factory = mocker.patch("mfd_connect.rshell.custom_path_factory", return_value="cp") diff --git a/tests/unit/test_mfd_connect/test_rshell_server.py b/tests/unit/test_mfd_connect/test_rshell_server.py index 13dfaea..e54151e 100644 --- a/tests/unit/test_mfd_connect/test_rshell_server.py +++ b/tests/unit/test_mfd_connect/test_rshell_server.py @@ -5,6 +5,8 @@ import importlib.util import runpy import sys +import threading +import time from pathlib import Path import pytest @@ -30,41 +32,64 @@ class TestRShellServerScript: def server_module(self): module = _load_server_module() module.output_queue.clear() + module.output_queue_timestamps.clear() + module.abandoned_command_ids.clear() module.command_dict_queue.clear() module.clients.clear() + module.client_last_seen.clear() + module.client_running_command.clear() return module def test_get_output_success(self, server_module): command_id = "cmd1" expected = server_module.output_object(output="hello", rc=0) server_module.output_queue[command_id] = expected + server_module.output_queue_timestamps[command_id] = 123.0 result = server_module.get_output(command_id, timeout=0) assert result == expected + assert command_id not in server_module.output_queue + assert command_id not in server_module.output_queue_timestamps def test_get_output_timeout(self, server_module): with pytest.raises(TimeoutError, match="Command timed out"): server_module.get_output("missing", timeout=-5) - def test_get_output_waits_then_returns(self, server_module, monkeypatch): - class _QueueProbe: - def __init__(self): - self.count = 0 + def test_get_output_timeout_marks_command_as_abandoned(self, server_module): + """A caller that gives up must mark its command so it is not executed/stored later.""" + with pytest.raises(TimeoutError): + server_module.get_output("gone", timeout=-5) - def get(self, _command_id, _default=None): - self.count += 1 - if self.count == 1: - return None - return server_module.output_object(output="later", rc=4) + assert "gone" in server_module.abandoned_command_ids - monkeypatch.setattr(server_module, "output_queue", _QueueProbe()) - monkeypatch.setattr(server_module.time, "sleep", lambda _x: None) + def test_get_output_is_woken_up_by_posted_result(self, server_module): + """Waiting is event driven - the waiter returns as soon as the result is stored.""" - result = server_module.get_output("cmd-later", timeout=0) + def _post_later(): + time.sleep(0.1) + server_module._store_output("cmd-later", "later", 4) + + threading.Thread(target=_post_later, daemon=True).start() + + started = time.monotonic() + result = server_module.get_output("cmd-later", timeout=10) + elapsed = time.monotonic() - started assert result.output == "later" assert result.rc == 4 + assert elapsed < 2, "waiter should be notified instead of polling" + + def test_store_output_drops_result_of_abandoned_command(self, server_module): + """Late results must not be kept in memory once nobody waits for them anymore.""" + server_module._abandon_command("dead-cmd") + + stored = server_module._store_output("dead-cmd", "late output", 0) + + assert stored is False + assert "dead-cmd" not in server_module.output_queue + assert "dead-cmd" not in server_module.output_queue_timestamps + assert "dead-cmd" not in server_module.abandoned_command_ids def test_add_command_to_queue_new_and_existing_queue(self, server_module): first_id = server_module.add_command_to_queue("echo 1", "10.0.0.1") @@ -74,6 +99,14 @@ def test_add_command_to_queue_new_and_existing_queue(self, server_module): queue_obj = server_module.command_dict_queue["10.0.0.1"] assert queue_obj.qsize() == 2 + def test_add_command_to_queue_is_bounded(self, server_module): + """Per-client queue must not grow without limits.""" + limit = server_module.MAX_PENDING_COMMANDS_PER_CLIENT + for index in range(limit + 25): + server_module.add_command_to_queue(f"echo {index}", "10.0.0.9") + + assert server_module.command_dict_queue["10.0.0.9"].qsize() <= limit + def test_health_check_endpoint(self, server_module): client = server_module.app.test_client() @@ -81,10 +114,130 @@ def test_health_check_endpoint(self, server_module): assert response_not_connected.status_code == 503 server_module.clients.append("10.0.0.1") + server_module.client_last_seen["10.0.0.1"] = time.monotonic() response_connected = client.get("/health/10.0.0.1") assert response_connected.status_code == 200 assert response_connected.get_data(as_text=True) == "OK" + def test_health_check_reports_client_that_stopped_polling(self, server_module): + """A client that was seen once but went silent must not look healthy forever.""" + client = server_module.app.test_client() + server_module.clients.append("10.0.0.1") + server_module.client_last_seen["10.0.0.1"] = ( + time.monotonic() - server_module.CLIENT_LIVENESS_TIMEOUT_SECONDS - 10 + ) + + response = client.get("/health/10.0.0.1") + + assert response.status_code == 503 + assert "stopped polling" in response.get_data(as_text=True) + + def test_health_check_requires_a_poll_before_reporting_healthy(self, server_module): + """Being present in 'clients' is not enough - the client has to have polled.""" + client = server_module.app.test_client() + server_module.clients.append("10.0.0.1") + + assert client.get("/health/10.0.0.1").status_code == 503 + + def test_poll_refreshes_liveness_and_tracks_running_command(self, server_module): + client = server_module.app.test_client() + command_id = server_module.add_command_to_queue("FS0:\\Tools\\nvmupdate64e.efi /i /l", "1.2.3.4") + + client.get("/getCommandToExecute", environ_base={"REMOTE_ADDR": "1.2.3.4"}) + + assert "1.2.3.4" in server_module.client_last_seen + running = server_module.client_running_command["1.2.3.4"] + assert running.command_id == command_id + assert running.command == "FS0:\\Tools\\nvmupdate64e.efi /i /l" + + def test_posting_result_clears_running_command(self, server_module): + client = server_module.app.test_client() + command_id = server_module.add_command_to_queue("ver", "1.2.3.4") + client.get("/getCommandToExecute", environ_base={"REMOTE_ADDR": "1.2.3.4"}) + + client.post( + "/post_result", + data=b"out", + headers={"CommandID": command_id, "rc": "0"}, + environ_base={"REMOTE_ADDR": "1.2.3.4"}, + ) + + assert "1.2.3.4" not in server_module.client_running_command + + def test_diagnose_timeout_for_unknown_client(self, server_module): + reason = server_module.diagnose_timeout("9.9.9.9", "ver", time.monotonic()) + + assert "never contacted the server" in reason + + def test_diagnose_timeout_detects_client_stuck_on_earlier_command(self, server_module): + """The exact situation from the field: the client blocks in os.system() and stops polling.""" + now = time.monotonic() + server_module.client_last_seen["1.2.3.4"] = now - 700 + server_module.client_running_command["1.2.3.4"] = server_module.RunningCommand( + command_id="old", + command="FS0:\\Tools\\nvmupdate64e.efi /i /l", + started_at=now - 700, + ) + + reason = server_module.diagnose_timeout("1.2.3.4", "ls FS0:\\Tools\\nut_upd1.cfg", waiting_since=now - 600) + + assert "stopped polling" in reason + assert "nvmupdate64e.efi /i /l" in reason + assert "hangs on the DUT" in reason + + def test_diagnose_timeout_for_slow_but_alive_client(self, server_module): + now = time.monotonic() + server_module.client_last_seen["1.2.3.4"] = now + + reason = server_module.diagnose_timeout("1.2.3.4", "slow command", waiting_since=now - 600) + + assert "alive and polling" in reason + assert "bigger timeout" in reason + + def test_diagnose_timeout_when_client_took_this_command_and_never_returned(self, server_module): + """First timeout of a hanging command: the client fetched it but never reported back.""" + now = time.monotonic() + server_module.client_last_seen["1.2.3.4"] = now - 600 + server_module.client_running_command["1.2.3.4"] = server_module.RunningCommand( + command_id="cid-1", + command="FS0:\\Tools\\nvmupdate64e.efi /i /l", + started_at=now - 600, + ) + + reason = server_module.diagnose_timeout( + "1.2.3.4", + "FS0:\\Tools\\nvmupdate64e.efi /i /l", + waiting_since=now - 605, + command_id="cid-1", + ) + + assert "picked this command up" in reason + assert "never reported back" in reason + + def test_execute_command_timeout_sends_reason_header_and_empty_body(self, server_module): + """The body must stay empty so the caller never stores an error page as stdout.""" + client = server_module.app.test_client() + + response = client.post("/execute_command", data={"command": "ver", "timeout": "-5", "ip": "1.2.3.4"}) + + assert response.status_code == 504 + assert response.get_data() == b"" + assert response.headers["rc"] == "-1" + assert response.headers[server_module.TIMEOUT_REASON_HEADER] + + def test_disconnect_client_clears_liveness_state(self, server_module): + """A reconnecting client must not inherit the liveness state of the previous session.""" + client = server_module.app.test_client() + server_module.add_command_to_queue("ver", "1.2.3.4") + client.get("/getCommandToExecute", environ_base={"REMOTE_ADDR": "1.2.3.4"}) + assert "1.2.3.4" in server_module.client_last_seen + + client.post("/disconnect_client/1.2.3.4") + + assert "1.2.3.4" not in server_module.client_last_seen + assert "1.2.3.4" not in server_module.client_running_command + assert client.get("/health/1.2.3.4").status_code == 503 + def test_get_command_to_execute_endpoint(self, server_module): client = server_module.app.test_client() @@ -99,6 +252,29 @@ def test_get_command_to_execute_endpoint(self, server_module): assert response_with_command.get_data(as_text=True) == "echo hi" assert response_with_command.headers["CommandID"] == command_id + def test_get_command_to_execute_skips_abandoned_commands(self, server_module): + """Abandoned commands must never reach the client, otherwise it stays one command behind.""" + client = server_module.app.test_client() + abandoned_id = server_module.add_command_to_queue("slow command", "1.2.3.4") + server_module._abandon_command(abandoned_id) + live_id = server_module.add_command_to_queue("echo alive", "1.2.3.4") + + response = client.get("/getCommandToExecute", environ_base={"REMOTE_ADDR": "1.2.3.4"}) + + assert response.status_code == 200 + assert response.headers["CommandID"] == live_id + assert response.get_data(as_text=True) == "echo alive" + assert abandoned_id not in server_module.abandoned_command_ids + + def test_get_command_to_execute_returns_204_when_only_abandoned(self, server_module): + client = server_module.app.test_client() + abandoned_id = server_module.add_command_to_queue("slow command", "1.2.3.4") + server_module._abandon_command(abandoned_id) + + response = client.get("/getCommandToExecute", environ_base={"REMOTE_ADDR": "1.2.3.4"}) + + assert response.status_code == 204 + def test_post_exception_endpoint(self, server_module): client = server_module.app.test_client() response = client.post("/exception", data=b"boom", headers={"CommandID": "cid-1"}) @@ -106,6 +282,7 @@ def test_post_exception_endpoint(self, server_module): assert response.status_code == 200 assert server_module.output_queue["cid-1"].output == "boom" assert server_module.output_queue["cid-1"].rc == -1 + assert "cid-1" in server_module.output_queue_timestamps def test_execute_command_endpoint_paths(self, server_module, monkeypatch): client = server_module.app.test_client() @@ -140,17 +317,70 @@ def test_execute_command_endpoint_paths(self, server_module, monkeypatch): assert response_normal.headers["Content-type"].startswith("text/plain") assert response_normal.headers["CommandID"] + def test_execute_command_returns_gateway_timeout_instead_of_html_error(self, server_module): + """A timeout must not return a Flask HTML 500 page, which the caller stores as stdout.""" + client = server_module.app.test_client() + + response = client.post( + "/execute_command", + data={ + "command": "FS0:\\Tools\\nvmupdate64e.efi /i /l", + "timeout": "-5", + "ip": "1.1.1.1", + }, + ) + + assert response.status_code == 504 + assert response.get_data(as_text=True) == "" + assert response.headers["rc"] == "-1" + assert "" not in response.get_data(as_text=True) + + def test_timed_out_command_does_not_desync_next_command(self, server_module): + """After a timeout the next command must succeed - the server has to re-sync.""" + client = server_module.app.test_client() + ip = "10.102.23.150" + + timed_out = client.post( + "/execute_command", + data={"command": "slow", "timeout": "-5", "ip": ip}, + ) + assert timed_out.status_code == 504 + + # The stale command must not be handed out to the client anymore. + assert client.get("/getCommandToExecute", environ_base={"REMOTE_ADDR": ip}).status_code == 204 + + next_id = server_module.add_command_to_queue("ver", ip) + handed_out = client.get("/getCommandToExecute", environ_base={"REMOTE_ADDR": ip}) + assert handed_out.status_code == 200 + assert handed_out.headers["CommandID"] == next_id + + server_module._store_output(next_id, "UEFI Shell", 0) + assert server_module.get_output(next_id, timeout=1).output == "UEFI Shell" + def test_disconnect_client_endpoint(self, server_module): client = server_module.app.test_client() server_module.clients.append("2.2.2.2") + server_module.add_command_to_queue("echo hi", "2.2.2.2") response_existing = client.post("/disconnect_client/2.2.2.2") assert response_existing.status_code == 200 assert "2.2.2.2" not in server_module.clients + assert "2.2.2.2" not in server_module.command_dict_queue response_missing = client.post("/disconnect_client/8.8.8.8") assert response_missing.status_code == 200 + def test_disconnect_client_abandons_pending_commands(self, server_module): + """Pending commands of a disconnected client must not be executed after reconnect.""" + client = server_module.app.test_client() + server_module.clients.append("3.3.3.3") + pending_id = server_module.add_command_to_queue("echo hi", "3.3.3.3") + + client.post("/disconnect_client/3.3.3.3") + + assert pending_id in server_module.abandoned_command_ids + assert server_module._store_output(pending_id, "late", 0) is False + def test_post_result_endpoint(self, server_module): client = server_module.app.test_client() @@ -158,6 +388,7 @@ def test_post_result_endpoint(self, server_module): assert response_default_rc.status_code == 200 assert server_module.output_queue["cmd-a"].output == "output-a" assert server_module.output_queue["cmd-a"].rc == -1 + assert "cmd-a" in server_module.output_queue_timestamps response_given_rc = client.post( "/post_result", @@ -167,6 +398,63 @@ def test_post_result_endpoint(self, server_module): assert response_given_rc.status_code == 200 assert server_module.output_queue["cmd-b"].output == "output-b" assert server_module.output_queue["cmd-b"].rc == 3 + assert "cmd-b" in server_module.output_queue_timestamps + + def test_cleanup_stale_outputs_removes_expired_entries_only(self, server_module): + # Timestamps simulate time.monotonic() values (seconds since arbitrary boot reference). + # now=4000 s, ttl=3600 s -> stale (t=10) is evicted, fresh (t=500) is kept. + server_module.output_queue["stale"] = server_module.output_object(output="old", rc=-1) + server_module.output_queue_timestamps["stale"] = 10.0 + server_module.output_queue["fresh"] = server_module.output_object(output="new", rc=0) + server_module.output_queue_timestamps["fresh"] = 500.0 + + server_module._cleanup_stale_outputs(now=4000.0, ttl=3600) + + assert "stale" not in server_module.output_queue + assert "stale" not in server_module.output_queue_timestamps + assert server_module.output_queue["fresh"].output == "new" + + def test_cleanup_stale_outputs_expires_abandoned_ids(self, server_module): + server_module.abandoned_command_ids["old"] = 10.0 + server_module.abandoned_command_ids["recent"] = 3900.0 + + server_module._cleanup_stale_outputs(now=4000.0, ttl=600) + + assert "old" not in server_module.abandoned_command_ids + assert "recent" in server_module.abandoned_command_ids + + def test_output_queue_is_hard_capped(self, server_module): + """Even without TTL expiry the stored outputs must stay bounded.""" + limit = server_module.MAX_STORED_OUTPUTS + for index in range(limit + 120): + server_module._store_output(f"cid-{index}", "payload", 0) + + assert len(server_module.output_queue) <= limit + assert len(server_module.output_queue_timestamps) == len(server_module.output_queue) + + def test_concurrent_access_is_thread_safe(self, server_module): + """Werkzeug serves requests in threads - shared state must not raise or corrupt.""" + errors = [] + + def _worker(worker_id): + try: + for index in range(30): + command_id = f"w{worker_id}-{index}" + server_module._store_output(command_id, "data", 0) + server_module.get_output(command_id, timeout=1) + server_module.add_command_to_queue(f"cmd{index}", f"10.0.1.{worker_id}") + server_module._cleanup_stale_outputs() + except Exception as exc: # noqa: BLE001 + errors.append(f"{type(exc).__name__}: {exc}") + + threads = [threading.Thread(target=_worker, args=(worker_id,)) for worker_id in range(12)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert not errors + assert not server_module.output_queue def test_run_function_starts_flask(self, server_module, monkeypatch): """Test that the run() function starts Flask with correct host and port.""" From c3232d0ab2e2966317222aa0ae5f3ec4b4d131dc Mon Sep 17 00:00:00 2001 From: Arkadiusz Baczek <86051266+abaczek@users.noreply.github.com> Date: Wed, 14 Oct 2026 12:30:18 +0000 Subject: [PATCH 2/2] fix: Bound rshell result storage and clean up completed commands Signed-off-by: Arkadiusz Baczek <86051266+abaczek@users.noreply.github.com> --- README.md | 29 +- mfd_connect/rshell.py | 8 +- mfd_connect/rshell_client.py | 9 +- mfd_connect/rshell_server/rshell_server.py | 204 ++++++--- tests/unit/test_mfd_connect/test_rshell.py | 14 +- .../test_mfd_connect/test_rshell_client.py | 14 + .../test_mfd_connect/test_rshell_server.py | 400 ++++++++++++++++-- 7 files changed, 590 insertions(+), 88 deletions(-) diff --git a/README.md b/README.md index d295c39..f2813ae 100644 --- a/README.md +++ b/README.md @@ -1177,9 +1177,12 @@ If code cannot establish connection, it will start deployment of python using [D Responses: * `200` - command output in the body, `rc` header holds the return code. * `400` - no command provided. + * `413` - command output exceeds 16 MiB. + * `429` - a client or command admission limit was reached; existing commands are not evicted. + * `503` - pending results exceed the storage limit. * `504` - the result did not arrive within `timeout` (+ client poll grace period). The body is empty - and `rc` is `-1`. The command is marked as *abandoned*, so it is never executed late and its result - is discarded on arrival - this keeps the caller and the EFI client in sync after a timeout. + and `rc` is `-1`. Queued commands are removed and late results are discarded. + A command already picked up by the client may still execute; a timeout does not stop it. The `X-RShell-Timeout-Reason` header explains *why* the wait failed (see below). * `/post_result` - Endpoint to post results back to the host. Headers fields: @@ -1196,6 +1199,20 @@ If code cannot establish connection, it will start deployment of python using [D * `/health/` - Reports `200` only when the client is known **and** polled within `CLIENT_LIVENESS_TIMEOUT_SECONDS`. A client that was seen once but went silent is reported as `503`, so a dead client no longer looks healthy for the rest of the run. * `/disconnect_client/` - Removes the client and drops everything still queued for it. +Results are removed after collection. Unknown, duplicate and abandoned results are not retained. +The server accepts up to 256 clients, 512 active commands and 256 queued commands per client. +Stored result payloads are limited to 64 MiB in total and 16 MiB per result, measured as UTF-8 bytes. +Overflow is reported without truncating results or evicting another caller's result. +Disconnecting clears the client's pending results and wakes its waiting callers. +Idle client records are reclaimed during subsequent requests; actively awaited commands keep their client records. +Unfinished-command diagnostics are retained separately, up to 256 records. Records without an active waiter +expire one hour after pickup; the oldest such records may be evicted earlier to enforce the cap. +Receiving the matching result, a new poll, or an explicit disconnect clears the diagnostic record. + +HTTP errors other than `504` raise `requests.HTTPError` in `RShellConnection.execute_command()`. +Server-generated HTTP errors have an empty body, `rc: -1` and an `X-RShell-Error` header; +they are not returned as successful command output. + ### Diagnosing a timeout The EFI client runs every command through a blocking `os.system()` call and asks for new work only @@ -1206,6 +1223,8 @@ From the caller side that is indistinguishable from a slow command: **every** la To make that visible the server tracks when each client last polled and which command it picked up, and returns the conclusion in the `X-RShell-Timeout-Reason` header, which `RShellConnection` logs: +The header value is escaped to single-line ASCII and limited to 1024 characters. + | Situation | Reported reason | | --- | --- | | Client never contacted the server | `has never contacted the server - check that rshell_client.py is running on the DUT` | @@ -1228,6 +1247,10 @@ from mfd_connect import RShellConnection conn = RShellConnection(ip="10.10.10.10", server_ip="10.10.10.1", default_timeout=2500) ``` +The fire-and-forget commands `end` and `reset*` return immediately, but must be picked up within +the requested timeout (default 600 seconds) plus the 5-second client poll grace period. +Undelivered commands are removed during subsequent requests, releasing their command and client slots. + ### Server resource limits The server keeps its in-memory state bounded, so a long test session cannot exhaust RAM: @@ -1239,6 +1262,8 @@ The server keeps its in-memory state bounded, so a long test session cannot exha | `MAX_ABANDONED_COMMAND_IDS` | `512` | Hard cap on remembered abandoned command IDs. | | `MAX_PENDING_COMMANDS_PER_CLIENT` | `256` | Hard cap on commands queued for a single client. | | `CLIENT_LIVENESS_TIMEOUT_SECONDS` | `300` | Silence after which `/health` reports the client as gone. | +| `RUNNING_COMMAND_TTL_SECONDS` | `3600` | Age since pickup after which an unfinished-command diagnostic without an active waiter is evicted. | +| `MAX_TIMEOUT_REASON_LENGTH` | `1024` | Maximum length of the ASCII-escaped timeout reason header. | Waiting for a result is event driven - the caller is woken up as soon as the output is posted, and all shared state is guarded by a lock because Werkzeug serves requests in threads. diff --git a/mfd_connect/rshell.py b/mfd_connect/rshell.py index 6793dfe..c5de937 100644 --- a/mfd_connect/rshell.py +++ b/mfd_connect/rshell.py @@ -242,16 +242,18 @@ def execute_command( waited = effective_timeout if effective_timeout is not None else SERVER_DEFAULT_TIMEOUT_SECONDS message = ( f"RShell server timed out after {waited}s waiting for the result of '{command}'. " - f"The command was dropped so it will not be executed later by the client." + f"Any queued command was removed; a command already picked up by the client may still be running." ) if reason: message = f"{message} Reason: {reason}." logger.log(level=log_levels.MODULE_DEBUG, msg=message) - elif response.status_code >= 500: + elif response.status_code >= 400: logger.log( level=log_levels.MODULE_DEBUG, - msg=f"RShell server returned an internal error ({response.status_code}) for '{command}'.", + msg=f"RShell request failed ({response.status_code}): " + f"{response.headers.get('X-RShell-Error', response.reason)}", ) + response.raise_for_status() completed_process = ConnectionCompletedProcess( args=command, stdout=response.text, diff --git a/mfd_connect/rshell_client.py b/mfd_connect/rshell_client.py index 5408e63..0e8def1 100644 --- a/mfd_connect/rshell_client.py +++ b/mfd_connect/rshell_client.py @@ -61,7 +61,11 @@ def _sleep(interval): # noqa: ANN001, ANN202 def _close(connection): # noqa: ANN001, ANN202 - """Close a connection without letting a broken socket kill the client.""" + """ + Close a connection without letting a broken socket kill the client. + + :param connection: HTTP connection to close, or None when no connection was created. + """ if connection is None: return try: @@ -79,6 +83,9 @@ def _get_command(): # noqa: ANN202 _id = rsp.getheader("CommandID") if status == 204: return None + if status != 200: + rsp.read() + raise RuntimeError("RShell command request failed with HTTP {}".format(status)) print("Waiting for command from server: ") data_received = rsp.read() diff --git a/mfd_connect/rshell_server/rshell_server.py b/mfd_connect/rshell_server/rshell_server.py index 5b0e23c..ddaf0e3 100644 --- a/mfd_connect/rshell_server/rshell_server.py +++ b/mfd_connect/rshell_server/rshell_server.py @@ -29,11 +29,12 @@ import threading import time from collections import OrderedDict -from queue import Empty, Queue -from typing import NamedTuple +from queue import Queue +from typing import NamedTuple, NoReturn from uuid import uuid4 -from flask import Flask, Response, request +from flask import Flask, Response, abort, request +from werkzeug.exceptions import HTTPException, RequestEntityTooLarge, ServiceUnavailable, TooManyRequests __version__ = "1.3.0" @@ -45,13 +46,20 @@ MAX_STORED_OUTPUTS = 512 MAX_ABANDONED_COMMAND_IDS = 512 MAX_PENDING_COMMANDS_PER_CLIENT = 256 +MAX_ACTIVE_COMMANDS = 512 +MAX_CLIENTS = 256 +MAX_OUTPUT_BYTES = 16 * 1024 * 1024 +MAX_STORED_OUTPUT_BYTES = 64 * 1024 * 1024 # Longest single wait inside get_output() - keeps the waiting loop responsive. OUTPUT_WAIT_SLICE_SECONDS = 1.0 # How long a client may stay silent before /health reports it as gone. The client does not poll # while it executes a command, so this has to comfortably exceed a normal command duration. CLIENT_LIVENESS_TIMEOUT_SECONDS = 300 +# Keep unfinished-command diagnostics after their waiter and client admission slot are gone. +RUNNING_COMMAND_TTL_SECONDS = 3600 # Header carrying the human readable explanation of a 504 back to RShellConnection. TIMEOUT_REASON_HEADER = "X-RShell-Timeout-Reason" +MAX_TIMEOUT_REASON_LENGTH = 1024 # Global command queue @@ -83,6 +91,10 @@ class RunningCommand(NamedTuple): # Results waiting to be collected by their /execute_command caller. output_queue: "OrderedDict[str, OutputObject]" = OrderedDict() output_queue_timestamps: dict[str, float] = dict() +output_sizes: dict[str, int] = dict() +output_errors: dict[str, tuple[int, str]] = dict() +active_commands: dict[str, str] = dict() +control_command_deadlines: dict[str, float] = dict() # Commands whose caller already gave up - results must not be stored and they must not run. abandoned_command_ids: "OrderedDict[str, float]" = OrderedDict() # Per client (IP) queue of commands waiting to be picked up. @@ -98,13 +110,15 @@ class RunningCommand(NamedTuple): _output_available = threading.Condition(_state_lock) app = Flask(__name__) +app.config["MAX_CONTENT_LENGTH"] = MAX_OUTPUT_BYTES def _cleanup_stale_outputs(now: float | None = None, ttl: int = STALE_OUTPUT_TTL_SECONDS) -> None: """ Remove orphaned command outputs that have been kept longer than the configured TTL. - Also enforces the hard caps on stored outputs and abandoned command IDs. + Also expires undelivered control commands and bounds retained running-command diagnostics, + stored outputs and abandoned command IDs. :param now: Reference time (``time.monotonic()`` based). Defaults to the current time. :param ttl: Maximum age, in seconds, of an uncollected output. @@ -112,22 +126,40 @@ def _cleanup_stale_outputs(now: float | None = None, ttl: int = STALE_OUTPUT_TTL with _state_lock: current_time = time.monotonic() if now is None else now + for command_id, deadline in list(control_command_deadlines.items()): + if current_time >= deadline: + print(f"Control command {command_id} expired before delivery") + _abandon_command(command_id) + stale_outputs = [cid for cid, created in list(output_queue_timestamps.items()) if current_time - created >= ttl] for command_id in stale_outputs: - output_queue.pop(command_id, None) - output_queue_timestamps.pop(command_id, None) + _abandon_command(command_id) stale_abandoned = [cid for cid, created in list(abandoned_command_ids.items()) if current_time - created >= ttl] for command_id in stale_abandoned: abandoned_command_ids.pop(command_id, None) while len(output_queue) > MAX_STORED_OUTPUTS: - oldest_id, _ = output_queue.popitem(last=False) - output_queue_timestamps.pop(oldest_id, None) + _abandon_command(next(iter(output_queue))) while len(abandoned_command_ids) > MAX_ABANDONED_COMMAND_IDS: abandoned_command_ids.popitem(last=False) + for ip, running in list(client_running_command.items()): + if running.command_id not in active_commands and ( + current_time - running.started_at >= RUNNING_COMMAND_TTL_SECONDS + or len(client_running_command) > MAX_CLIENTS + ): + client_running_command.pop(ip, None) + + busy_clients = set(active_commands.values()) + for ip, last_seen in list(client_last_seen.items()): + if current_time - last_seen > CLIENT_LIVENESS_TIMEOUT_SECONDS and ip not in busy_clients: + client_last_seen.pop(ip, None) + command_dict_queue.pop(ip, None) + if ip in clients: + clients.remove(ip) + def _abandon_command(command_id: str) -> None: """ @@ -135,11 +167,28 @@ def _abandon_command(command_id: str) -> None: :param command_id: The ID of the command whose caller gave up. """ - with _state_lock: + with _output_available: abandoned_command_ids[command_id] = time.monotonic() + ip = active_commands.pop(command_id, None) + control_command_deadlines.pop(command_id, None) output_queue.pop(command_id, None) output_queue_timestamps.pop(command_id, None) - _cleanup_stale_outputs() + output_sizes.pop(command_id, None) + output_errors.pop(command_id, None) + client_queue = command_dict_queue.get(ip) + if client_queue is not None: + pending = [] + while not client_queue.empty(): + queued = client_queue.get_nowait() + if queued.command_id != command_id: + pending.append(queued) + for queued in pending: + client_queue.put_nowait(queued) + if client_queue.empty(): + command_dict_queue.pop(ip, None) + while len(abandoned_command_ids) > MAX_ABANDONED_COMMAND_IDS: + abandoned_command_ids.popitem(last=False) + _output_available.notify_all() def _store_output(command_id: str, output: str, rc: int) -> bool: @@ -154,16 +203,40 @@ def _store_output(command_id: str, output: str, rc: int) -> bool: :return: True when the output was stored, False when it was dropped as abandoned. """ with _output_available: - if abandoned_command_ids.pop(command_id, None) is not None: - print(f"Dropping output of abandoned command {command_id} - caller already gave up") + _cleanup_stale_outputs() + if command_id not in active_commands or command_id in output_queue or command_id in output_errors: + abandoned_command_ids.pop(command_id, None) + print(f"Dropping unsolicited or duplicate output for command {command_id}") return False + size = len(output.encode("utf-8")) + if size > MAX_OUTPUT_BYTES: + _fail_output(command_id, RequestEntityTooLarge(description="RShell command output exceeds the size limit")) + if len(output_queue) >= MAX_STORED_OUTPUTS or sum(output_sizes.values()) + size > MAX_STORED_OUTPUT_BYTES: + _fail_output(command_id, ServiceUnavailable(description="RShell result storage is full")) output_queue[command_id] = output_object(output=output, rc=rc) output_queue_timestamps[command_id] = time.monotonic() + output_sizes[command_id] = size _cleanup_stale_outputs() _output_available.notify_all() return True +def _fail_output(command_id: str, error: HTTPException) -> NoReturn: + """Reject an output and notify its waiting caller without retaining the payload.""" + with _output_available: + if command_id in active_commands and command_id not in output_queue: + output_errors[command_id] = (error.get_response().status_code, error.description) + _output_available.notify_all() + raise error + + +def _check_client_limit(ip_address: str) -> None: + """Reject new clients when all client slots are occupied.""" + known_clients = set(clients) | set(client_last_seen) | set(active_commands.values()) + if ip_address not in known_clients and len(known_clients) >= MAX_CLIENTS: + raise TooManyRequests(description="RShell client limit reached") + + def _mark_client_finished(ip_address: str, command_id: str) -> None: """ Record that a client reported back, so it is no longer considered busy nor silent. @@ -172,10 +245,11 @@ def _mark_client_finished(ip_address: str, command_id: str) -> None: :param command_id: The ID of the command the client reported. """ with _state_lock: - client_last_seen[ip_address] = time.monotonic() running = client_running_command.get(ip_address) if running is not None and running.command_id == command_id: client_running_command.pop(ip_address, None) + if ip_address in clients: + client_last_seen[ip_address] = time.monotonic() def diagnose_timeout(ip_address: str, command: str, waiting_since: float, command_id: str | None = None) -> str: @@ -190,11 +264,13 @@ def diagnose_timeout(ip_address: str, command: str, waiting_since: float, comman :param command: The command that timed out. :param waiting_since: Monotonic timestamp of the moment the caller started waiting. :param command_id: The ID of the command that timed out, when known. - :return: Human readable reason, safe to send as an HTTP header value. + :return: Human readable reason. """ with _state_lock: last_seen = client_last_seen.get(ip_address) running = client_running_command.get(ip_address) + if last_seen is None and running is not None: + last_seen = running.started_at if last_seen is None: return ( @@ -245,10 +321,17 @@ def get_output(command_id: str, timeout: float = 600) -> OutputObject: deadline = time.monotonic() + timeout + CLIENT_LOOP_WAIT_SECONDS with _output_available: while True: + error = output_errors.pop(command_id, None) + if error is not None: + _abandon_command(command_id) + abort(error[0], description=error[1]) result = output_queue.pop(command_id, None) if result is not None: - output_queue_timestamps.pop(command_id, None) + _abandon_command(command_id) return result + if command_id not in active_commands: + _abandon_command(command_id) + raise TimeoutError("Command is no longer pending") remaining = deadline - time.monotonic() if remaining <= 0: _abandon_command(command_id) @@ -256,28 +339,31 @@ def get_output(command_id: str, timeout: float = 600) -> OutputObject: _output_available.wait(min(remaining, OUTPUT_WAIT_SLICE_SECONDS)) -def add_command_to_queue(command: str, ip_address: str) -> str: +def add_command_to_queue(command: str, ip_address: str, timeout: float = 600) -> str: """ Add a command to the global command queue. :param command: The command to add to the queue. :param ip_address: The IP address of the client. + :param timeout: Delivery timeout for control commands, in seconds, before the client poll grace period. :return: The ID of the added command. """ print("Adding command to queue:", command) _id = str(uuid4().int) with _state_lock: + _cleanup_stale_outputs() + _check_client_limit(ip_address) + if len(active_commands) >= MAX_ACTIVE_COMMANDS: + raise TooManyRequests(description="RShell active command limit reached") client_queue = command_dict_queue.get(ip_address) if client_queue is None: client_queue = Queue() command_dict_queue[ip_address] = client_queue - while client_queue.qsize() >= MAX_PENDING_COMMANDS_PER_CLIENT: - try: - dropped = client_queue.get_nowait() - except Empty: - break - print(f"Dropping queued command {dropped.command_id} for {ip_address} - queue limit reached") - abandoned_command_ids[dropped.command_id] = time.monotonic() + if client_queue.qsize() >= MAX_PENDING_COMMANDS_PER_CLIENT: + raise TooManyRequests(description="RShell client command queue is full") + active_commands[_id] = ip_address + if command == "end" or command.startswith("reset"): + control_command_deadlines[_id] = time.monotonic() + timeout + CLIENT_LOOP_WAIT_SECONDS client_queue.put(command_object(command_id=_id, str=command)) return _id @@ -314,6 +400,8 @@ def get_command_to_execute() -> Response: """ ip_address = str(request.remote_addr) with _state_lock: + _cleanup_stale_outputs() + _check_client_limit(ip_address) if ip_address not in clients: print(f"Client connected: {ip_address}") clients.append(ip_address) @@ -322,14 +410,21 @@ def get_command_to_execute() -> Response: client_queue = command_dict_queue.get(ip_address) while client_queue is not None and not client_queue.empty(): queued_command = client_queue.get() - if abandoned_command_ids.pop(queued_command.command_id, None) is not None: + if queued_command.command_id not in active_commands: + abandoned_command_ids.pop(queued_command.command_id, None) print(f"Skipping abandoned command {queued_command.command_id} for {ip_address}") continue + if client_queue.empty(): + command_dict_queue.pop(ip_address, None) + if queued_command.str == "end" or queued_command.str.startswith("reset"): + active_commands.pop(queued_command.command_id, None) + control_command_deadlines.pop(queued_command.command_id, None) client_running_command[ip_address] = RunningCommand( command_id=queued_command.command_id, command=queued_command.str, started_at=time.monotonic(), ) + _cleanup_stale_outputs() return Response( queued_command.str, status=200, @@ -348,13 +443,32 @@ def post_exception() -> Response: :param CommandID: The ID of the command that caused the exception. :return: A response indicating the exception was received. """ - read_data = request.data + return _receive_output(-1, "Exception received") + + +def _receive_output(rc: int, message: str) -> Response: + """Receive a size-limited result or exception for a pending command.""" command_id = str(request.headers.get("CommandID")) + try: + read_data = request.data + except RequestEntityTooLarge as error: + _fail_output(command_id, error) print("CommandID: ", command_id) print(str(read_data, encoding="utf-8")) _mark_client_finished(str(request.remote_addr), command_id) - _store_output(command_id, str(read_data, encoding="utf-8"), rc=-1) - return Response("Exception received", status=200) + stored = _store_output(command_id, str(read_data, encoding="utf-8"), rc=rc) + return Response(message if stored else "Output no longer awaited", status=200) + + +@app.errorhandler(HTTPException) +def report_http_error(error: HTTPException) -> Response: + """Report transport failures separately from command stdout.""" + app.logger.warning("RShell request failed: %s", error.description) + return Response( + b"", + status=error.code, + headers={"rc": "-1", "X-RShell-Error": error.description}, + ) @app.route("/execute_command", methods=["POST"]) @@ -373,7 +487,7 @@ def execute_command() -> Response: if not command: return Response("No command provided", status=400) - _id = add_command_to_queue(command, ip_address) + _id = add_command_to_queue(command, ip_address, timeout) if command == "end": return Response("No more commands available to run", status=200) if command.startswith("reset"): @@ -395,9 +509,12 @@ def execute_command() -> Response: "Content-type": "text/plain", "CommandID": _id, "rc": "-1", - TIMEOUT_REASON_HEADER: reason, + TIMEOUT_REASON_HEADER: reason.encode("unicode_escape").decode("ascii")[:MAX_TIMEOUT_REASON_LENGTH], }, ) + finally: + with _state_lock: + _abandon_command(_id) return Response( process.output.encode("utf-8"), @@ -417,34 +534,25 @@ def disconnect_client(ip_address: str) -> Response: :param ip_address: The IP address of the client to disconnect. """ - with _state_lock: + with _output_available: if ip_address in clients: clients.remove(ip_address) - client_last_seen.pop(ip_address, None) - client_running_command.pop(ip_address, None) - client_queue = command_dict_queue.pop(ip_address, None) - while client_queue is not None and not client_queue.empty(): - try: - pending = client_queue.get_nowait() - except Empty: - break - abandoned_command_ids[pending.command_id] = time.monotonic() - _cleanup_stale_outputs() - print(f"Client disconnected: {ip_address}") + client_last_seen.pop(ip_address, None) + client_running_command.pop(ip_address, None) + for command_id, ip in list(active_commands.items()): + if ip == ip_address: + _abandon_command(command_id) + command_dict_queue.pop(ip_address, None) + _cleanup_stale_outputs() + print(f"Client disconnected: {ip_address}") return Response("Client disconnected", status=200) @app.route("/post_result", methods=["POST"]) def post_result() -> Response: """Receive command execution results from the client.""" - read_data = request.data - command_id = str(request.headers.get("CommandID")) rc = int(request.headers.get("rc", -1)) - print("CommandID: ", command_id) - print(str(read_data, encoding="utf-8")) - _mark_client_finished(str(request.remote_addr), command_id) - _store_output(command_id, str(read_data, encoding="utf-8"), rc=rc) - return Response("Results received", status=200) + return _receive_output(rc, "Results received") def run() -> None: diff --git a/tests/unit/test_mfd_connect/test_rshell.py b/tests/unit/test_mfd_connect/test_rshell.py index 84e73ad..9ef8621 100644 --- a/tests/unit/test_mfd_connect/test_rshell.py +++ b/tests/unit/test_mfd_connect/test_rshell.py @@ -291,17 +291,21 @@ def test_execute_command_logs_server_timeout(self, rshell, mocker, caplog): assert result.stdout == "" assert "timed out" in caplog.text - def test_execute_command_logs_server_internal_error(self, rshell, mocker, caplog): + @pytest.mark.parametrize("status", [400, 413, 429, 500, 503]) + def test_execute_command_raises_on_http_error(self, rshell, mocker, caplog, status): caplog.set_level(0) + response = requests.Response() + response.status_code = status + response.headers["X-RShell-Error"] = "Request failed" mocker.patch( "mfd_connect.rshell.requests.post", - return_value=Mock(status_code=500, text="", headers={}), + return_value=response, ) - result = rshell.execute_command("echo hi") + with pytest.raises(requests.HTTPError): + rshell.execute_command("echo hi") - assert result.return_code == -1 - assert "internal error" in caplog.text + assert f"RShell request failed ({status}): Request failed" in caplog.text def test_execute_command_logs_server_timeout_reason(self, rshell, mocker, caplog): """The server explains WHY it timed out - that diagnosis must reach the log.""" diff --git a/tests/unit/test_mfd_connect/test_rshell_client.py b/tests/unit/test_mfd_connect/test_rshell_client.py index 71fa199..58ec560 100644 --- a/tests/unit/test_mfd_connect/test_rshell_client.py +++ b/tests/unit/test_mfd_connect/test_rshell_client.py @@ -107,6 +107,20 @@ def _fake_system(cmd): class TestRShellClientScript: """Tests for RShell client script behavior.""" + @pytest.mark.parametrize("status", [429, 500, 503]) + def test_client_retries_http_errors_without_executing_response(self, monkeypatch, status, capsys): + result = _run_client( + monkeypatch, + scenarios=[ + {"status": status, "body": b"echo not-a-command"}, + {"status": 200, "body": b"end", "command_id": "cid-end"}, + ], + argv=["rshell_client.py", "127.0.0.1"], + ) + assert not result["system_calls"] + assert result["instances"][0].closed + assert "HTTP {}".format(status) in capsys.readouterr().out + def test_client_flow_with_no_command_then_end(self, monkeypatch): result = _run_client( monkeypatch, diff --git a/tests/unit/test_mfd_connect/test_rshell_server.py b/tests/unit/test_mfd_connect/test_rshell_server.py index e54151e..93da477 100644 --- a/tests/unit/test_mfd_connect/test_rshell_server.py +++ b/tests/unit/test_mfd_connect/test_rshell_server.py @@ -10,6 +10,11 @@ from pathlib import Path import pytest +from werkzeug.exceptions import ( + RequestEntityTooLarge, + ServiceUnavailable, + TooManyRequests, +) SERVER_SCRIPT_PATH = Path(__file__).resolve().parents[3] / "mfd_connect" / "rshell_server" / "rshell_server.py" @@ -33,6 +38,10 @@ def server_module(self): module = _load_server_module() module.output_queue.clear() module.output_queue_timestamps.clear() + module.output_sizes.clear() + module.output_errors.clear() + module.active_commands.clear() + module.control_command_deadlines.clear() module.abandoned_command_ids.clear() module.command_dict_queue.clear() module.clients.clear() @@ -53,8 +62,9 @@ def test_get_output_success(self, server_module): assert command_id not in server_module.output_queue_timestamps def test_get_output_timeout(self, server_module): + command_id = server_module.add_command_to_queue("ver", "192.0.2.1") with pytest.raises(TimeoutError, match="Command timed out"): - server_module.get_output("missing", timeout=-5) + server_module.get_output(command_id, timeout=-5) def test_get_output_timeout_marks_command_as_abandoned(self, server_module): """A caller that gives up must mark its command so it is not executed/stored later.""" @@ -65,15 +75,16 @@ def test_get_output_timeout_marks_command_as_abandoned(self, server_module): def test_get_output_is_woken_up_by_posted_result(self, server_module): """Waiting is event driven - the waiter returns as soon as the result is stored.""" + command_id = server_module.add_command_to_queue("ver", "192.0.2.1") def _post_later(): time.sleep(0.1) - server_module._store_output("cmd-later", "later", 4) + server_module._store_output(command_id, "later", 4) threading.Thread(target=_post_later, daemon=True).start() started = time.monotonic() - result = server_module.get_output("cmd-later", timeout=10) + result = server_module.get_output(command_id, timeout=10) elapsed = time.monotonic() - started assert result.output == "later" @@ -102,10 +113,13 @@ def test_add_command_to_queue_new_and_existing_queue(self, server_module): def test_add_command_to_queue_is_bounded(self, server_module): """Per-client queue must not grow without limits.""" limit = server_module.MAX_PENDING_COMMANDS_PER_CLIENT - for index in range(limit + 25): + for index in range(limit): server_module.add_command_to_queue(f"echo {index}", "10.0.0.9") + with pytest.raises(TooManyRequests): + server_module.add_command_to_queue("echo overflow", "10.0.0.9") - assert server_module.command_dict_queue["10.0.0.9"].qsize() <= limit + assert server_module.command_dict_queue["10.0.0.9"].qsize() == limit + assert server_module.command_dict_queue["10.0.0.9"].get_nowait().str == "echo 0" def test_health_check_endpoint(self, server_module): client = server_module.app.test_client() @@ -218,13 +232,130 @@ def test_execute_command_timeout_sends_reason_header_and_empty_body(self, server """The body must stay empty so the caller never stores an error page as stdout.""" client = server_module.app.test_client() - response = client.post("/execute_command", data={"command": "ver", "timeout": "-5", "ip": "1.2.3.4"}) + response = client.post( + "/execute_command", + data={"command": "ver", "timeout": "-5", "ip": "1.2.3.4"}, + ) assert response.status_code == 504 assert response.get_data() == b"" assert response.headers["rc"] == "-1" assert response.headers[server_module.TIMEOUT_REASON_HEADER] + @pytest.mark.parametrize( + "command", + ["echo first\r\nsecond", "echo \u03bb", "echo " + "\u03bb\n" * 10000, "echo " + "x" * 10000], + ids=["multiline", "unicode", "oversized_unicode", "oversized_ascii"], + ) + def test_timeout_reason_is_single_line_ascii_and_bounded(self, server_module, monkeypatch, command): + monkeypatch.setattr(server_module.time, "monotonic", lambda: 1000.0) + ip = "192.0.2.1" + server_module.client_last_seen[ip] = 1000.0 + response = server_module.app.test_client().post( + "/execute_command", data={"command": command, "timeout": "-5", "ip": ip} + ) + + assert response.status_code == 504 + assert response.data == b"" + reason = response.headers[server_module.TIMEOUT_REASON_HEADER] + expected = server_module.diagnose_timeout(ip, command, waiting_since=1000.0) + assert reason == expected.encode("unicode_escape").decode("ascii")[:1024] + assert len(reason.encode("ascii")) <= 1024 + assert "\r" not in reason and "\n" not in reason + assert not server_module.active_commands + + def test_later_timeout_remembers_command_abandoned_after_default_wait(self, server_module, monkeypatch): + now = 1000.0 + monkeypatch.setattr(server_module.time, "monotonic", lambda: now) + client = server_module.app.test_client() + ip = "192.0.2.1" + picked_up = False + + def advance_wait(seconds): + nonlocal now, picked_up + if not picked_up: + assert client.get("/getCommandToExecute", environ_base={"REMOTE_ADDR": ip}).status_code == 200 + picked_up = True + now += seconds + + monkeypatch.setattr(server_module._output_available, "wait", advance_wait) + first = client.post("/execute_command", data={"command": "hanging-tool", "ip": ip}) + assert first.status_code == 504 + assert now == 1605.0 + assert "picked this command up" in first.headers[server_module.TIMEOUT_REASON_HEADER] + + second = client.post("/execute_command", data={"command": "ver", "ip": ip, "timeout": "-5"}) + assert second.status_code == 504 + assert "still blocked executing 'hanging-tool'" in second.headers[server_module.TIMEOUT_REASON_HEADER] + assert ip not in server_module.client_last_seen + assert ip not in server_module.clients + + def test_running_diagnostics_expire_without_removing_active_commands(self, server_module, monkeypatch): + now = 1000.0 + monkeypatch.setattr(server_module.time, "monotonic", lambda: now) + client = server_module.app.test_client() + for ip in ("192.0.2.1", "192.0.2.2"): + server_module.add_command_to_queue("tool", ip) + client.get("/getCommandToExecute", environ_base={"REMOTE_ADDR": ip}) + abandoned = server_module.client_running_command["192.0.2.1"].command_id + server_module._abandon_command(abandoned) + + now += 3599 + server_module._cleanup_stale_outputs() + assert "192.0.2.1" in server_module.client_running_command + now += 1 + server_module._cleanup_stale_outputs() + assert "192.0.2.1" not in server_module.client_running_command + assert "192.0.2.2" in server_module.client_running_command + assert server_module.client_running_command["192.0.2.2"].command_id in server_module.active_commands + + def test_running_diagnostics_cap_evicts_abandoned_not_active_commands(self, server_module, monkeypatch): + monkeypatch.setattr(server_module, "MAX_CLIENTS", 2) + now = 1000.0 + monkeypatch.setattr(server_module.time, "monotonic", lambda: now) + client = server_module.app.test_client() + for ip in ("192.0.2.1", "192.0.2.2"): + server_module.add_command_to_queue("tool", ip) + client.get("/getCommandToExecute", environ_base={"REMOTE_ADDR": ip}) + server_module._abandon_command(server_module.client_running_command["192.0.2.2"].command_id) + now += server_module.CLIENT_LIVENESS_TIMEOUT_SECONDS + 1 + server_module._cleanup_stale_outputs() + assert set(server_module.client_running_command) == {"192.0.2.1", "192.0.2.2"} + server_module.add_command_to_queue("next tool", "192.0.2.3") + assert client.get("/getCommandToExecute", environ_base={"REMOTE_ADDR": "192.0.2.3"}).status_code == 200 + assert set(server_module.client_running_command) == {"192.0.2.1", "192.0.2.3"} + + @pytest.mark.parametrize("completion", ["result", "poll", "disconnect"]) + def test_retained_running_diagnostic_is_cleared_when_client_returns(self, server_module, monkeypatch, completion): + now = 1000.0 + monkeypatch.setattr(server_module.time, "monotonic", lambda: now) + client = server_module.app.test_client() + ip = "192.0.2.1" + command_id = server_module.add_command_to_queue("tool", ip) + client.get("/getCommandToExecute", environ_base={"REMOTE_ADDR": ip}) + server_module._abandon_command(command_id) + now += server_module.CLIENT_LIVENESS_TIMEOUT_SECONDS + 1 + server_module._cleanup_stale_outputs() + assert ip in server_module.client_running_command + assert ip not in server_module.clients + + if completion == "result": + assert ( + client.post( + "/post_result", + data=b"late output", + headers={"CommandID": command_id}, + environ_base={"REMOTE_ADDR": ip}, + ).status_code + == 200 + ) + elif completion == "poll": + assert client.get("/getCommandToExecute", environ_base={"REMOTE_ADDR": ip}).status_code == 204 + else: + assert client.post(f"/disconnect_client/{ip}").status_code == 200 + assert ip not in server_module.client_running_command + assert not server_module.output_queue + def test_disconnect_client_clears_liveness_state(self, server_module): """A reconnecting client must not inherit the liveness state of the previous session.""" client = server_module.app.test_client() @@ -264,7 +395,7 @@ def test_get_command_to_execute_skips_abandoned_commands(self, server_module): assert response.status_code == 200 assert response.headers["CommandID"] == live_id assert response.get_data(as_text=True) == "echo alive" - assert abandoned_id not in server_module.abandoned_command_ids + assert abandoned_id not in server_module.active_commands def test_get_command_to_execute_returns_204_when_only_abandoned(self, server_module): client = server_module.app.test_client() @@ -277,12 +408,13 @@ def test_get_command_to_execute_returns_204_when_only_abandoned(self, server_mod def test_post_exception_endpoint(self, server_module): client = server_module.app.test_client() - response = client.post("/exception", data=b"boom", headers={"CommandID": "cid-1"}) + command_id = server_module.add_command_to_queue("ver", "192.0.2.1") + response = client.post("/exception", data=b"boom", headers={"CommandID": command_id}) assert response.status_code == 200 - assert server_module.output_queue["cid-1"].output == "boom" - assert server_module.output_queue["cid-1"].rc == -1 - assert "cid-1" in server_module.output_queue_timestamps + assert server_module.output_queue[command_id].output == "boom" + assert server_module.output_queue[command_id].rc == -1 + assert command_id in server_module.output_queue_timestamps def test_execute_command_endpoint_paths(self, server_module, monkeypatch): client = server_module.app.test_client() @@ -317,6 +449,70 @@ def test_execute_command_endpoint_paths(self, server_module, monkeypatch): assert response_normal.headers["Content-type"].startswith("text/plain") assert response_normal.headers["CommandID"] + @pytest.mark.parametrize("command", ["end", "reset -c"]) + @pytest.mark.parametrize("timeout", [None, 10]) + def test_undelivered_control_commands_expire_and_release_client_slots( + self, server_module, monkeypatch, command, timeout + ): + now = 1000.0 + monkeypatch.setattr(server_module.time, "monotonic", lambda: now) + client = server_module.app.test_client() + for index in range(server_module.MAX_CLIENTS): + data = {"command": command, "ip": f"192.0.2.{index}"} + if timeout is not None: + data["timeout"] = str(timeout) + assert client.post("/execute_command", data=data).status_code == 200 + old_ids = set(server_module.active_commands) + assert len(old_ids) == server_module.MAX_CLIENTS + new_request = {"command": command, "ip": "198.51.100.1"} + now += (600 if timeout is None else timeout) + server_module.CLIENT_LOOP_WAIT_SECONDS - 0.1 + assert client.post("/execute_command", data=new_request).status_code == 429 + now += 0.1 + assert client.post("/execute_command", data=new_request).status_code == 200 + assert len(server_module.active_commands) == 1 + assert set(server_module.command_dict_queue) == {"198.51.100.1"} + assert not old_ids.intersection(server_module.control_command_deadlines) + assert client.get("/getCommandToExecute", environ_base={"REMOTE_ADDR": "192.0.2.1"}).status_code == 204 + assert not server_module._store_output(next(iter(old_ids)), "late", 0) + + @pytest.mark.parametrize("command", ["end", "reset -c"]) + def test_control_command_deadline_is_cleared_after_delivery_or_disconnect( + self, server_module, monkeypatch, command + ): + now = 1000.0 + monkeypatch.setattr(server_module.time, "monotonic", lambda: now) + client = server_module.app.test_client() + ip = "192.0.2.1" + assert client.post("/execute_command", data={"command": command, "ip": ip, "timeout": "10"}).status_code == 200 + response = client.get("/getCommandToExecute", environ_base={"REMOTE_ADDR": ip}) + assert response.status_code == 200 + assert response.get_data(as_text=True) == command + assert not server_module.control_command_deadlines + assert not server_module.active_commands + now += 15 + server_module._cleanup_stale_outputs() + assert response.headers["CommandID"] not in server_module.abandoned_command_ids + + client.post("/execute_command", data={"command": command, "ip": ip}) + assert server_module.control_command_deadlines + client.post(f"/disconnect_client/{ip}") + assert not server_module.control_command_deadlines + assert not server_module.active_commands + + def test_control_expiry_preserves_other_commands_for_same_client(self, server_module, monkeypatch): + now = 1000.0 + monkeypatch.setattr(server_module.time, "monotonic", lambda: now) + client = server_module.app.test_client() + ip = "192.0.2.1" + client.post("/execute_command", data={"command": "end", "ip": ip, "timeout": "10"}) + normal_id = server_module.add_command_to_queue("ver", ip) + now += 15 + response = client.get("/getCommandToExecute", environ_base={"REMOTE_ADDR": ip}) + assert response.status_code == 200 + assert response.headers["CommandID"] == normal_id + assert normal_id in server_module.active_commands + assert not server_module.control_command_deadlines + def test_execute_command_returns_gateway_timeout_instead_of_html_error(self, server_module): """A timeout must not return a Flask HTML 500 page, which the caller stores as stdout.""" client = server_module.app.test_client() @@ -338,7 +534,7 @@ def test_execute_command_returns_gateway_timeout_instead_of_html_error(self, ser def test_timed_out_command_does_not_desync_next_command(self, server_module): """After a timeout the next command must succeed - the server has to re-sync.""" client = server_module.app.test_client() - ip = "10.102.23.150" + ip = "192.0.2.1" timed_out = client.post( "/execute_command", @@ -383,22 +579,24 @@ def test_disconnect_client_abandons_pending_commands(self, server_module): def test_post_result_endpoint(self, server_module): client = server_module.app.test_client() + first_id = server_module.add_command_to_queue("ver", "192.0.2.1") + second_id = server_module.add_command_to_queue("help", "192.0.2.1") - response_default_rc = client.post("/post_result", data=b"output-a", headers={"CommandID": "cmd-a"}) + response_default_rc = client.post("/post_result", data=b"output-a", headers={"CommandID": first_id}) assert response_default_rc.status_code == 200 - assert server_module.output_queue["cmd-a"].output == "output-a" - assert server_module.output_queue["cmd-a"].rc == -1 - assert "cmd-a" in server_module.output_queue_timestamps + assert server_module.output_queue[first_id].output == "output-a" + assert server_module.output_queue[first_id].rc == -1 + assert first_id in server_module.output_queue_timestamps response_given_rc = client.post( "/post_result", data=b"output-b", - headers={"CommandID": "cmd-b", "rc": "3"}, + headers={"CommandID": second_id, "rc": "3"}, ) assert response_given_rc.status_code == 200 - assert server_module.output_queue["cmd-b"].output == "output-b" - assert server_module.output_queue["cmd-b"].rc == 3 - assert "cmd-b" in server_module.output_queue_timestamps + assert server_module.output_queue[second_id].output == "output-b" + assert server_module.output_queue[second_id].rc == 3 + assert second_id in server_module.output_queue_timestamps def test_cleanup_stale_outputs_removes_expired_entries_only(self, server_module): # Timestamps simulate time.monotonic() values (seconds since arbitrary boot reference). @@ -423,14 +621,16 @@ def test_cleanup_stale_outputs_expires_abandoned_ids(self, server_module): assert "old" not in server_module.abandoned_command_ids assert "recent" in server_module.abandoned_command_ids - def test_output_queue_is_hard_capped(self, server_module): + def test_output_queue_is_hard_capped(self, server_module, monkeypatch): """Even without TTL expiry the stored outputs must stay bounded.""" - limit = server_module.MAX_STORED_OUTPUTS - for index in range(limit + 120): - server_module._store_output(f"cid-{index}", "payload", 0) - - assert len(server_module.output_queue) <= limit - assert len(server_module.output_queue_timestamps) == len(server_module.output_queue) + monkeypatch.setattr(server_module, "MAX_STORED_OUTPUTS", 1) + first_id = server_module.add_command_to_queue("ver", "192.0.2.1") + second_id = server_module.add_command_to_queue("help", "192.0.2.1") + server_module._store_output(first_id, "payload", 0) + with pytest.raises(ServiceUnavailable): + server_module._store_output(second_id, "overflow", 0) + assert list(server_module.output_queue) == [first_id] + assert len(server_module.output_queue_timestamps) == 1 def test_concurrent_access_is_thread_safe(self, server_module): """Werkzeug serves requests in threads - shared state must not raise or corrupt.""" @@ -439,10 +639,9 @@ def test_concurrent_access_is_thread_safe(self, server_module): def _worker(worker_id): try: for index in range(30): - command_id = f"w{worker_id}-{index}" + command_id = server_module.add_command_to_queue(f"cmd{index}", f"192.0.2.{worker_id + 1}") server_module._store_output(command_id, "data", 0) server_module.get_output(command_id, timeout=1) - server_module.add_command_to_queue(f"cmd{index}", f"10.0.1.{worker_id}") server_module._cleanup_stale_outputs() except Exception as exc: # noqa: BLE001 errors.append(f"{type(exc).__name__}: {exc}") @@ -455,6 +654,8 @@ def _worker(worker_id): assert not errors assert not server_module.output_queue + assert not server_module.active_commands + assert not server_module.output_sizes def test_run_function_starts_flask(self, server_module, monkeypatch): """Test that the run() function starts Flask with correct host and port.""" @@ -469,6 +670,147 @@ def _fake_run(self, host, port): assert captured == {"host": "0.0.0.0", "port": 80} + def test_late_results_stay_rejected_after_tombstone_eviction(self, server_module, monkeypatch): + monkeypatch.setattr(server_module, "MAX_ABANDONED_COMMAND_IDS", 1) + command_id = server_module.add_command_to_queue("ver", "192.0.2.1") + server_module._abandon_command(command_id) + server_module._abandon_command("another-command") + + assert command_id not in server_module.abandoned_command_ids + assert server_module._store_output(command_id, "late", 0) is False + assert server_module._store_output(command_id, "duplicate", 0) is False + client = server_module.app.test_client() + assert client.get("/getCommandToExecute", environ_base={"REMOTE_ADDR": "192.0.2.1"}).status_code == 204 + assert not server_module.output_queue + assert not server_module.output_sizes + + def test_output_is_only_collected_once(self, server_module): + command_id = server_module.add_command_to_queue("ver", "192.0.2.1") + assert server_module._store_output(command_id, "first", 0) + assert not server_module._store_output(command_id, "duplicate", 1) + assert server_module.get_output(command_id).output == "first" + assert not server_module._store_output(command_id, "late duplicate", 1) + assert not server_module.active_commands + assert not server_module.output_queue_timestamps + assert not server_module.output_sizes + assert not server_module.command_dict_queue + + @pytest.mark.parametrize("connected", [False, True]) + def test_disconnect_clears_all_command_states(self, server_module, connected): + client = server_module.app.test_client() + ip = "192.0.2.1" + running = server_module.add_command_to_queue("ver", ip) + if connected: + client.get("/getCommandToExecute", environ_base={"REMOTE_ADDR": ip}) + ready = server_module.add_command_to_queue("help", ip) + server_module._store_output(ready, "result", 0) + pending = server_module.add_command_to_queue("echo pending", ip) + + assert client.post(f"/disconnect_client/{ip}").status_code == 200 + for command_id in (running, ready, pending): + assert not server_module._store_output(command_id, "late", 0) + with pytest.raises(TimeoutError): + server_module.get_output(command_id) + assert not server_module.active_commands + assert not server_module.output_queue + assert not server_module.output_sizes + assert not server_module.output_queue_timestamps + assert not server_module.command_dict_queue + assert not server_module.client_running_command + + def test_active_command_limit_rejects_without_eviction(self, server_module, monkeypatch): + monkeypatch.setattr(server_module, "MAX_ACTIVE_COMMANDS", 1) + first_id = server_module.add_command_to_queue("ver", "192.0.2.1") + response = server_module.app.test_client().post("/execute_command", data={"command": "help", "ip": "192.0.2.2"}) + assert response.status_code == 429 + assert response.headers["rc"] == "-1" + assert response.headers["X-RShell-Error"] + assert list(server_module.active_commands) == [first_id] + assert "192.0.2.2" not in server_module.command_dict_queue + + def test_client_limit_and_idle_cleanup(self, server_module, monkeypatch): + monkeypatch.setattr(server_module, "MAX_CLIENTS", 1) + client = server_module.app.test_client() + assert client.get("/getCommandToExecute", environ_base={"REMOTE_ADDR": "192.0.2.1"}).status_code == 204 + assert client.get("/getCommandToExecute", environ_base={"REMOTE_ADDR": "192.0.2.2"}).status_code == 429 + server_module.client_last_seen["192.0.2.1"] -= server_module.CLIENT_LIVENESS_TIMEOUT_SECONDS + 1 + assert client.get("/getCommandToExecute", environ_base={"REMOTE_ADDR": "192.0.2.2"}).status_code == 204 + assert server_module.clients == ["192.0.2.2"] + + def test_idle_cleanup_preserves_active_commands(self, server_module): + command_id = server_module.add_command_to_queue("ver", "192.0.2.1") + server_module.client_last_seen["192.0.2.1"] = 0 + server_module._cleanup_stale_outputs(now=1000) + assert command_id in server_module.active_commands + assert "192.0.2.1" in server_module.client_last_seen + + def test_output_byte_limits(self, server_module, monkeypatch): + assert server_module.MAX_OUTPUT_BYTES == 16 * 1024 * 1024 + assert server_module.MAX_STORED_OUTPUT_BYTES == 64 * 1024 * 1024 + monkeypatch.setattr(server_module, "MAX_OUTPUT_BYTES", 4) + monkeypatch.setattr(server_module, "MAX_STORED_OUTPUT_BYTES", 8) + ids = [server_module.add_command_to_queue("ver", "192.0.2.1") for _ in range(4)] + assert server_module._store_output(ids[0], "\u00e9\u00e9", 0) + assert server_module._store_output(ids[1], "1234", 0) + assert sum(server_module.output_sizes.values()) == 8 + with pytest.raises(ServiceUnavailable): + server_module._store_output(ids[2], "x", 0) + with pytest.raises(RequestEntityTooLarge): + server_module._store_output(ids[3], "12345", 0) + with pytest.raises(ServiceUnavailable): + server_module.get_output(ids[2]) + with pytest.raises(RequestEntityTooLarge): + server_module.get_output(ids[3]) + assert server_module.get_output(ids[0]).output == "\u00e9\u00e9" + assert server_module.get_output(ids[1]).output == "1234" + assert not server_module.output_sizes + assert not server_module.output_errors + assert not server_module.active_commands + + @pytest.mark.parametrize("endpoint", ["/post_result", "/exception"]) + def test_oversized_request_notifies_waiter(self, server_module, endpoint): + server_module.app.config["MAX_CONTENT_LENGTH"] = 4 + command_id = server_module.add_command_to_queue("ver", "192.0.2.1") + response = server_module.app.test_client().post( + endpoint, data=b"12345", headers={"CommandID": command_id, "rc": "0"} + ) + assert response.status_code == 413 + with pytest.raises(RequestEntityTooLarge): + server_module.get_output(command_id) + assert not server_module.output_queue + assert not server_module.output_errors + assert not server_module.active_commands + + def test_disconnect_wakes_waiting_caller(self, server_module): + command_id = server_module.add_command_to_queue("ver", "192.0.2.1") + errors = [] + started = threading.Event() + + def wait_for_output(): + started.set() + try: + server_module.get_output(command_id, timeout=60) + except TimeoutError as error: + errors.append(error) + + thread = threading.Thread(target=wait_for_output, daemon=True) + thread.start() + assert started.wait(1) + server_module.app.test_client().post("/disconnect_client/192.0.2.1") + thread.join(timeout=2) + assert not thread.is_alive() + assert len(errors) == 1 + + def test_repeated_results_release_payload_memory(self, server_module): + for _ in range(32): + command_id = server_module.add_command_to_queue("ver", "192.0.2.1") + server_module._store_output(command_id, "x" * 65536, 0) + assert len(server_module.get_output(command_id).output) == 65536 + assert not server_module.output_queue + assert not server_module.output_sizes + assert not server_module.active_commands + assert not server_module.command_dict_queue + def test_main_block_starts_flask(self, monkeypatch): captured = {}