5 Commits
Author SHA1 Message Date
XiaodoudouandEdouard Durand 6041a1d72a RigShare 1.4.0: unload models after an idle period (#4)
New admin option (auto_unload_minutes, off by default): once nothing has run for that many minutes, the hub sets ComfyUI's unload_models queue flag, as the Unload models button does. A run or a waiting prompt restarts the clock; it unloads once per idle period. auto_unload_cache also clears the node cache (free_memory).

Co-authored-by: Edouard Durand <xiaodoudou@users.noreply.github.com>
2026-10-26 01:26:28 +08:00
Xiaodoudou 3de3cc32e2 RigShare 1.3.3: trust follows the real client behind a proxy, login throttling, cross-site and framing protection
* RigShare 1.3.1: refuse folder paths outside the workspace (Windows)

* Tests: escape test no longer assumes Windows path semantics

* RigShare 1.3.2: access checks without a live room, refuse Windows alias paths, same-origin login redirect

* RigShare 1.3.3: trust follows the real client behind a proxy, login throttling, cross-site and framing protection
2026-10-08 19:32:03 +08:00
Xiaodoudou 0b315f6647 RigShare 1.3.2: access checks without a live room, refuse Windows alias paths, same-origin login redirect
* RigShare 1.3.1: refuse folder paths outside the workspace (Windows)

* Tests: escape test no longer assumes Windows path semantics

* RigShare 1.3.2: access checks without a live room, refuse Windows alias paths, same-origin login redirect
2026-10-08 19:29:42 +08:00
Xiaodoudou 1befdca2e7 RigShare 1.3.1: refuse folder paths outside the workspace (Windows)
* RigShare 1.3.1: refuse folder paths outside the workspace (Windows)

* Tests: escape test no longer assumes Windows path semantics
2026-10-08 16:35:36 +08:00
Edouard Durand d68307cde4 README: GIF of two people editing one workflow live 2026-10-04 07:58:07 +00:00
16 changed files with 655 additions and 45 deletions
+48
View File
@@ -1,5 +1,53 @@
# Changelog
## 1.4.0 (2026-10-25)
Models can leave the GPU on their own once the server is idle.
### Added
- Admin option *Unload models after idle*: once nothing has run for that many minutes, RigShare unloads the models, as the *Unload models* button does, so the VRAM is free for other programs. Off by default (`auto_unload_minutes`: `0`). A run or a waiting prompt restarts the clock, and it unloads once per idle period
- With *Also clear the cache when unloading* (`auto_unload_cache`), it clears ComfyUI's cache too, like *Free models & cache*: RAM is freed as well, but the next run computes every node again
## 1.3.3 (2026-10-08)
Login and request hardening, mainly for servers behind a reverse proxy.
### Fixed
- A reverse proxy on the same machine (nginx, Caddy, a tunnel) made every visitor look like `127.0.0.1`, which is a trusted address by default, so login and API protection were skipped for everyone. Behind a trusted proxy, trust now follows the visitor's own address (the last `X-Forwarded-For` / `X-Real-IP` entry), and a `Forwarded` header RigShare cannot read means no trust
- Login throttling counted only the address, so users behind one proxy locked each other out, one account could be guessed from many addresses, and the counters never shrank. It now limits per visitor address (5 failures a minute) and per account (10 a minute across all addresses), answers with `Retry-After`, and forgets old entries. People already signed in are never affected. Failed logins are logged
### Added
- Requests that carry the session cookie and an `Origin` from another site are refused: writes and the live connection (hostnames compared, so a proxy that drops the port or sets `X-Forwarded-Host` still works). API clients without a cookie are not affected
- The login page is sent with `X-Frame-Options: DENY`, a `frame-ancestors 'none'` policy, `nosniff` and a same-origin referrer policy, so another site cannot frame it
## 1.3.2 (2026-10-08)
Access rules tightened, following the review of the ComfyUI Manager listing.
### Fixed
- Snapshots of a workflow could be listed and read by an account with no access to its folder whenever nobody had that workflow open. Snapshot access now follows the workflow's folder (private and restricted shared folders included), open or not
- Joining a live workflow seeded a new room from the caller's document before the access check. Access is checked first, and only someone who can edit the workflow can open a new room
- Windows path segments that end in a dot or space, or use an 8.3 short name, are now refused instead of being silently dropped, so a path can no longer be classified as a different folder than the one Windows opens
- After signing in, the login page only follows its `next` redirect when it resolves to the same origin
## 1.3.1 (2026-10-08)
Security hardening of the workspace folder routes, mainly for Windows servers.
### Fixed
- Folder rename, move, delete and create (and the workspace listing) could be pointed outside the workflows folder with a Windows drive prefix (`C:`), which let `rmtree` / `rename` act on folders such as models or custom_nodes. Every path is now checked before it touches the disk: segments with a colon, backslash tricks or `..` are refused, and the resolved path must stay inside the workflows folder
- On Windows, case and aliases of a path were not seen by the access checks: `Users/alice`, `users./alice` or `USERS~1/alice` opened Alice's private folder as if it were common, and `Shared/<restricted folder>` skipped its access list. Paths are now matched to the name on disk and trailing dots, spaces and 8.3 short names are dropped before any permission check
- New usernames and folder names can no longer end in a dot or space, start with a dot, or be a Windows device name (`CON`, `NUL`, `COM1`...). Usernames that differ only by case count as the same name, since they share one folder on Windows. Existing accounts and folders are untouched
### Added
- Tests for the JWT module, the server stats, and the folder routes refusing paths outside the workspace
## 1.3.0 (2026-10-04)
A Performance panel, a floating monitor that docks, and Manager in its own tab.
+8
View File
@@ -15,6 +15,14 @@ A custom node that turns a single ComfyUI server into a shared workspace. Everyo
---
<div align="center">
<img src="docs/images/live-cursors.gif" width="900" alt="Two people editing the same workflow at once, each seeing the other's named cursor and selection">
<sub>Two accounts, one workflow: each person sees the other's cursor, selection and edits live.</sub>
</div>
## What it does
- asks everyone to sign in before ComfyUI loads, and creates the admin account on first run
Binary file not shown.

After

Width:  |  Height:  |  Size: 4.0 MiB

+2
View File
@@ -36,6 +36,8 @@ Admins change these under **Admin → Server settings**. They are stored in `rig
| `snapshot_interval_sec` | `300` | Automatic snapshot interval while editing; `0` turns it off |
| `snapshot_keep` | `30` | Automatic snapshots kept per workflow, and as many on-save ones; named ones are always kept |
| `chat_keep` | `0` | Chat messages kept: `0` keeps the whole history, a number trims it to the latest ones at startup. Admins can also clear it under *Admin → Chat* |
| `auto_unload_minutes` | `0` | Unload models from VRAM once nothing has run for this many minutes, as the *Unload models* button does; `0` turns it off. A run or a waiting prompt restarts the clock |
| `auto_unload_cache` | `false` | When unloading, also clear ComfyUI's cache (*Free models & cache*): frees RAM too, but the next run computes every node again |
## Environment variables
+1 -1
View File
@@ -1,7 +1,7 @@
[project]
name = "comfyui-rigshare"
description = "Turn one ComfyUI server into a shared team workspace. Sign-in with accounts, permissions (edit, queue, Manager, admin) and personal API keys, enforced by the server. A private folder for each person plus shared folders with access lists, with Save asking where. Live co-editing of every shared workflow, with cursors, follow mode and conflict-free merges, in the graph or App view. Per-workflow history with a snapshot on every save, team chat with its whole history, and a CPU, RAM, GPU and queue monitor."
version = "1.3.0"
version = "1.4.0"
license = { file = "LICENSE" }
dependencies = []
+53 -10
View File
@@ -110,6 +110,8 @@ class Hub:
self.workspace = None # set by __init__ (folders, private spaces)
self.queue_owners = OrderedDict() # prompt_id -> {user, name, workflow, at}
self.run_started = {} # prompt_id -> ms timestamp of its execution_start (set by routes)
self.last_activity = time.time() # last run event or busy queue, for auto unload
self.models_idle = True # nothing to unload until a run loads models
# ----- lifecycle ----------------------------------------------------
@@ -135,6 +137,7 @@ class Hub:
if expiry > 0 and now - room.updated > expiry and not self.members(room.key):
self.rooms.pop(room.key, None)
self.store.delete_room(room.key)
self.auto_unload(now)
except Exception as e:
log.error(f"[RigShare] housekeeping error: {e}")
@@ -183,6 +186,8 @@ class Hub:
prompt_id = data.get("prompt_id") if isinstance(data, dict) else None
if not prompt_id:
return
self.last_activity = time.time()
self.models_idle = False
stamp = data.get("timestamp") or int(time.time() * 1000)
if event == "execution_start":
self.run_started[prompt_id] = stamp
@@ -221,6 +226,31 @@ class Hub:
return {"running": running, "recent": list(reversed(self.store.runs()[-recent:])),
"etas": etas, "now": now}
def auto_unload(self, now):
"""Unload models once nothing has run for ``auto_unload_minutes`` (0 = off).
Sets the same queue flags as ComfyUI's /free, so its prompt worker does the
unloading between runs. Happens once per idle period.
"""
try:
minutes = float(self.store.config.get("auto_unload_minutes") or 0)
except (TypeError, ValueError):
minutes = 0
if minutes <= 0 or self.models_idle or not self.server:
return
queue = self.server.prompt_queue
running, pending = queue.get_current_queue_volatile()
if running or pending:
self.last_activity = now
return
if now - self.last_activity < minutes * 60:
return
queue.set_flag("unload_models", True)
if self.store.config.get("auto_unload_cache"):
queue.set_flag("free_memory", True)
self.models_idle = True
log.info(f"[RigShare] Unloaded models after {minutes:g} min without runs.")
def running_prompt_ids(self):
try:
running, _ = self.server.prompt_queue.get_current_queue_volatile()
@@ -334,8 +364,12 @@ class Hub:
Access follows the workflow's folder (restricted shared folders), within
the person's account permissions.
"""
if self.workspace and room.key.startswith("file:workflows/"):
return self.workspace.role(room.key[len("file:workflows/"):], kind, key, perms)
return self.access_for_key(room.key, kind, key, perms)
def access_for_key(self, room_key, kind, key, perms):
"""Same as access(), from a room key: works before (or without) a live room."""
if self.workspace and room_key.startswith("file:workflows/"):
return self.workspace.role(room_key[len("file:workflows/"):], kind, key, perms)
return "edit" if (perms.get("edit") or perms.get("admin")) else "view"
def client_access(self, room, client):
@@ -556,16 +590,25 @@ class Hub:
else:
room = self.rooms.get(key)
created = False
if room is None:
doc = data.get("doc") if isinstance(data.get("doc"), dict) else {}
room = Room(key, clean_text(data.get("name"), 120) or key, "file", doc)
room.dirty = True
self.rooms[key] = room
created = True
if self.client_access(room, client) is None:
# Check access before anything is created: only someone who may edit the
# workflow can open a room and seed it with their document.
role = self.access_for_key(key, client.kind, client.key, client.perms)
if role is None:
client.room = None
await self.send(client, {"type": "room_denied", "room": key, "name": room.name})
await self.send(client, {"type": "room_denied", "room": key,
"name": room.name if room else clean_text(data.get("name"), 120) or key})
elif room is None and role != "edit":
client.room = None
await self.send(client, {"type": "room_denied", "room": key,
"name": clean_text(data.get("name"), 120) or key,
"reason": "Nobody has opened this workflow for editing yet."})
else:
if room is None:
doc = data.get("doc") if isinstance(data.get("doc"), dict) else {}
room = Room(key, clean_text(data.get("name"), 120) or key, "file", doc)
room.dirty = True
self.rooms[key] = room
created = True
client.room = key
room.updated = time.time()
await self.send(client, {"type": "room", "room": key, "name": room.name, "version": room.version,
+8 -1
View File
@@ -84,6 +84,13 @@
const err = document.getElementById("e");
const btn = document.getElementById("b");
const next = new URLSearchParams(location.search).get("next") || "/";
// Follow `next` only when it resolves to this same origin (blocks //host, /\host, scheme tricks).
const safeNext = () => {
try {
const url = new URL(next, location.origin);
return url.origin === location.origin ? url.pathname + url.search + url.hash : "/";
} catch { return "/"; }
};
if (SETUP) {
document.title = "Set up · {{SERVER_NAME}}";
document.getElementById("sub").textContent = "First-time setup";
@@ -112,7 +119,7 @@
if (res.status === 409) { location.reload(); return; } // someone finished setup meanwhile
if (!res.ok) throw new Error(data.error || "Sign in failed");
try { localStorage.setItem("rigshare.token", data.token); } catch {}
location.replace(next.startsWith("/") && !next.startsWith("//") ? next : "/");
location.replace(safeNext());
} catch (e) {
err.textContent = e.message;
btn.disabled = false;
+89 -21
View File
@@ -9,12 +9,12 @@ import logging
import os
import posixpath
import time
from urllib.parse import unquote
from urllib.parse import unquote, urlsplit
from collections import defaultdict
from aiohttp import web
from .workspace import Workspace, norm
from .workspace import InvalidPath, Workspace, norm as _norm
log = logging.getLogger("ComfyUI-RigShare")
@@ -94,6 +94,9 @@ BROADCAST_EVENTS = {"execution_start", "execution_cached", "executing", "execute
COOKIE = "rigshare_session"
# The login page must not be framed by another site (clickjacking) or sniffed into another type.
SECURITY_HEADERS = {"X-Frame-Options": "DENY", "Content-Security-Policy": "frame-ancestors 'none'",
"X-Content-Type-Options": "nosniff", "Referrer-Policy": "same-origin"}
LOGIN_PAGE = os.path.join(os.path.dirname(__file__), "login.html")
ICON_SVG = os.path.join(os.path.dirname(os.path.dirname(__file__)), "web", "rigshare", "assets", "icon.svg")
# Reachable without logging in.
@@ -112,7 +115,8 @@ def token_from(request):
def setup(server, store, hub):
routes = web.RouteTableDef()
failures = defaultdict(list)
failures = defaultdict(list) # client address -> times of failed logins
failures_by_name = defaultdict(list) # username -> times of failed logins
def current_user(request):
return store.user_for_token(token_from(request))
@@ -148,6 +152,13 @@ def setup(server, store, hub):
def error(message, status=400):
return web.json_response({"error": message}, status=status)
def norm(path):
"""workspace.norm, answering 403 for a path that is refused instead of adjusted."""
try:
return _norm(path)
except InvalidPath:
raise web.HTTPForbidden(text='{"error": "Invalid path"}', content_type="application/json")
def set_cookie(request, response, token):
secure = request.secure or request.headers.get("X-Forwarded-Proto") == "https"
response.set_cookie(COOKIE, token, max_age=30 * 24 * 3600, path="/", httponly=True,
@@ -184,19 +195,30 @@ def setup(server, store, hub):
except OSError:
icon = ""
page = page.replace("{{ICON_SVG}}", icon).replace("{{ICON_B64}}", base64.b64encode(icon.encode()).decode())
return web.Response(text=page, content_type="text/html", headers={"Cache-Control": "no-store"})
return web.Response(text=page, content_type="text/html", headers={"Cache-Control": "no-store", **SECURITY_HEADERS})
@routes.post("/rigshare/api/login")
async def login(request):
ip = request.remote or "?"
now = time.time()
failures[ip] = [t for t in failures[ip] if now - t < 60]
if len(failures[ip]) >= 5:
return error("Too many attempts, wait a minute", 429)
ip, now = client_ip(request) or "?", time.time()
body = await request.json()
user = store.authenticate(str(body.get("username", "")).strip(), str(body.get("password", "")))
name = str(body.get("username", "")).strip()[:64].casefold()
for table in (failures, failures_by_name): # forget old failures, so the tables stay small
for key in [k for k, times in table.items() if not times or now - times[-1] >= 60]:
del table[key]
recent = [t for t in failures[ip] if now - t < 60]
recent_name = [t for t in failures_by_name[name] if now - t < 60]
failures[ip], failures_by_name[name] = recent, recent_name
# Per address: 5 a minute. Per account: 10 a minute across all addresses, which slows
# guessing from many places without locking the owner out for long. Signed-in people
# never come through here.
if len(recent) >= 5 or len(recent_name) >= 10:
return web.json_response({"error": "Too many attempts, wait a minute"}, status=429,
headers={"Retry-After": "60"})
user = store.authenticate(name_raw := str(body.get("username", "")).strip(), str(body.get("password", "")))
if not user:
failures[ip].append(now)
failures_by_name[name].append(now)
log.warning(f"[RigShare] Failed login for '{name_raw[:32]}' from {ip}")
return error("Invalid username or password", 401)
token = store.create_session(user["username"])
response = web.json_response({"token": token, "user": store.public_user(user)})
@@ -580,8 +602,7 @@ def setup(server, store, hub):
if not key:
require_admin(request)
return web.json_response(store.list_snapshots())
room = hub.rooms.get(key)
if room and not room_access(request, room):
if hub.access_for_key(key, *who(request)) is None:
return error("No such room", 404)
return web.json_response(store.list_snapshots(key))
@@ -605,8 +626,9 @@ def setup(server, store, hub):
snap = store.get_snapshot(request.match_info["snap_id"])
if not snap:
return error("Not found", 404)
room = hub.rooms.get(snap["meta"].get("room", ""))
if room and not room_access(request, room):
# Access follows the workflow's folder whether or not anyone has it open right now.
room_key = snap["meta"].get("room", "")
if room_key and hub.access_for_key(room_key, *who(request)) is None:
return error("Not found", 404)
return web.json_response(snap)
@@ -644,6 +666,50 @@ def setup(server, store, hub):
continue
return False
def client_ip(request):
"""The caller's address. Behind a trusted proxy (one of trusted_ips), the address that
proxy appended to X-Forwarded-For / X-Real-IP; otherwise the socket's. A proxy on the
same machine used to make every visitor look like 127.0.0.1, so everyone was trusted."""
remote = request.remote or ""
if not trusted(remote):
return remote # an untrusted peer's forwarding headers are just claims
forwarded = request.headers.get("X-Forwarded-For")
if forwarded:
return forwarded.split(",")[-1].strip()
real = request.headers.get("X-Real-IP")
if real:
return real.strip()
if request.headers.get("Forwarded"):
return "" # a proxy we cannot read: do not trust the visitor
return remote
def cross_site(request):
"""A browser sending its session cookie from another site. Browsers always send Origin on
cross-site writes and websockets; API clients (no cookie, usually no Origin) are not
affected."""
origin = request.headers.get("Origin")
if not origin or COOKIE not in request.cookies:
return False
if request.method in ("GET", "HEAD", "OPTIONS") and request.path != "/rigshare/ws":
return False
# Hostnames only: a proxy often passes Host without the port the browser used.
def hostname(value):
try:
return (urlsplit(value if "//" in value else "//" + value).hostname or "").lower()
except ValueError:
return ""
origin_host = hostname(origin)
allowed = {hostname(request.host)}
allowed.update(hostname(h.strip()) for h in (request.headers.get("X-Forwarded-Host") or "").split(",") if h.strip())
if origin_host and origin_host in allowed:
return False
log.warning(f"[RigShare] Refused cross-site {request.method} {request.path}: Origin {origin[:80]!r} "
f"does not match Host {request.host[:80]!r}")
return True
def trusted_request(request):
return trusted(client_ip(request))
def parse_userdata(request):
"""('list'|'get'|'save'|'delete'|'move', path, dest) for ComfyUI userdata calls under workflows/."""
raw = request.raw_path.split("?", 1)[0]
@@ -699,7 +765,7 @@ def setup(server, store, hub):
"""Cancelling a queued prompt or stopping the running one: your own, unless admin."""
if request.method != "POST" or api_path(request) not in ("/queue", "/interrupt"):
return None
if trusted(request.remote or "") or perms_for(request).get("admin"):
if trusted_request(request) or perms_for(request).get("admin"):
return None
user = current_user(request)
me = user["username"] if user else None
@@ -798,7 +864,9 @@ def setup(server, store, hub):
@web.middleware
async def protect(request, handler):
path = request.path
if store.config.get("require_login") and path not in PUBLIC_PATHS and not trusted(request.remote or ""):
if cross_site(request):
return web.json_response({"error": "RigShare: cross-site request refused"}, status=403)
if store.config.get("require_login") and path not in PUBLIC_PATHS and not trusted_request(request):
if not current_user(request):
wants_page = request.method == "GET" and (
path in ("/", "/index.html") or "text/html" in request.headers.get("Accept", ""))
@@ -809,20 +877,20 @@ def setup(server, store, hub):
"'Authorization: Bearer <api key>'."}, status=401,
headers={"X-RigShare-Login": "required"})
userdata = parse_userdata(request)
if userdata and not trusted(request.remote or ""):
if userdata and not trusted_request(request):
refusal = check_userdata(request, *userdata)
if refusal:
return web.json_response({"error": f"RigShare: {refusal}"}, status=403)
cfg = store.config.get("api_protection", {})
if cfg.get("enabled"):
path = request.path[4:] if request.path.startswith("/api/") else request.path
if needs_manager(request.method, path) and not trusted(request.remote or ""):
if needs_manager(request.method, path) and not trusted_request(request):
if not perms_for(request).get("manager"):
return web.json_response({"error": "RigShare: you need the 'manager' permission to use ComfyUI Manager"},
status=403)
for method, prefix, need in PROTECTED:
if (not method or request.method == method) and (path == prefix or path.startswith(prefix)):
if trusted(request.remote or ""):
if trusted_request(request):
break
perms = perms_for(request)
if not any(perms.get(p) for p in need.split("|")):
@@ -834,7 +902,7 @@ def setup(server, store, hub):
break
for method, prefix, need, why in PLUGIN_RULES:
if (method is None or request.method == method) and api_path(request).startswith(prefix):
if not trusted(request.remote or "") and not perms_for(request).get(need) and not perms_for(request).get("admin"):
if not trusted_request(request) and not perms_for(request).get(need) and not perms_for(request).get("admin"):
return web.json_response({"error": f"RigShare: {why}"}, status=403)
break
refusal = await check_queue_ownership(request)
@@ -855,7 +923,7 @@ def setup(server, store, hub):
snapshot_saved_file(request, userdata[1], await request.read())
await hub.files_changed(userdata[1])
# Workflow listings only show what this person may open.
if userdata and userdata[0] == "list" and not trusted(request.remote or "") and hub.workspace:
if userdata and userdata[0] == "list" and not trusted_request(request) and hub.workspace:
response = filter_listing(request, response, userdata[1])
# A file moved on disk: its live room follows it.
if userdata and userdata[0] == "move" and getattr(response, "status", 0) == 200:
+11 -3
View File
@@ -24,7 +24,7 @@ SESSION_TTL = 30 * 24 * 3600 # a login lasts at most this long
TOKEN_TTL = 12 * 3600 # each JWT is short-lived...
TOKEN_REFRESH = 6 * 3600 # ...and re-issued while in use once older than this
ISSUER = "rigshare"
USERNAME_RE = re.compile(r"^[A-Za-z0-9_.-]{2,32}$")
USERNAME_RE = re.compile(r"^(?!(?i:con|prn|aux|nul|com\d|lpt\d)(?:\.|$))[A-Za-z0-9_][A-Za-z0-9_.-]{0,30}[A-Za-z0-9_-]\Z") # no leading/trailing dot, no Windows device names
PERMS = ("edit", "queue", "manager", "admin")
API_KEY_PREFIX = "rs_"
@@ -58,6 +58,10 @@ DEFAULT_CONFIG = {
"snapshot_interval_sec": 300,
"snapshot_keep": 30,
"chat_keep": 0, # 0 keeps the whole chat; N trims it to the last N messages at startup
# Unload models once the queue has been empty this many minutes (0 = off),
# and with auto_unload_cache also clear ComfyUI's node cache.
"auto_unload_minutes": 0,
"auto_unload_cache": False,
}
@@ -216,10 +220,14 @@ class Store:
_verify_password(password, _hash_password("x", b"0" * 16))
return None
def _name_taken(self, name, ignore=None):
"""Names differing only by case share one folder on Windows, so they count as the same."""
return any(u.casefold() == name.casefold() for u in self.users if u != ignore)
def create_user(self, username, password, display_name=None, perms=None):
if not USERNAME_RE.match(username or ""):
raise ValueError("Username must be 2-32 characters: letters, digits, _ . -")
if username in self.users:
if self._name_taken(username):
raise ValueError("Username already exists")
if len(password or "") < 8:
raise ValueError("Password must be at least 8 characters")
@@ -239,7 +247,7 @@ class Store:
raise KeyError(username)
if not USERNAME_RE.match(new_username):
raise ValueError("Username must be 2-32 characters: letters, digits, _ . -")
if new_username in self.users:
if self._name_taken(new_username, ignore=username):
raise ValueError("Username already exists")
with self._lock:
user = self.users.pop(username)
+67 -9
View File
@@ -20,7 +20,7 @@ import time
from urllib.parse import unquote
WORKFLOWS = "workflows"
FOLDER_NAME_RE = re.compile(r"^[\w][\w .()-]{0,63}$")
FOLDER_NAME_RE = re.compile(r"^(?!(?i:con|prn|aux|nul|com\d|lpt\d)(?:\.|$))\w(?:[\w .()-]{0,62}[\w()-])?\Z") # no trailing dot/space, no Windows device names
ROLE_RANK = {None: 0, "view": 1, "edit": 2}
@@ -35,10 +35,66 @@ def full_unquote(value):
return value
_WIN = os.name == "nt"
_SHORT_NAME_RE = re.compile(r"~\d") # Windows 8.3 alias of a long name (USERS~1)
_case_root = None # workflows directory used to restore on-disk casing (Windows only)
def _on_disk_case(parts, base):
"""Windows ignores case: map each segment to the casing it has on disk, so 'Users/Bob' is
judged as the same path it will open."""
out = []
for part in parts:
try:
names = os.listdir(base) if base else []
except OSError:
names = []
match = part if part in names else next((n for n in names if n.casefold() == part.casefold()), part)
out.append(match)
base = os.path.join(base, match) if base and match in names else None
return out
class InvalidPath(PermissionError):
"""A path Windows would read as a different one than the checks do: refused, never adjusted."""
def norm(path):
"""Decode and normalise a userdata-relative path ('' for the root)."""
"""Decode and normalise a userdata-relative path ('' for the root).
A path with a segment Windows would read as something else is refused (InvalidPath), not
adjusted: a colon (drive prefix 'C:', NTFS stream), a 8.3 short name, or trailing dots or
spaces (Windows strips them, so 'users./bob' would open 'users/bob'). Casing is restored from
disk on Windows. Anything that still reaches the disk goes through safe_join().
"""
path = full_unquote(str(path or "")).replace("\\", "/")
return posixpath.normpath("/" + path).lstrip("/") if path else ""
if not path:
return ""
clean = posixpath.normpath("/" + path).lstrip("/")
parts = [p for p in clean.split("/") if p]
for p in parts:
if ":" in p or (_WIN and (_SHORT_NAME_RE.search(p) or p != p.rstrip(" ."))):
raise InvalidPath("Invalid path")
if _WIN and _case_root and parts:
lead = [WORKFLOWS] if parts[0].casefold() == WORKFLOWS else []
parts = lead + _on_disk_case(parts[len(lead):], _case_root)
return "/".join(parts)
def safe_join(root, *parts):
"""Join relative parts under root; refuse anything that resolves outside it."""
segments = [p for part in parts for p in str(part).replace("\\", "/").split("/") if p]
if any(":" in p or p in (".", "..") or (_WIN and p != p.rstrip(" .")) for p in segments):
raise PermissionError("Invalid path")
full = os.path.join(root, *segments)
base = os.path.realpath(root)
try:
inside = os.path.commonpath([base, os.path.realpath(full)]) == base
except ValueError: # different drives
inside = False
if not inside:
raise PermissionError("Invalid path")
return full
def has_app(doc):
@@ -56,6 +112,8 @@ class Workspace:
def __init__(self, store, root):
self.store = store
self.root = root # absolute path of the workflows directory
global _case_root
_case_root = root
self.folders = store._read("folders.json", {}) # "shared/<name>" -> {owner, acl, created}
os.makedirs(os.path.join(root, "users"), exist_ok=True)
os.makedirs(os.path.join(root, "shared"), exist_ok=True)
@@ -147,7 +205,7 @@ class Workspace:
return self.role(rel, kind, key, perms) is not None
def _files(self, rel_dir):
base = os.path.join(self.root, *rel_dir.split("/")) if rel_dir else self.root
base = safe_join(self.root, rel_dir) if rel_dir else self.root
out = []
if not os.path.isdir(base):
return out
@@ -198,7 +256,7 @@ class Workspace:
# ----- browsing (Files tab) -------------------------------------------
def _abs(self, rel):
return os.path.join(self.root, *[p for p in rel.split("/") if p])
return safe_join(self.root, rel)
def _file_has_app(self, full):
"""Read a workflow file only when it changed since the last check."""
@@ -372,7 +430,7 @@ class Workspace:
name = (name or "").strip()
if not FOLDER_NAME_RE.match(name) or name in (".", ".."):
raise ValueError("Folder names: up to 64 letters, digits, spaces and . _ - ( )")
path = os.path.join(self.root, "shared", name)
path = safe_join(self.root, "shared", name)
if os.path.exists(path):
raise ValueError("A folder with that name already exists")
os.makedirs(path)
@@ -385,8 +443,8 @@ class Workspace:
if not FOLDER_NAME_RE.match(new_name):
raise ValueError("Folder names: up to 64 letters, digits, spaces and . _ - ( )")
new_folder = f"shared/{new_name}"
src = os.path.join(self.root, *folder.split("/"))
dst = os.path.join(self.root, "shared", new_name)
src = safe_join(self.root, folder)
dst = safe_join(self.root, "shared", new_name)
if os.path.exists(dst):
raise ValueError("A folder with that name already exists")
os.rename(src, dst)
@@ -395,7 +453,7 @@ class Workspace:
return new_folder
def delete_folder(self, folder):
path = os.path.join(self.root, *folder.split("/"))
path = safe_join(self.root, folder)
if self._files(folder):
raise ValueError("Move or delete the workflows inside it first")
shutil.rmtree(path, ignore_errors=True)
+4
View File
@@ -38,10 +38,14 @@ class FakeQueue:
def __init__(self):
self.running, self.pending, self.number = [], [], 0
self.flags = {}
def get_current_queue_volatile(self):
return list(self.running), list(self.pending)
def set_flag(self, name, data):
self.flags[name] = data
history = {}
def get_history(self, prompt_id=None, max_items=None, offset=-1, map_function=None):
+60
View File
@@ -0,0 +1,60 @@
"""HS256 tokens: round trip, tampering, expiry, issuer, algorithm confusion."""
import base64
import json
import time
import pytest
from rigshare_core import jwt
SECRET = b"s3cret"
def claims(**extra):
return {"iss": "rig", "exp": time.time() + 60, **extra}
def b64(obj):
return base64.urlsafe_b64encode(json.dumps(obj).encode()).rstrip(b"=").decode()
def test_round_trip():
token = jwt.encode(claims(sub="bob"), SECRET)
assert jwt.decode(token, SECRET, "rig")["sub"] == "bob"
def test_wrong_secret_and_tampered_payload():
token = jwt.encode(claims(sub="bob"), SECRET)
with pytest.raises(jwt.InvalidToken):
jwt.decode(token, b"other", "rig")
head, _, sig = token.split(".")
forged = f"{head}.{b64(claims(sub='admin'))}.{sig}"
with pytest.raises(jwt.InvalidToken):
jwt.decode(forged, SECRET, "rig")
def test_expiry_leeway_and_not_before():
with pytest.raises(jwt.InvalidToken):
jwt.decode(jwt.encode({"iss": "rig", "exp": time.time() - 3600}, SECRET), SECRET, "rig")
jwt.decode(jwt.encode({"iss": "rig", "exp": time.time() - 5}, SECRET), SECRET, "rig") # inside leeway
assert jwt.decode(jwt.encode({"iss": "rig", "exp": time.time() - 3600}, SECRET), SECRET, "rig", verify_exp=False)
with pytest.raises(jwt.InvalidToken):
jwt.decode(jwt.encode(claims(nbf=time.time() + 3600), SECRET), SECRET, "rig")
def test_wrong_issuer():
with pytest.raises(jwt.InvalidToken):
jwt.decode(jwt.encode(claims(), SECRET), SECRET, "someone-else")
def test_alg_none_is_refused():
token = f"{b64({'alg': 'none', 'typ': 'JWT'})}.{b64(claims())}."
with pytest.raises(jwt.InvalidToken):
jwt.decode(token, SECRET, "rig")
@pytest.mark.parametrize("junk", ["", "a.b", "a.b.c.d", "!!!.???.***", "e30.e30.e30"])
def test_malformed_tokens(junk):
with pytest.raises(jwt.InvalidToken):
jwt.decode(junk, SECRET, "rig")
+188
View File
@@ -5,6 +5,8 @@ import json
import os
from urllib.parse import quote
from aiohttp import ClientSession
from helpers import PASSWORD, enc, rig_server
BASE_DOC = {"last_node_id": 2, "last_link_id": 0, "nodes": [
@@ -646,3 +648,189 @@ def test_performance_run_log_eta_history_and_clear(run):
assert (await r.json())["cleared"] == 3
assert hub.etas() == {"video": {"ms": 90000, "runs": 1}}, "the ETA starts over"
run(scenario())
def test_auto_unload_after_idle(run):
async def scenario():
async with rig_server() as rig:
import time
admin = await rig.setup_admin()
hub, queue = rig.hub, rig.ps.prompt_queue
def run(pid):
hub.run_event("execution_start", {"prompt_id": pid, "timestamp": 1})
hub.run_event("execution_success", {"prompt_id": pid, "timestamp": 2})
r = await rig.http.patch(rig.url("/rigshare/api/config"), json={"auto_unload_minutes": 15}, headers=admin)
assert (await r.json())["auto_unload_minutes"] == 15
hub.auto_unload(time.time() + 3600)
assert queue.flags == {}, "nothing ran since startup, nothing to unload"
run("a")
now = time.time()
hub.auto_unload(now + 14 * 60)
assert queue.flags == {}, "not idle long enough"
queue.pending.append((1, "b", {}, {}, []))
hub.auto_unload(now + 3600)
assert queue.flags == {}, "a waiting prompt keeps models loaded and restarts the clock"
queue.pending.clear()
hub.auto_unload(now + 3600 + 14 * 60)
assert queue.flags == {}
hub.auto_unload(now + 3600 + 15 * 60)
assert queue.flags == {"unload_models": True}
queue.flags.clear()
hub.auto_unload(now + 4 * 3600)
assert queue.flags == {}, "once per idle period"
await rig.http.patch(rig.url("/rigshare/api/config"), json={"auto_unload_cache": True}, headers=admin)
run("c")
hub.auto_unload(time.time() + 15 * 60)
assert queue.flags == {"unload_models": True, "free_memory": True}
queue.flags.clear()
await rig.http.patch(rig.url("/rigshare/api/config"), json={"auto_unload_minutes": 0}, headers=admin)
run("d")
hub.auto_unload(time.time() + 24 * 3600)
assert queue.flags == {}, "0 turns it off"
run(scenario())
def test_folder_routes_refuse_paths_outside_the_workspace(run, tmp_path):
async def scenario():
async with rig_server() as rig:
await rig.setup_admin()
await rig.add_user("alice", edit=True)
outside = tmp_path / "models"
outside.mkdir()
(outside / "keep.bin").write_bytes(b"x")
H, h = rig.headers, rig.http
targets = [str(outside), "C:" + str(outside)[2:], r"..\..\models", "shared/C:/evil"]
for route, extra in (("delete-folder", {}), ("rename-folder", {"name": "x"}), ("move-folder", {"dest": "@shared"})):
for target in targets:
r = await h.post(rig.url(f"/rigshare/api/workspace/{route}"), json={"path": target, **extra}, headers=H["alice"])
assert r.status in (400, 403, 404), (route, target, r.status)
r = await h.post(rig.url("/rigshare/api/workspace/mkdir"), json={"parent": "C:/", "name": "x"}, headers=H["alice"])
assert r.status in (400, 403, 404)
assert (outside / "keep.bin").exists()
run(scenario())
def test_access_is_checked_without_a_live_room(run):
async def scenario():
async with rig_server() as rig:
await rig.setup_admin()
alice = await rig.add_user("alice", edit=True)
bob = await rig.add_user("bob", edit=True)
await rig.add_user("viewer")
h = rig.http
r = await h.post(rig.url("/rigshare/api/workspace/mkdir"), json={"parent": "@shared", "name": "Team"}, headers=alice)
assert r.status == 200
rig.hub.workspace.set_acl("shared/Team", {"members": {"alice": "edit"}}, "alice")
path = "workflows/shared/Team/plan.json"
doc = json.dumps({"nodes": [{"id": 1}], "links": []}).encode()
assert (await h.post(rig.url(f"/api/userdata/{enc(path)}"), data=doc, headers=alice)).status == 200
key = "file:" + path
assert key not in rig.hub.rooms, "nobody has it open"
snap_id = rig.store.list_snapshots(key)[0]["id"]
for who_, status in ((alice, 200), (bob, 404)):
r = await h.get(rig.url(f"/rigshare/api/snapshots?room={enc(key)}"), headers=who_)
assert r.status == status, who_
r = await h.get(rig.url(f"/rigshare/api/snapshots/{snap_id}"), headers=who_)
assert r.status == status, who_
# Joining a workflow you cannot open must not create (or seed) a room.
sock = await rig.ws("bob")
await sock.send({"type": "join", "room": key, "name": "plan", "doc": BASE_DOC})
await sock.until("room_denied")
assert key not in rig.hub.rooms
# A viewer cannot be the one who opens a room, so cannot seed it with their own document.
open_key = "file:workflows/common.json"
viewer_sock = await rig.ws("viewer")
await viewer_sock.send({"type": "join", "room": open_key, "name": "common", "doc": BASE_DOC})
assert "editing" in (await viewer_sock.until("room_denied"))["reason"]
assert open_key not in rig.hub.rooms
run(scenario())
def test_a_local_proxy_does_not_make_visitors_trusted(run):
async def scenario():
async with rig_server() as rig:
await rig.setup_admin()
rig.store.update_config({"api_protection": {"enabled": True, "trusted_ips": ["127.0.0.1"]}})
anon = {"Accept": "application/json"}
async with ClientSession() as plain:
# Straight from the machine itself: trusted, no login asked.
assert (await plain.get(rig.url("/api/prompt"), headers=anon)).status != 401
# The same machine relaying an outside visitor (nginx, Caddy, a tunnel): not trusted.
for header in ({"X-Forwarded-For": "203.0.113.9"},
{"X-Forwarded-For": "10.0.0.1, 203.0.113.9"},
{"X-Real-IP": "203.0.113.9"}, {"Forwarded": "for=203.0.113.9"}):
r = await plain.get(rig.url("/api/prompt"), headers={**anon, **header})
assert r.status == 401, header
# A proxy relaying someone on the trusted network keeps that trust.
r = await plain.get(rig.url("/api/prompt"), headers={**anon, "X-Forwarded-For": "127.0.0.1"})
assert r.status != 401
run(scenario())
def test_login_attempts_are_throttled_per_address_and_per_account(run):
async def scenario():
async with rig_server() as rig:
await rig.setup_admin()
await rig.add_user("alice", edit=True)
h, login = rig.http, rig.url("/rigshare/api/login")
bad = {"username": "alice", "password": "wrong-password"}
for _ in range(5):
assert (await h.post(login, json=bad)).status == 401
r = await h.post(login, json=bad)
assert r.status == 429 and r.headers["Retry-After"] == "60"
assert (await h.post(login, json={"username": "alice", "password": PASSWORD})).status == 429, "even the right password waits"
# Someone already signed in is not affected: no login involved.
assert (await h.get(rig.url("/rigshare/api/status"), headers=rig.headers["alice"])).status == 200
# Many addresses (a trusted proxy relaying them) against one account.
rig.store.update_config({"api_protection": {"enabled": True, "trusted_ips": ["127.0.0.1"]}})
statuses = []
for n in range(12):
r = await h.post(login, json={"username": "Alice", "password": "wrong-password"},
headers={"X-Forwarded-For": f"203.0.113.{n + 1}"})
statuses.append(r.status)
# The 5 failures above already count against the account: 5 more fit, then it waits.
assert statuses[:5] == [401] * 5 and statuses[5:] == [429] * 7
r = await h.post(login, json={"username": "someone-else", "password": "x"}, headers={"X-Forwarded-For": "198.51.100.7"})
assert r.status == 401, "other accounts are not slowed down"
run(scenario())
def test_cross_site_requests_with_a_session_cookie_are_refused(run):
async def scenario():
async with rig_server() as rig:
await rig.setup_admin()
h = rig.http
assert any(c.key == "rigshare_session" for c in h.cookie_jar), "logged in by cookie"
own = {"Origin": rig.url("")}
evil = {"Origin": "https://evil.example"}
r = await h.post(rig.url("/rigshare/api/snapshots"), json={}, headers=evil)
assert r.status == 403 and "cross-site" in (await r.json())["error"]
assert (await h.post(rig.url("/rigshare/api/snapshots"), json={}, headers=own)).status != 403
assert (await h.post(rig.url("/rigshare/api/snapshots"), json={})).status != 403, "no Origin: API client"
assert (await h.get(rig.url("/rigshare/api/status"), headers=evil)).status == 200, "reads are not writes"
r = await h.get(rig.url("/rigshare/ws"), headers={**evil, "Upgrade": "websocket", "Connection": "Upgrade"})
assert r.status == 403
# Behind a proxy that rewrites Host, X-Forwarded-Host names the public host.
r = await h.post(rig.url("/rigshare/api/snapshots"), json={},
headers={"Origin": "https://comfy.example.com", "X-Forwarded-Host": "comfy.example.com"})
assert r.status != 403
run(scenario())
def test_login_page_cannot_be_framed(run):
async def scenario():
async with rig_server() as rig:
r = await rig.http.get(rig.url("/rigshare/login"))
assert r.headers["X-Frame-Options"] == "DENY"
assert "frame-ancestors 'none'" in r.headers["Content-Security-Policy"]
assert r.headers["X-Content-Type-Options"] == "nosniff"
run(scenario())
+45
View File
@@ -0,0 +1,45 @@
"""Server stats: every metric is optional and never raises."""
from rigshare_core import stats
class Queue:
def __init__(self, running, pending):
self.running, self.pending = running, pending
def get_current_queue_volatile(self):
return self.running, self.pending
class Server:
def __init__(self, queue):
self.prompt_queue = queue
def test_queue_items_order_and_node_count():
server = Server(Queue([(1, "r1", {"1": {}, "2": {}})], [(5, "p5", {"1": {}}), (3, "p3", {})]))
items = stats.queue_items(server)
assert [(i["id"], i["status"], i["nodes"]) for i in items] == [("r1", "running", 2), ("p3", "pending", 0), ("p5", "pending", 1)]
def test_queue_items_survive_odd_entries_and_missing_queue():
assert stats.queue_items(Server(Queue([(1,)], [(9,)]))) == []
assert stats.queue_items(object()) is None
def test_queue_counts_with_fallback():
assert stats._queue(Server(Queue([1], [2, 3]))) == {"running": 1, "pending": 2}
class Old:
def get_tasks_remaining(self):
return 4
assert stats._queue(Server(Old())) == {"running": 0, "pending": 4}
assert stats._queue(object()) is None
def test_collect_stats_without_psutil_or_gpu(monkeypatch):
monkeypatch.setattr(stats, "psutil", None)
monkeypatch.setattr(stats, "pynvml", None)
out = stats.collect_stats(Server(Queue([], [])))
assert out["gpus"] == [] and "cpu" not in out and out["queue"] == {"running": 0, "pending": 0}
assert stats.collect_stats() == {"gpus": []}
+65
View File
@@ -131,3 +131,68 @@ def test_rename_user_moves_private_folder(ws):
ws.rename_user("alice", "alicia")
assert os.path.exists(os.path.join(ws.root, "users", "alicia", "keep.json"))
assert ws.folders["shared/Team"]["owner"] == "alicia"
@pytest.mark.parametrize("raw", ["C:/models", "C:", "a/D:x/b", "c%3A%5Cmodels"])
def test_norm_refuses_drive_segments(raw):
with pytest.raises(PermissionError):
norm(raw)
def test_norm_still_resolves_dotdot_inside_the_tree():
assert norm(r"..\..\x") == "x"
assert norm("a/../b") == "b"
@pytest.mark.parametrize("bad", ["C:/models", "C:", "..", "a/../..", "x:y"])
def test_safe_join_refuses_escapes(ws, bad):
from rigshare_core.workspace import safe_join
with pytest.raises(PermissionError):
safe_join(ws.root, bad)
def test_folder_ops_cannot_leave_the_workspace(ws, tmp_path):
outside = tmp_path / "models"
outside.mkdir()
(outside / "keep.txt").write_text("x")
# A colon segment (drive prefix) is refused outright; a plain absolute path is only ever
# read relative to the workflows root (a leading slash is dropped), so it can't reach `outside`.
for bad in ("C:" + str(outside)[2:], "shared/C:evil"):
with pytest.raises(PermissionError):
ws.delete_folder(bad)
with pytest.raises(PermissionError):
ws.rename_folder(bad, "renamed")
for plain in (str(outside), str(outside).replace("\\", "/")):
try:
ws.delete_folder(plain)
except (PermissionError, ValueError):
pass
assert (outside / "keep.txt").exists()
@pytest.mark.skipif(os.name != "nt", reason="trailing dots, spaces and short names only alias on Windows")
@pytest.mark.parametrize("raw", ["users./bob/a.json", "users /bob", "USERS~1/bob", "users/bob./a.json"])
def test_norm_refuses_windows_aliases(raw):
with pytest.raises(PermissionError):
norm(raw)
@pytest.mark.skipif(os.name != "nt", reason="case-insensitive filesystem")
def test_case_variants_cannot_bypass_private_folders(ws):
touch(ws, "users/alice/a.json")
assert norm("Users/Alice/a.json") == "users/alice/a.json"
assert norm("Workflows/USERS/ALICE") == "workflows/users/alice"
assert ws.role(norm("USERS/alice/a.json"), "user", "bob", EDITOR) is None
def test_name_rules():
from rigshare_core.workspace import FOLDER_NAME_RE
from rigshare_core.store import USERNAME_RE
for bad in ("con", "NUL", "com1.txt", "art.", "art ", "bad\n"):
assert not FOLDER_NAME_RE.match(bad), bad
for bad in ("..", "bob.", ".bob", "con", "aux", "ab\n", "a"):
assert not USERNAME_RE.match(bad), bad
for ok in ("Art (v2)", "a", "my-folder"):
assert FOLDER_NAME_RE.match(ok), ok
for ok in ("bob", "al.ice", "x_y-z", "ab"):
assert USERNAME_RE.match(ok), ok
+6
View File
@@ -861,6 +861,8 @@ export class RigSharePanel {
trusted: text((cfg.api_protection?.trusted_ips || []).join(", "), "127.0.0.1, 172.17.0.0/16"),
interval: h("input", { class: "rs-input rs-num", type: "number", min: 0, value: Math.round((cfg.snapshot_interval_sec ?? 300) / 60) }),
keep: h("input", { class: "rs-input rs-num", type: "number", min: 1, value: cfg.snapshot_keep ?? 30 }),
unloadMinutes: h("input", { class: "rs-input rs-num", type: "number", min: 0, value: cfg.auto_unload_minutes ?? 0 }),
unloadCache: box(cfg.auto_unload_cache),
};
const line = (input, label, hint) => h("label", { class: "rs-check rs-block" }, input,
h("span", {}, label, hint ? h("div", { class: "rs-muted rs-small" }, hint) : null));
@@ -880,6 +882,8 @@ export class RigSharePanel {
api_protection: { enabled: f.protect.checked, trusted_ips: f.trusted.value.split(/[\s,]+/).filter(Boolean) },
snapshot_interval_sec: Math.max(0, +f.interval.value || 0) * 60,
snapshot_keep: Math.max(1, +f.keep.value || 30),
auto_unload_minutes: Math.max(0, +f.unloadMinutes.value || 0),
auto_unload_cache: f.unloadCache.checked,
}), "Settings saved");
if (ok) { this.client.server.name = ok.server_name; this.renderHeader(); }
},
@@ -896,6 +900,8 @@ export class RigSharePanel {
h("label", { class: "rs-label" }, "Trusted IPs / networks (skip login and checks)"), f.trusted,
h("div", { class: "rs-row" }, h("label", { class: "rs-label rs-grow" }, "Auto snapshot every (min, 0 = off)"), f.interval),
h("div", { class: "rs-row" }, h("label", { class: "rs-label rs-grow" }, "Automatic snapshots kept"), f.keep),
h("div", { class: "rs-row" }, h("label", { class: "rs-label rs-grow" }, "Unload models after idle (min, 0 = off)"), f.unloadMinutes),
line(f.unloadCache, "Also clear the cache when unloading", "Frees RAM as well, but the next run computes every node again"),
h("button", { class: "rs-btn rs-btn-primary", type: "submit" }, h("i", { class: "pi pi-save" }), "Save settings"));
}
}