diff --git a/README.md b/README.md index 6d129e1..f2813ae 100644 --- a/README.md +++ b/README.md @@ -1173,6 +1173,17 @@ 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. + * `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`. 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: * `CommandID` - Unique identifier for the command. @@ -1184,8 +1195,78 @@ 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. + +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 +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: + +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` | +| 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) +``` + +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: + +| 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. | +| `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. `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 +1275,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..c5de937 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,42 @@ 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"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 >= 400: + logger.log( + level=log_levels.MODULE_DEBUG, + 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, @@ -336,7 +382,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..0e8def1 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,20 @@ 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. + + :param connection: HTTP connection to close, or None when no connection was created. + """ + 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 @@ -50,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() @@ -61,27 +97,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 +166,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..ddaf0e3 100644 --- a/mfd_connect/rshell_server/rshell_server.py +++ b/mfd_connect/rshell_server/rshell_server.py @@ -5,31 +5,311 @@ 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 collections import OrderedDict 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" + +# 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 +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 -__version__ = "1.1.0" # 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 + + +class RunningCommand(NamedTuple): + """Store the command a client picked up and has not reported back yet.""" -output_queue: dict[str, output_object] = dict() + 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() +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. 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__) +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 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. + """ + 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: + _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: + _abandon_command(next(iter(output_queue))) + while len(abandoned_command_ids) > MAX_ABANDONED_COMMAND_IDS: + abandoned_command_ids.popitem(last=False) -def get_output(command_id: str, timeout: float = 600) -> output_object: + 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: """ - Retrieve the output for a given command ID. + 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 _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) + 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: + """ + 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. + """ + with _output_available: + _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. + + :param ip_address: The IP address the result came from. + :param command_id: The ID of the command the client reported. + """ + with _state_lock: + 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: + """ + 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. + """ + 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 ( + 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,39 +318,73 @@ 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") - - -def add_command_to_queue(command: str, ip_address: str) -> str: + 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: + _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) + raise TimeoutError("Command timed out") + _output_available.wait(min(remaining, OUTPUT_WAIT_SLICE_SECONDS)) + + +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) - 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: + _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 + 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 @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 +392,46 @@ 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: + _cleanup_stale_outputs() + _check_client_limit(ip_address) + 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 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, + 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"]) @@ -106,12 +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")) - output_queue[command_id] = output_object(output=str(read_data, encoding="utf-8"), rc=-1) - return Response("Exception received", status=200) + _mark_client_finished(str(request.remote_addr), command_id) + 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"]) @@ -127,35 +484,66 @@ 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, timeout) + 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.encode("unicode_escape").decode("ascii")[:MAX_TIMEOUT_REASON_LENGTH], }, ) - else: - return Response("No command provided", status=400) + finally: + with _state_lock: + _abandon_command(_id) + + 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) + 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) + 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) @@ -163,13 +551,8 @@ def disconnect_client(ip_address: str) -> Response: @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")) - output_queue[command_id] = output_object(output=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 9a507d1..9ef8621 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,99 @@ 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 + + @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=response, + ) + + with pytest.raises(requests.HTTPError): + rshell.execute_command("echo hi") + + 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.""" + 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_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 13dfaea..93da477 100644 --- a/tests/unit/test_mfd_connect/test_rshell_server.py +++ b/tests/unit/test_mfd_connect/test_rshell_server.py @@ -5,9 +5,16 @@ import importlib.util import runpy import sys +import threading +import time 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" @@ -30,41 +37,70 @@ class TestRShellServerScript: 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() + 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): + 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_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.""" + command_id = server_module.add_command_to_queue("ver", "192.0.2.1") - result = server_module.get_output("cmd-later", timeout=0) + def _post_later(): + time.sleep(0.1) + server_module._store_output(command_id, "later", 4) + + threading.Thread(target=_post_later, daemon=True).start() + + started = time.monotonic() + result = server_module.get_output(command_id, 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 +110,17 @@ 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): + 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"].get_nowait().str == "echo 0" + def test_health_check_endpoint(self, server_module): client = server_module.app.test_client() @@ -81,10 +128,247 @@ 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] + + @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() + 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,13 +383,38 @@ 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.active_commands + + 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"}) + 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 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() @@ -140,33 +449,213 @@ 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() + + 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 = "192.0.2.1" + + 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() + 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 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 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). + # 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, monkeypatch): + """Even without TTL expiry the stored outputs must stay bounded.""" + 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.""" + errors = [] + + def _worker(worker_id): + try: + for index in range(30): + 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._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 + 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.""" @@ -181,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 = {}