Harden: fail-closed+const-time auth, id validation, FILES_DIR reaper+trim, server size cap, copy-outside-lock, api_poll .get(), no analytics

app.py

@@ -21,6 +21,8 @@ second so a leaked HF read-token alone can't drive your shells.
 """
 
 import os
+import re
+import hmac
 import time
 import uuid
 import shutil
@@ -36,10 +38,17 @@ CLIENT_STALE_AFTER = 30                                # seconds w/o a poll => "
 MAX_OUTPUT_CHARS = 200_000                             # cap stored output per command
 MAX_HISTORY_PER_CLIENT = 200                           # ring-buffer of commands
 MAX_FILE_BYTES = 100 * 1024 * 1024                     # 100 MB cap on transfers
+FILE_TTL = 3600                                        # delete stored files older than 1h
 
 FILES_DIR = pathlib.Path(tempfile.gettempdir()) / "relay_files"
 FILES_DIR.mkdir(exist_ok=True)
 
+_HEX32 = re.compile(r"\A[0-9a-f]{32}\Z")
+
+
+def _valid_id(s) -> bool:
+    return isinstance(s, str) and bool(_HEX32.match(s))
+
 LOCK = threading.Lock()
 CLIENTS: dict[str, dict] = {}                          # cid -> {last_seen, meta}
 COMMANDS: dict[str, dict] = {}                         # command_id -> record
@@ -59,8 +68,10 @@ def _human(n: float) -> str:
 
 
 def _check(secret: str):
-    if SECRET and secret != SECRET:
-        raise gr.Error("bad secret")
+    # Fail closed: an empty/unset RELAY_SECRET rejects everything (no silent
+    # auth-off). Constant-time compare to avoid a timing oracle on the secret.
+    if not SECRET or not hmac.compare_digest(secret or "", SECRET):
+        raise gr.Error("unauthorized")
 
 
 def _touch_client(cid: str, meta: str = ""):
@@ -70,6 +81,57 @@ def _touch_client(cid: str, meta: str = ""):
         c["meta"] = meta
 
 
+def _unlink(path: str):
+    try:
+        os.unlink(path)
+    except OSError:
+        pass
+
+
+def _drop_command(cmd_id: str):
+    """Remove a command record and any file it owns on disk. Call under LOCK."""
+    rec = COMMANDS.pop(cmd_id, None)
+    if not rec:
+        return
+    if rec.get("stored_path"):
+        _unlink(rec["stored_path"])
+    fid = rec.get("file_id")
+    if fid:
+        info = PUT_FILES.pop(fid, None)
+        if info:
+            _unlink(info["path"])
+
+
+def _append_and_trim(cid: str, cmd_id: str):
+    """Append to a client's history and ring-buffer it, deleting evicted files.
+    Call under LOCK. Used by every append site so growth is bounded uniformly."""
+    hist = BY_CLIENT[cid]
+    hist.append(cmd_id)
+    while len(hist) > MAX_HISTORY_PER_CLIENT:
+        _drop_command(hist.pop(0))
+
+
+def _file_reaper():
+    """Backstop against disk growth on the ephemeral Space: delete stored files
+    older than FILE_TTL and prune PUT_FILES whose file is gone."""
+    while True:
+        time.sleep(600)
+        cutoff = _now() - FILE_TTL
+        try:
+            for p in FILES_DIR.iterdir():
+                try:
+                    if p.is_file() and p.stat().st_mtime < cutoff:
+                        p.unlink()
+                except OSError:
+                    pass
+        except OSError:
+            pass
+        with LOCK:
+            for fid in list(PUT_FILES):
+                if not os.path.isfile(PUT_FILES[fid]["path"]):
+                    PUT_FILES.pop(fid, None)
+
+
 # --------------------------------------------------------------------------- #
 # Machine API (called by the daemon via gradio_client)                         #
 # --------------------------------------------------------------------------- #
@@ -90,8 +152,10 @@ def api_poll(client_id: str, secret: str):
     out = []
     with LOCK:
         _touch_client(client_id)
-        for cmd_id in BY_CLIENT[client_id]:
-            rec = COMMANDS[cmd_id]
+        for cmd_id in list(BY_CLIENT.get(client_id, [])):
+            rec = COMMANDS.get(cmd_id)
+            if rec is None:
+                continue
             if rec["status"] == "pending":
                 rec["status"] = "acknowledged"
                 rec["acknowledged_at"] = _now()
@@ -113,6 +177,8 @@ def api_poll(client_id: str, secret: str):
 def api_result(command_id: str, client_id: str, secret: str,
                exit_code: float, stdout: str, session: str):
     _check(secret)
+    if not _valid_id(command_id):
+        raise gr.Error("bad command id")
     with LOCK:
         rec = COMMANDS.get(command_id)
         if rec is None:
@@ -122,7 +188,7 @@ def api_result(command_id: str, client_id: str, secret: str,
                 "command": "(unknown — relay restarted)", "created": _now(),
                 "status": "pending",
             }
-            BY_CLIENT[client_id].append(command_id)
+            _append_and_trim(client_id, command_id)
         rec["status"] = "done"
         rec["finished_at"] = _now()
         rec["exit_code"] = int(exit_code)
@@ -134,6 +200,8 @@ def api_result(command_id: str, client_id: str, secret: str,
 def api_fetch_file(file_id: str, secret: str):
     """Daemon calls this to download a server->client (push) payload."""
     _check(secret)
+    if not _valid_id(file_id):
+        raise gr.Error("bad file id")
     info = PUT_FILES.get(file_id)
     if not info or not os.path.isfile(info["path"]):
         raise gr.Error("unknown or expired file_id")
@@ -144,6 +212,24 @@ def api_upload_file(command_id: str, client_id: str, secret: str, fileobj, error
     """Daemon calls this to post the result of a client->server (pull) request.
     `fileobj` is the uploaded file path on success, or None when `error` is set."""
     _check(secret)
+    if not _valid_id(command_id):
+        raise gr.Error("bad command id")
+
+    # Do filesystem work BEFORE taking the lock (a 100 MB copy must not stall
+    # every concurrent poll/UI handler). Enforce the size cap server-side too —
+    # don't trust the client.
+    saved = None
+    fname = sz = None
+    too_big = False
+    if fileobj:
+        sz = os.path.getsize(fileobj)
+        if sz > MAX_FILE_BYTES:
+            too_big = True
+        else:
+            fname = os.path.basename(fileobj)
+            saved = FILES_DIR / f"{command_id}__{fname}"
+            shutil.copy(fileobj, saved)
+
     with LOCK:
         rec = COMMANDS.get(command_id)
         if rec is None:  # relay restarted — keep the file anyway
@@ -152,15 +238,14 @@ def api_upload_file(command_id: str, client_id: str, secret: str, fileobj, error
                 "type": "get", "command": "(file fetch — relay restarted)",
                 "created": _now(), "status": "pending",
             }
-            BY_CLIENT[client_id].append(command_id)
+            _append_and_trim(client_id, command_id)
         rec["status"] = "done"
         rec["finished_at"] = _now()
         _touch_client(client_id)
-        if fileobj:
-            fname = os.path.basename(fileobj)
-            saved = FILES_DIR / f"{command_id}__{fname}"
-            shutil.copy(fileobj, saved)
-            sz = saved.stat().st_size
+        if too_big:
+            rec.update(exit_code=1,
+                       stdout=f"file exceeds the {_human(MAX_FILE_BYTES)} limit")
+        elif saved is not None:
             rec.update(stored_path=str(saved), filename=fname, size=sz,
                        exit_code=0, stdout=f"fetched {fname} ({_human(sz)})")
         else:
@@ -200,12 +285,7 @@ def ui_send(client_label: str, session: str, command: str, timeout: int):
             "command": command, "timeout": int(timeout),
             "status": "pending", "created": _now(),
         }
-        hist = BY_CLIENT[cid]
-        hist.append(cmd_id)
-        # trim ring buffer
-        while len(hist) > MAX_HISTORY_PER_CLIENT:
-            old = hist.pop(0)
-            COMMANDS.pop(old, None)
+        _append_and_trim(cid, cmd_id)
     return f"➡️ sent to `{cid}` [{session}]", ""
 
 
@@ -233,7 +313,7 @@ def ui_put_file(client_label: str, fileobj, dest: str):
             "command": f"📤 send {fname} ({_human(sz)}) → {dest}",
             "timeout": 300, "status": "pending", "created": _now(),
         }
-        BY_CLIENT[cid].append(cmd_id)
+        _append_and_trim(cid, cmd_id)
     return f"➡️ sending `{fname}` to `{cid}`:`{dest}`", None
 
 
@@ -252,7 +332,7 @@ def ui_get_file(client_label: str, remote_path: str):
             "remote_path": remote_path, "command": f"📥 fetch {remote_path}",
             "timeout": 300, "status": "pending", "created": _now(),
         }
-        BY_CLIENT[cid].append(cmd_id)
+        _append_and_trim(cid, cmd_id)
     return f"➡️ requested `{remote_path}` from `{cid}`"
 
 
@@ -330,7 +410,7 @@ def ui_tick(client_label: str, last_choices: list[str], last_hist: str,
 # --------------------------------------------------------------------------- #
 # Layout                                                                        #
 # --------------------------------------------------------------------------- #
-with gr.Blocks(title="NII Relay") as demo:
+with gr.Blocks(title="NII Relay", analytics_enabled=False) as demo:
     gr.Markdown(
         "# 🛰️ NII Relay\n"
         "Send shell commands to no-SSH containers via a polling daemon. "
@@ -421,7 +501,12 @@ with gr.Blocks(title="NII Relay") as demo:
                        a_out, api_name="upload_result")
 
 
+threading.Thread(target=_file_reaper, daemon=True).start()
+
 if __name__ == "__main__":
+    if not SECRET:
+        print("WARNING: RELAY_SECRET is empty — all daemon API calls will be rejected "
+              "(fail-closed). Set it as a Space secret.", flush=True)
     demo.queue(default_concurrency_limit=16).launch(
         server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)),
         theme=gr.themes.Soft(), max_file_size="150mb",