diff --git a/README.md b/README.md index 6d129e1..b22723b 100644 --- a/README.md +++ b/README.md @@ -1173,6 +1173,13 @@ 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. * `/post_result` - Endpoint to post results back to the host. Headers fields: * `CommandID` - Unique identifier for the command. @@ -1184,8 +1191,23 @@ 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. + * `/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. * `/health/` - Endpoint to check the health status of the connection. + * `/disconnect_client/` - Removes the client and drops everything still queued for it. + +### 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. | + +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. diff --git a/mfd_connect/rshell.py b/mfd_connect/rshell.py index 3b366b3..9e2f1d7 100644 --- a/mfd_connect/rshell.py +++ b/mfd_connect/rshell.py @@ -76,10 +76,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 +100,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") @@ -201,13 +208,27 @@ def execute_command( 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}") + 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}, proxies={"no_proxy": "*"}, ) + if response.status_code == 504: + logger.log( + level=log_levels.MODULE_DEBUG, + msg=f"RShell server timed out waiting for the result of '{command}'. " + f"The command was dropped so it will not be executed later by the client.", + ) + 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 +357,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_server/rshell_server.py b/mfd_connect/rshell_server/rshell_server.py index 36e1646..c638b8a 100644 --- a/mfd_connect/rshell_server/rshell_server.py +++ b/mfd_connect/rshell_server/rshell_server.py @@ -5,31 +5,147 @@ 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. """ +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.2.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 + # 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_object = OutputObject +command_object = CommandObject -output_queue: dict[str, output_object] = dict() +# 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 = [] +# 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: """ - Retrieve the output for a given command ID. + 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: + 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 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 +154,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,16 +178,28 @@ 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: + with _state_lock: + connected = ip in clients + if connected: return Response("OK", status=200) else: return Response("Client not connected", status=503) @@ -78,23 +210,28 @@ 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. + :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_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 + 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 +247,7 @@ 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) + _store_output(command_id, str(read_data, encoding="utf-8"), rc=-1) return Response("Exception received", status=200) @@ -127,36 +264,57 @@ 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) + + 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. + print(f"Command {_id} timed out after {timeout}s - marked as abandoned") return Response( - process.output.encode("utf-8"), - status=200, - headers={ - "Content-type": "text/plain", - "CommandID": _id, - "rc": process.rc, - }, + b"", + status=504, + headers={"Content-type": "text/plain", "CommandID": _id, "rc": "-1"}, ) - 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_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 +326,7 @@ 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) + _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..951a8f7 100644 --- a/tests/unit/test_mfd_connect/test_rshell.py +++ b/tests/unit/test_mfd_connect/test_rshell.py @@ -222,7 +222,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 +249,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 +260,42 @@ 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_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..4c1d408 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,6 +32,8 @@ 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() return module @@ -38,33 +42,52 @@ 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 +97,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() @@ -99,6 +130,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 +160,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 +195,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 +266,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 +276,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."""