fix: prevent rshell server unbounded server memory growth

Signed-off-by: Ziomek, Kamil <kamil.ziomek@intel.com>

Author: k-ziomek

mfd_connect/rshell_server.py

@@ -8,26 +8,63 @@
 """
 
 import time
-from collections import namedtuple
 from queue import Queue
+from typing import NamedTuple
 from uuid import uuid4
 
 from flask import Flask, Response, request
 
 __version__ = "1.1.0"
 
+CLIENT_LOOP_WAIT_SECONDS = 5
+STALE_OUTPUT_TTL_SECONDS = 3600
+
+
 # Global command queue
-output_object = namedtuple("OutputObject", ["output", "rc"])
-command_object = namedtuple("CommandObject", ["command_id", "str"])
+class OutputObject(NamedTuple):
+    """Store command output together with its return code."""
+
+    output: str
+    rc: int
+
+
+class CommandObject(NamedTuple):
+    """Store queued command metadata sent to a specific client."""
+
+    command_id: str
+    str: str
 
-output_queue: dict[str, output_object] = dict()
+
+output_object = OutputObject
+command_object = CommandObject
+
+output_queue: dict[str, OutputObject] = dict()
+output_queue_timestamps: dict[str, float] = dict()
 command_dict_queue: dict[str, Queue] = dict()
 clients: list = []
 
 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."""
+    current_time = time.monotonic() if now is None else now
+    stale_command_ids = [
+        command_id for command_id, created_at in output_queue_timestamps.items() if current_time - created_at >= ttl
+    ]
+    for command_id in stale_command_ids:
+        output_queue.pop(command_id, None)
+        output_queue_timestamps.pop(command_id, None)
+
+
+def _store_output(command_id: str, output: str, rc: int) -> None:
+    """Persist command output together with its insertion timestamp."""
+    _cleanup_stale_outputs()
+    output_queue[command_id] = output_object(output=output, rc=rc)
+    output_queue_timestamps[command_id] = time.time()
+
+
+def get_output(command_id: str, timeout: float = 600) -> OutputObject:
     """
     Retrieve the output for a given command ID.
 
@@ -38,11 +75,13 @@ 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
+    timeout = timeout + CLIENT_LOOP_WAIT_SECONDS  # add time for client loop waiting
     while timeout > 0:
-        result = output_queue.get(command_id, None)
+        result = output_queue.pop(command_id, None)
         if result is not None:
+            output_queue_timestamps.pop(command_id, None)
             return result
+        _cleanup_stale_outputs()
         time.sleep(1)
         timeout -= 1
     raise TimeoutError("Command timed out")
@@ -110,7 +149,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)
 
 
@@ -140,7 +179,7 @@ def execute_command() -> Response:
             headers={
                 "Content-type": "text/plain",
                 "CommandID": _id,
-                "rc": process.rc,
+                "rc": str(process.rc),
             },
         )
     else:
@@ -156,6 +195,7 @@ def disconnect_client(ip_address: str) -> Response:
     """
     if ip_address in clients:
         clients.remove(ip_address)
+        command_dict_queue.pop(ip_address, None)
         print(f"Client disconnected: {ip_address}")
     return Response("Client disconnected", status=200)
 
@@ -168,7 +208,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)
 
 

tests/unit/test_mfd_connect/test_rshell_server.py

@@ -27,6 +27,7 @@ class TestRShellServerScript:
     def server_module(self):
         module = _load_server_module()
         module.output_queue.clear()
+        module.output_queue_timestamps.clear()
         module.command_dict_queue.clear()
         module.clients.clear()
         return module
@@ -35,10 +36,13 @@ 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"):
@@ -49,13 +53,14 @@ class _QueueProbe:
             def __init__(self):
                 self.count = 0
 
-            def get(self, _command_id, _default=None):
+            def pop(self, _command_id, _default=None):
                 self.count += 1
                 if self.count == 1:
                     return None
                 return server_module.output_object(output="later", rc=4)
 
         monkeypatch.setattr(server_module, "output_queue", _QueueProbe())
+        monkeypatch.setattr(server_module, "output_queue_timestamps", {})
         monkeypatch.setattr(server_module.time, "sleep", lambda _x: None)
 
         result = server_module.get_output("cmd-later", timeout=0)
@@ -103,6 +108,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,10 +146,12 @@ def test_execute_command_endpoint_paths(self, server_module, monkeypatch):
     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
@@ -155,6 +163,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",
@@ -164,6 +173,19 @@ 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):
+        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_main_block_starts_flask(self, monkeypatch):
         captured = {}