Xiaodoudou
3de3cc32e2
* 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
982 lines
45 KiB
Python
982 lines
45 KiB
Python
"""HTTP API, API-protection middleware and execution-event broadcasting."""
|
|
|
|
import asyncio
|
|
import base64
|
|
import html
|
|
import ipaddress
|
|
import json
|
|
import logging
|
|
import os
|
|
import posixpath
|
|
import time
|
|
from urllib.parse import unquote, urlsplit
|
|
from collections import defaultdict
|
|
|
|
from aiohttp import web
|
|
|
|
from .workspace import InvalidPath, Workspace, norm as _norm
|
|
|
|
log = logging.getLogger("ComfyUI-RigShare")
|
|
|
|
# Runs in the ComfyUI page before anything else: any request answered "login
|
|
# required" (session ended, cookie cleared) sends the browser to the login page.
|
|
LOGIN_WATCH = """<script id="rigshare-login-watch">(() => {
|
|
const fetch0 = window.fetch;
|
|
let leaving = false;
|
|
window.fetch = async (...args) => {
|
|
const res = await fetch0(...args);
|
|
if (res.status === 401 && !leaving && res.headers.get("X-RigShare-Login") === "required") {
|
|
leaving = true;
|
|
location.href = "/rigshare/login?next=" + encodeURIComponent(location.pathname + location.search);
|
|
}
|
|
return res;
|
|
};
|
|
})();</script>"""
|
|
|
|
# Other plugins' server routes RigShare guards: (method or None, path prefix,
|
|
# permission, why). Settings they keep on the server apply to everyone, like
|
|
# ComfyUI Manager, so they need the same "manager" permission; Pixaroma's workflow
|
|
# browser reads the workflows folder straight from disk, private folders included.
|
|
PLUGIN_RULES = [
|
|
("POST", "/rgthree/api/", "manager", "changing rgthree-comfy's settings needs the 'manager' permission"),
|
|
("POST", "/crystools/", "manager", "changing Crystools' settings needs the 'manager' permission"),
|
|
(None, "/pixaroma/api/workflows", "admin",
|
|
"Pixaroma's workflow browser shows every folder, private ones included; use RigShare's Files tab"),
|
|
]
|
|
|
|
# (method, path) prefixes that need a permission when api_protection is on.
|
|
PROTECTED = [
|
|
("POST", "/prompt", "queue"),
|
|
("POST", "/queue", "queue"),
|
|
("POST", "/interrupt", "queue"),
|
|
("POST", "/free", "admin"), # unloading models affects everyone's runs
|
|
("POST", "/upload/", "edit|queue"),
|
|
("POST", "/api/jobs", "queue"),
|
|
]
|
|
|
|
# ComfyUI Manager: everything needs the "manager" permission except a few
|
|
# read-only lookups the ComfyUI frontend makes for everyone (missing nodes...).
|
|
MANAGER_PREFIXES = ("/manager/", "/customnode/", "/snapshot/", "/externalmodel/", "/comfyui_manager/")
|
|
MANAGER_READONLY = {"/manager/version", "/manager/notice", "/manager/startup_alerts", "/manager/is_legacy_manager_ui",
|
|
"/customnode/getmappings", "/customnode/installed", "/customnode/import_fail_info",
|
|
"/customnode/get_node_types_in_workflows", "/customnode/alternatives"}
|
|
|
|
|
|
def is_api_template(template):
|
|
"""Templates that need paid API nodes (Comfy.org partner nodes)."""
|
|
return ("API" in (template.get("tags") or []) or template.get("openSource") is False
|
|
or str(template.get("name", "")).startswith("api_"))
|
|
|
|
|
|
def filter_templates(index):
|
|
out = []
|
|
for category in index:
|
|
if not isinstance(category, dict) or not isinstance(category.get("templates"), list):
|
|
out.append(category)
|
|
continue
|
|
kept = [t for t in category["templates"] if not is_api_template(t)]
|
|
if kept:
|
|
out.append({**category, "templates": kept})
|
|
return out
|
|
|
|
|
|
def needs_manager(method, path):
|
|
if path.startswith("/v2/"):
|
|
path = path[3:]
|
|
if not path.startswith(MANAGER_PREFIXES):
|
|
return False
|
|
return not (method in ("GET", "HEAD") and path in MANAGER_READONLY)
|
|
|
|
# Execution events ComfyUI normally sends only to the prompt's owner.
|
|
BROADCAST_EVENTS = {"execution_start", "execution_cached", "executing", "executed", "progress",
|
|
"progress_state", "execution_success", "execution_error", "execution_interrupted",
|
|
"progress_text"}
|
|
|
|
|
|
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.
|
|
PUBLIC_PATHS = {"/rigshare/login", "/rigshare/api/login", "/rigshare/api/logout", "/rigshare/api/setup",
|
|
"/rigshare/api/status", "/favicon.ico"}
|
|
|
|
|
|
def token_from(request):
|
|
"""Session token or API key from header, cookie or ?token= (for websockets)."""
|
|
auth = request.headers.get("Authorization", "")
|
|
if auth.startswith("Bearer "):
|
|
return auth[7:].strip()
|
|
return (request.headers.get("X-RigShare-Token") or request.cookies.get(COOKIE)
|
|
or request.query.get("token"))
|
|
|
|
|
|
def setup(server, store, hub):
|
|
routes = web.RouteTableDef()
|
|
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))
|
|
|
|
def perms_for(request):
|
|
user = current_user(request)
|
|
if user:
|
|
perms = dict(user["perms"])
|
|
if perms.get("admin"):
|
|
perms.update(edit=True, queue=True, manager=True)
|
|
return perms
|
|
guest = request.headers.get("X-RigShare-Guest")
|
|
if guest and store.config.get("allow_guests") and not store.config.get("require_login"):
|
|
return store.guest_perms(guest)
|
|
return {"edit": False, "queue": False, "manager": False, "admin": False}
|
|
|
|
def who(request):
|
|
"""(kind, key, perms) of the caller, for room access checks."""
|
|
user = current_user(request)
|
|
if user:
|
|
return "user", user["username"], perms_for(request)
|
|
return "guest", request.headers.get("X-RigShare-Guest"), perms_for(request)
|
|
|
|
def room_access(request, room):
|
|
return hub.access(room, *who(request))
|
|
|
|
def require_admin(request):
|
|
user = current_user(request)
|
|
if not user or not user["perms"].get("admin"):
|
|
raise web.HTTPForbidden(text='{"error": "Admin only"}', content_type="application/json")
|
|
return user
|
|
|
|
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,
|
|
samesite="Lax", secure=secure)
|
|
|
|
# ----- websocket + session -------------------------------------------
|
|
|
|
@routes.get("/rigshare/ws")
|
|
async def ws_handler(request):
|
|
hub.start(asyncio.get_running_loop())
|
|
if request.headers.get("Upgrade", "").lower() != "websocket":
|
|
log.warning(f"[RigShare] WebSocket request from {request.headers.get('X-Forwarded-For', request.remote)} "
|
|
"arrived without 'Upgrade: websocket'. If ComfyUI is behind a reverse proxy, enable "
|
|
"WebSocket support there (Nginx Proxy Manager: 'Websockets Support').")
|
|
return error("WebSocket upgrade headers missing: enable WebSocket support in your reverse proxy", 400)
|
|
return await hub.handle(request)
|
|
|
|
@routes.get("/rigshare/api/status")
|
|
async def status(request):
|
|
return web.json_response({"server_name": store.config.get("server_name"),
|
|
"require_login": store.config.get("require_login"),
|
|
"needs_setup": store.needs_setup,
|
|
"allow_guests": store.config.get("allow_guests"),
|
|
"online": len(hub.clients)})
|
|
|
|
@routes.get("/rigshare/login")
|
|
async def login_page(request):
|
|
with open(LOGIN_PAGE, "r", encoding="utf-8") as f:
|
|
page = f.read().replace("{{SERVER_NAME}}", html.escape(store.config.get("server_name") or "RigShare"))
|
|
page = page.replace("{{NEEDS_SETUP}}", "true" if store.needs_setup else "false")
|
|
try:
|
|
with open(ICON_SVG, "r", encoding="utf-8") as f:
|
|
icon = f.read()
|
|
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", **SECURITY_HEADERS})
|
|
|
|
@routes.post("/rigshare/api/login")
|
|
async def login(request):
|
|
ip, now = client_ip(request) or "?", time.time()
|
|
body = await request.json()
|
|
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)})
|
|
set_cookie(request, response, token)
|
|
return response
|
|
|
|
@routes.post("/rigshare/api/setup")
|
|
async def setup_admin(request):
|
|
if not store.needs_setup:
|
|
return error("Setup is already complete", 409)
|
|
body = await request.json()
|
|
try:
|
|
user = store.setup_admin(str(body.get("username", "")).strip(), str(body.get("password", "")),
|
|
body.get("display_name"))
|
|
except PermissionError as e:
|
|
return error(str(e), 409)
|
|
except ValueError as e:
|
|
return error(str(e))
|
|
log.warning(f"[RigShare] First-run setup: admin '{user['username']}' created from {request.remote}")
|
|
token = store.create_session(user["username"])
|
|
response = web.json_response({"token": token, "user": store.public_user(user)})
|
|
set_cookie(request, response, token)
|
|
return response
|
|
|
|
@routes.post("/rigshare/api/logout")
|
|
async def logout(request):
|
|
token = token_from(request)
|
|
if token:
|
|
store.revoke_session(token)
|
|
response = web.json_response({"ok": True})
|
|
response.del_cookie(COOKIE, path="/")
|
|
return response
|
|
|
|
@routes.get("/rigshare/api/me")
|
|
async def me(request):
|
|
user = current_user(request)
|
|
if not user:
|
|
return error("Not logged in", 401)
|
|
return web.json_response(store.public_user(user))
|
|
|
|
@routes.post("/rigshare/api/me")
|
|
async def update_me(request):
|
|
user = current_user(request)
|
|
if not user:
|
|
return error("Not logged in", 401)
|
|
body = await request.json()
|
|
try:
|
|
if body.get("new_password"):
|
|
if not store.authenticate(user["username"], body.get("current_password", "")):
|
|
return error("Current password is wrong", 403)
|
|
store.update_user(user["username"], password=body["new_password"])
|
|
# Changing password revokes sessions; keep the caller signed in.
|
|
token = store.create_session(user["username"])
|
|
else:
|
|
token = None
|
|
if "display_name" in body:
|
|
store.update_user(user["username"], display_name=str(body["display_name"]))
|
|
except ValueError as e:
|
|
return error(str(e))
|
|
await hub.refresh_identities(user["username"])
|
|
response = web.json_response({"user": store.public_user(store.users[user["username"]]), "token": token})
|
|
if token:
|
|
set_cookie(request, response, token)
|
|
return response
|
|
|
|
@routes.get("/rigshare/api/me/keys")
|
|
async def list_keys(request):
|
|
user = current_user(request)
|
|
if not user:
|
|
return error("Not logged in", 401)
|
|
return web.json_response(store.list_api_keys(user["username"]))
|
|
|
|
@routes.post("/rigshare/api/me/keys")
|
|
async def create_key(request):
|
|
user = current_user(request)
|
|
if not user:
|
|
return error("Not logged in", 401)
|
|
body = await request.json()
|
|
key, meta = store.create_api_key(user["username"], body.get("name"))
|
|
return web.json_response({"key": key, **meta})
|
|
|
|
@routes.delete("/rigshare/api/me/keys/{key_id}")
|
|
async def revoke_key(request):
|
|
user = current_user(request)
|
|
if not user:
|
|
return error("Not logged in", 401)
|
|
if not store.revoke_api_key(user["username"], request.match_info["key_id"]):
|
|
return error("No such key", 404)
|
|
return web.json_response({"ok": True})
|
|
|
|
# ----- admin: accounts, guests, config -------------------------------
|
|
|
|
@routes.get("/rigshare/api/users")
|
|
async def list_users(request):
|
|
require_admin(request)
|
|
return web.json_response([store.public_user(u) for u in store.users.values()])
|
|
|
|
@routes.post("/rigshare/api/users")
|
|
async def create_user(request):
|
|
require_admin(request)
|
|
body = await request.json()
|
|
try:
|
|
user = store.create_user(str(body.get("username", "")).strip(), str(body.get("password", "")),
|
|
body.get("display_name"), body.get("perms") or {})
|
|
except ValueError as e:
|
|
return error(str(e))
|
|
return web.json_response(store.public_user(user))
|
|
|
|
@routes.patch("/rigshare/api/users/{username}")
|
|
async def update_user(request):
|
|
admin = require_admin(request)
|
|
username = request.match_info["username"]
|
|
body = await request.json()
|
|
try:
|
|
user = store.update_user(username, body.get("display_name"), body.get("perms"), body.get("password") or None)
|
|
if body.get("username"):
|
|
user = store.rename_user(username, str(body["username"]))
|
|
except KeyError:
|
|
return error("No such user", 404)
|
|
except ValueError as e:
|
|
return error(str(e))
|
|
if user["username"] != username:
|
|
if hub.workspace:
|
|
hub.workspace.rename_user(username, user["username"])
|
|
await hub.path_moved(f"workflows/users/{username}", f"workflows/users/{user['username']}",
|
|
renamed={username: user["username"]})
|
|
log.info(f"[RigShare] {admin['username']} renamed @{username} to @{user['username']}")
|
|
await hub.refresh_identities()
|
|
else:
|
|
await hub.refresh_identities(username)
|
|
return web.json_response(store.public_user(user))
|
|
|
|
@routes.delete("/rigshare/api/users/{username}")
|
|
async def delete_user(request):
|
|
require_admin(request)
|
|
username = request.match_info["username"]
|
|
try:
|
|
store.delete_user(username)
|
|
if hub.workspace:
|
|
hub.workspace.forget_member(username)
|
|
except KeyError:
|
|
return error("No such user", 404)
|
|
except ValueError as e:
|
|
return error(str(e))
|
|
await hub.refresh_identities(username)
|
|
return web.json_response({"ok": True})
|
|
|
|
@routes.post("/rigshare/api/guests/{guest_id}")
|
|
async def guest_perms(request):
|
|
require_admin(request)
|
|
store.set_guest_perms(request.match_info["guest_id"], await request.json())
|
|
await hub.refresh_identities(request.match_info["guest_id"])
|
|
return web.json_response({"ok": True})
|
|
|
|
@routes.get("/rigshare/api/config")
|
|
async def get_config(request):
|
|
require_admin(request)
|
|
return web.json_response(store.config)
|
|
|
|
@routes.patch("/rigshare/api/config")
|
|
async def patch_config(request):
|
|
require_admin(request)
|
|
config = store.update_config(await request.json())
|
|
await hub.refresh_identities()
|
|
await hub.broadcast_server()
|
|
return web.json_response(config)
|
|
|
|
# ----- rooms + snapshots ---------------------------------------------
|
|
|
|
def display(request):
|
|
user = current_user(request)
|
|
if user:
|
|
return user.get("display_name") or user["username"]
|
|
return "guest"
|
|
|
|
@routes.get("/rigshare/api/chat")
|
|
async def chat_history(request):
|
|
"""Older chat messages: ``before`` is the seq of the oldest one already shown."""
|
|
try:
|
|
before = request.query.get("before")
|
|
before = int(before) if before not in (None, "") else None
|
|
limit = int(request.query.get("limit", 50))
|
|
except ValueError:
|
|
return error("before and limit must be numbers")
|
|
messages, more = hub.chat.page(before, limit)
|
|
return web.json_response({"messages": messages, "more": more})
|
|
|
|
@routes.delete("/rigshare/api/chat")
|
|
async def clear_chat(request):
|
|
admin = require_admin(request)
|
|
await hub.clear_chat(admin.get("display_name") or admin["username"])
|
|
log.info(f"[RigShare] {admin['username']} cleared the chat history")
|
|
return web.json_response({"ok": True})
|
|
|
|
@routes.get("/rigshare/api/performance/history")
|
|
async def performance_history(request):
|
|
"""Finished runs, newest first (?limit=, at most 500; ?workflow= to filter)."""
|
|
try:
|
|
limit = max(1, min(500, int(request.query.get("limit", 100))))
|
|
except ValueError:
|
|
return error("limit must be a number")
|
|
workflow = request.query.get("workflow")
|
|
runs = [r for r in reversed(store.runs()) if not workflow or r.get("workflow") == workflow]
|
|
return web.json_response({"runs": runs[:limit], "total": len(runs)})
|
|
|
|
@routes.delete("/rigshare/api/performance")
|
|
async def performance_clear(request):
|
|
"""Forget one workflow's timings, so its ETA starts over."""
|
|
if not perms_for(request).get("edit"):
|
|
return error("You need the Edit permission to clear timings", 403)
|
|
workflow = request.query.get("workflow", "").strip()
|
|
if not workflow:
|
|
return error("Which workflow?")
|
|
gone = store.clear_runs(workflow)
|
|
log.info(f"[RigShare] {display(request)} cleared the timings of {workflow} ({gone} runs)")
|
|
return web.json_response({"ok": True, "cleared": gone})
|
|
|
|
@routes.get("/rigshare/api/rooms")
|
|
async def list_rooms(request):
|
|
return web.json_response([r for r in hub.room_list() if room_access(request, hub.rooms[r["key"]])])
|
|
|
|
# ----- workspace folders ------------------------------------------
|
|
|
|
def ws():
|
|
if not hub.workspace:
|
|
raise web.HTTPServiceUnavailable(text='{"error": "Workspace unavailable"}', content_type="application/json")
|
|
return hub.workspace
|
|
|
|
@routes.get("/rigshare/api/workspace")
|
|
async def workspace_tree(request):
|
|
if not current_user(request):
|
|
return error("Not logged in", 401)
|
|
return web.json_response(ws().tree(*who(request)))
|
|
|
|
def ws_call(fn):
|
|
try:
|
|
return None, fn()
|
|
except PermissionError as e:
|
|
return error(str(e), 403), None
|
|
except FileNotFoundError as e:
|
|
return error(str(e), 404), None
|
|
except ValueError as e:
|
|
return error(str(e)), None
|
|
|
|
@routes.get("/rigshare/api/workspace/list")
|
|
async def ws_list(request):
|
|
if not current_user(request):
|
|
return error("Not logged in", 401)
|
|
err, result = ws_call(lambda: ws().list_dir(request.query.get("path", ""), *who(request)))
|
|
return err or web.json_response(result)
|
|
|
|
@routes.post("/rigshare/api/workspace/mkdir")
|
|
async def ws_mkdir(request):
|
|
body = await request.json()
|
|
err, path = ws_call(lambda: ws().make_dir(body.get("parent", ""), body.get("name"), *who(request)))
|
|
if err:
|
|
return err
|
|
await hub.files_changed("workflows/" + posixpath.dirname(path))
|
|
return web.json_response({"path": path})
|
|
|
|
@routes.post("/rigshare/api/workspace/rename-folder")
|
|
async def ws_rename_dir(request):
|
|
body = await request.json()
|
|
old = norm(body.get("path"))
|
|
seen_by = hub.who_sees("workflows/" + old)
|
|
err, new = ws_call(lambda: ws().rename_dir(old, body.get("name"), *who(request)))
|
|
if err:
|
|
return err
|
|
await hub.path_moved("workflows/" + old, "workflows/" + new, seen_by=seen_by)
|
|
return web.json_response({"path": new})
|
|
|
|
@routes.post("/rigshare/api/workspace/move-folder")
|
|
async def ws_move_dir(request):
|
|
body = await request.json()
|
|
old = norm(body.get("path"))
|
|
seen_by = hub.who_sees("workflows/" + old)
|
|
err, new = ws_call(lambda: ws().move_dir(old, body.get("dest", ""), *who(request)))
|
|
if err:
|
|
return err
|
|
await hub.path_moved("workflows/" + old, "workflows/" + new, seen_by=seen_by)
|
|
return web.json_response({"path": new})
|
|
|
|
@routes.post("/rigshare/api/workspace/delete-folder")
|
|
async def ws_delete_dir(request):
|
|
body = await request.json()
|
|
path = norm(body.get("path", ""))
|
|
err, _ = ws_call(lambda: ws().delete_dir(path, *who(request)))
|
|
if err:
|
|
return err
|
|
await hub.drop_rooms("file:workflows/" + path, display(request))
|
|
await hub.files_changed("workflows/" + posixpath.dirname(path))
|
|
return web.json_response({"ok": True})
|
|
|
|
@routes.get("/rigshare/api/workspace/folder")
|
|
async def folder_acl(request):
|
|
folder = norm(request.query.get("path"))
|
|
if Workspace.area(folder)[0] != "shared" or folder.count("/") != 1:
|
|
return error("Not a shared folder")
|
|
if not ws().can_manage(folder, *who(request)):
|
|
return error("Only the folder's owner or an admin can see its access list", 403)
|
|
meta = ws().folders.get(folder) or {}
|
|
return web.json_response({"path": folder, "owner": meta.get("owner"), "restricted": bool(meta.get("acl")),
|
|
"members": (meta.get("acl") or {}).get("members") or {}})
|
|
|
|
@routes.post("/rigshare/api/workspace/folders")
|
|
async def create_folder(request):
|
|
kind, key, perms = who(request)
|
|
if kind != "user" or not (perms.get("edit") or perms.get("admin")):
|
|
return error("You need the Edit permission to create shared folders", 403)
|
|
body = await request.json()
|
|
try:
|
|
folder = ws().create_folder(body.get("name"), key)
|
|
except ValueError as e:
|
|
return error(str(e))
|
|
log.info(f"[RigShare] {key} created folder {folder}")
|
|
await hub.files_changed("workflows/shared")
|
|
return web.json_response({"path": folder})
|
|
|
|
@routes.patch("/rigshare/api/workspace/folders")
|
|
async def update_folder(request):
|
|
body = await request.json()
|
|
folder = norm(body.get("path"))
|
|
if Workspace.area(folder)[0] != "shared" or folder.count("/") != 1:
|
|
return error("Not a shared folder")
|
|
if not ws().can_manage(folder, *who(request)):
|
|
return error("Only the folder's owner or an admin can change it", 403)
|
|
try:
|
|
if "restricted" in body:
|
|
acl = None
|
|
if body.get("restricted"):
|
|
owner = (ws().folders.get(folder) or {}).get("owner")
|
|
members = {u: r for u, r in (body.get("members") or {}).items()
|
|
if u in store.users and r in ("view", "edit") and u != owner}
|
|
acl = {"members": members}
|
|
kind, key, _ = who(request)
|
|
ws().set_acl(folder, acl, fallback_owner=key if kind == "user" else None)
|
|
if body.get("name"):
|
|
seen_by = hub.who_sees(f"workflows/{folder}")
|
|
new_folder = ws().rename_folder(folder, body["name"])
|
|
await hub.path_moved(f"workflows/{folder}", f"workflows/{new_folder}", seen_by=seen_by)
|
|
folder = new_folder
|
|
except ValueError as e:
|
|
return error(str(e))
|
|
await hub.refresh_rooms_access()
|
|
await hub.files_changed("workflows/shared") # the folder's access changed who sees it
|
|
return web.json_response({"ok": True, "path": folder})
|
|
|
|
@routes.delete("/rigshare/api/workspace/folders")
|
|
async def delete_folder(request):
|
|
folder = norm(request.query.get("path"))
|
|
if Workspace.area(folder)[0] != "shared" or folder.count("/") != 1:
|
|
return error("Not a shared folder")
|
|
if not ws().can_manage(folder, *who(request)):
|
|
return error("Only the folder's owner or an admin can delete it", 403)
|
|
try:
|
|
ws().delete_folder(folder)
|
|
except ValueError as e:
|
|
return error(str(e))
|
|
return web.json_response({"ok": True})
|
|
|
|
@routes.get("/rigshare/api/people")
|
|
async def people(request):
|
|
"""Accounts that can be added to a workflow's access list."""
|
|
if not current_user(request):
|
|
return error("Not logged in", 401)
|
|
return web.json_response([{"username": u["username"], "display_name": u.get("display_name") or u["username"],
|
|
"admin": bool(u["perms"].get("admin")), "edit": bool(u["perms"].get("edit"))}
|
|
for u in store.users.values()])
|
|
|
|
@routes.get("/rigshare/api/room")
|
|
async def get_room(request):
|
|
room = hub.rooms.get(request.query.get("key", ""))
|
|
if not room or not room_access(request, room):
|
|
return error("No such room", 404)
|
|
return web.json_response({"key": room.key, "name": room.name, "kind": room.kind,
|
|
"version": room.version, "doc": room.doc})
|
|
|
|
@routes.get("/rigshare/api/snapshots")
|
|
async def list_snapshots(request):
|
|
key = request.query.get("room")
|
|
if not key:
|
|
require_admin(request)
|
|
return web.json_response(store.list_snapshots())
|
|
if hub.access_for_key(key, *who(request)) is None:
|
|
return error("No such room", 404)
|
|
return web.json_response(store.list_snapshots(key))
|
|
|
|
@routes.post("/rigshare/api/snapshots")
|
|
async def create_snapshot(request):
|
|
if not perms_for(request).get("edit"):
|
|
return error("You do not have permission to edit", 403)
|
|
body = await request.json()
|
|
room = hub.rooms.get(body.get("room", ""))
|
|
if not room:
|
|
return error("This tab is not shared")
|
|
if room_access(request, room) != "edit":
|
|
return error("You can view this workflow but not edit it", 403)
|
|
snap = hub.snapshot(room, display(request), body.get("label") or "Manual")
|
|
if not snap:
|
|
return error("Nothing to snapshot yet")
|
|
return web.json_response(snap)
|
|
|
|
@routes.get("/rigshare/api/snapshots/{snap_id}")
|
|
async def get_snapshot(request):
|
|
snap = store.get_snapshot(request.match_info["snap_id"])
|
|
if not snap:
|
|
return error("Not found", 404)
|
|
# 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)
|
|
|
|
@routes.post("/rigshare/api/snapshots/{snap_id}/restore")
|
|
async def restore_snapshot(request):
|
|
if not perms_for(request).get("edit"):
|
|
return error("You do not have permission to edit", 403)
|
|
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 not room:
|
|
return error("That workflow is no longer shared; open the snapshot as a copy instead", 404)
|
|
if room_access(request, room) != "edit":
|
|
return error("You can view this workflow but not edit it", 403)
|
|
by = display(request)
|
|
hub.snapshot(room, by, "Before restore")
|
|
await hub.replace_doc(room, snap["doc"], by, f"{by} restored a snapshot of {room.name}")
|
|
return web.json_response({"ok": True})
|
|
|
|
server.app.add_routes(routes)
|
|
|
|
# ----- API protection middleware -------------------------------------
|
|
|
|
def trusted(ip):
|
|
try:
|
|
addr = ipaddress.ip_address(ip)
|
|
except ValueError:
|
|
return False
|
|
for entry in store.config["api_protection"].get("trusted_ips", []):
|
|
try:
|
|
if addr in ipaddress.ip_network(entry, strict=False):
|
|
return True
|
|
except ValueError:
|
|
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]
|
|
segs = raw.split("/")
|
|
if "userdata" not in segs:
|
|
return None
|
|
i = segs.index("userdata")
|
|
if segs[i - 1] not in ("", "api", "v2"):
|
|
return None
|
|
if len(segs) == i + 1: # listing: /userdata?dir= or /v2/userdata?path=
|
|
if request.method != "GET":
|
|
return None
|
|
directory = norm(request.query.get("dir") or request.query.get("path") or "")
|
|
return ("list", directory, None) if Workspace.workflows_rel(directory) is not None else None
|
|
# ComfyUI encodes "/" inside names (one segment each), but accept plain
|
|
# slashes too so no spelling of a path can slip past the checks.
|
|
rest = segs[i + 1:]
|
|
if "move" in rest[1:]:
|
|
m = rest.index("move", 1)
|
|
src, dest = norm("/".join(rest[:m])), norm("/".join(rest[m + 1:]))
|
|
else:
|
|
src, dest = norm("/".join(rest)), None
|
|
if Workspace.workflows_rel(src) is None and (dest is None or Workspace.workflows_rel(dest) is None):
|
|
return None
|
|
if dest is not None:
|
|
return ("move", src, dest)
|
|
action = {"GET": "get", "HEAD": "get", "POST": "save", "DELETE": "delete"}.get(request.method)
|
|
return (action, src, None) if action else None
|
|
|
|
def watch_login(response):
|
|
try:
|
|
file_path = getattr(response, "_path", None)
|
|
if file_path:
|
|
with open(file_path, "r", encoding="utf-8") as f:
|
|
page = f.read()
|
|
elif isinstance(response, web.Response) and response.body and "html" in (response.content_type or ""):
|
|
page = response.body.decode("utf-8")
|
|
else:
|
|
return response
|
|
if "<head>" not in page or "rigshare-login-watch" in page:
|
|
return response
|
|
page = page.replace("<head>", "<head>" + LOGIN_WATCH, 1)
|
|
return web.Response(text=page, content_type="text/html",
|
|
headers={"Cache-Control": "no-store, must-revalidate"})
|
|
except Exception as e:
|
|
log.debug(f"[RigShare] could not add the login watch: {e}")
|
|
return response
|
|
|
|
def api_path(request):
|
|
return request.path[4:] if request.path.startswith("/api/") else request.path
|
|
|
|
async def check_queue_ownership(request):
|
|
"""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(request) or perms_for(request).get("admin"):
|
|
return None
|
|
user = current_user(request)
|
|
me = user["username"] if user else None
|
|
try:
|
|
body = await request.json() if request.body_exists else {}
|
|
except ValueError:
|
|
body = {}
|
|
body = body if isinstance(body, dict) else {}
|
|
if api_path(request) == "/queue":
|
|
if body.get("clear"):
|
|
return "only admins can clear the whole queue"
|
|
ids = body.get("delete") or []
|
|
else:
|
|
ids = [body["prompt_id"]] if body.get("prompt_id") else hub.running_prompt_ids()
|
|
for prompt_id in ids:
|
|
owner = hub.owner_of(prompt_id)
|
|
if owner and owner != me:
|
|
return "you can only cancel your own prompts"
|
|
return None
|
|
|
|
def snapshot_saved_file(request, path, body):
|
|
if not path.endswith(".json") or Workspace.workflows_rel(path) is None:
|
|
return
|
|
try:
|
|
doc = json.loads(body)
|
|
except ValueError:
|
|
return
|
|
if not isinstance(doc, dict) or "nodes" not in doc:
|
|
return
|
|
name = path.rsplit("/", 1)[-1].removesuffix(".json").removesuffix(".app")
|
|
# The panel shows the time in each viewer's own time zone; the label is for API clients.
|
|
label = "Saved at " + time.strftime("%Y-%m-%d %H:%M UTC", time.gmtime())
|
|
store.add_snapshot(doc, display(request), "file:" + path, name, label, kind="save")
|
|
|
|
def check_userdata(request, action, src, dest):
|
|
"""Folder rules (private spaces, restricted shared folders)."""
|
|
kind, key, perms = who(request)
|
|
w = hub.workspace
|
|
can_edit = bool(perms.get("edit") or perms.get("admin"))
|
|
|
|
def role(path):
|
|
rel = Workspace.workflows_rel(path)
|
|
if rel is None:
|
|
return "edit"
|
|
return w.role(rel, kind, key, perms) if w else ("edit" if can_edit else "view")
|
|
|
|
def deletable(path):
|
|
rel = Workspace.workflows_rel(path)
|
|
if rel is None:
|
|
return True
|
|
if w:
|
|
return w.can_delete(rel, kind, key, perms) and role(path) is not None
|
|
return bool(perms.get("admin"))
|
|
|
|
if action == "list":
|
|
return None
|
|
if action == "get":
|
|
return None if role(src) else "you do not have access to this workflow"
|
|
if action == "save":
|
|
return None if role(src) == "edit" else "you cannot save in this folder"
|
|
if action == "delete":
|
|
return None if deletable(src) else "only the folder's owner or an admin can delete this"
|
|
if action == "move":
|
|
if Workspace.workflows_rel(src) is None or Workspace.workflows_rel(dest) is None:
|
|
return "workflows cannot be moved in or out of the workflows folder"
|
|
if role(src) is None or role(dest) != "edit":
|
|
return "you cannot move this workflow there"
|
|
if w:
|
|
ok = w.can_move(Workspace.workflows_rel(src), Workspace.workflows_rel(dest), kind, key, perms)
|
|
else:
|
|
ok = role(src) == "edit" and (perms.get("admin") or src.rsplit("/", 1)[0] == dest.rsplit("/", 1)[0])
|
|
return None if ok else "you cannot move this workflow there"
|
|
return None
|
|
|
|
def filter_listing(request, response, directory):
|
|
try:
|
|
items = json.loads(response.body)
|
|
except Exception:
|
|
return response
|
|
if not isinstance(items, list):
|
|
return response
|
|
w = hub.workspace
|
|
kind, key, perms = who(request)
|
|
base = Workspace.workflows_rel(directory) or ""
|
|
|
|
def rel_of(item):
|
|
p = item.get("path") if isinstance(item, dict) else (item[0] if isinstance(item, list) else item)
|
|
p = norm(p)
|
|
if p.startswith("workflows/"):
|
|
return p[len("workflows/"):]
|
|
return f"{base}/{p}" if base else p
|
|
|
|
kept = [it for it in items if w.visible(rel_of(it), kind, key, perms)]
|
|
return web.json_response(kept, headers={"Cache-Control": "no-store"})
|
|
|
|
@web.middleware
|
|
async def protect(request, handler):
|
|
path = request.path
|
|
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", ""))
|
|
if wants_page:
|
|
raise web.HTTPFound(f"/rigshare/login?next={path}")
|
|
# The header lets the ComfyUI page (see LOGIN_WATCH) go back to the login screen.
|
|
return web.json_response({"error": "RigShare: login required. API clients: send "
|
|
"'Authorization: Bearer <api key>'."}, status=401,
|
|
headers={"X-RigShare-Login": "required"})
|
|
userdata = parse_userdata(request)
|
|
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(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(request):
|
|
break
|
|
perms = perms_for(request)
|
|
if not any(perms.get(p) for p in need.split("|")):
|
|
return web.json_response(
|
|
{"error": {"type": "rigshare_forbidden",
|
|
"message": f"RigShare: you need '{need.replace('|', ' or ')}' permission",
|
|
"details": "", "extra_info": {}}, "node_errors": {}},
|
|
status=403)
|
|
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(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)
|
|
if refusal:
|
|
return web.json_response({"error": f"RigShare: {refusal}"}, status=403)
|
|
if userdata and userdata[0] == "save":
|
|
await request.read() # keep the body for the snapshot below (aiohttp caches it)
|
|
response = await handler(request)
|
|
# Remember who queued each prompt, for the Queue tab.
|
|
if request.method == "POST" and api_path(request) == "/prompt" and getattr(response, "status", 0) == 200:
|
|
try:
|
|
prompt_id = json.loads(response.body).get("prompt_id")
|
|
except (ValueError, TypeError, AttributeError):
|
|
prompt_id = None
|
|
hub.record_prompt(prompt_id, current_user(request), unquote(request.headers.get("X-RigShare-Workflow", "")))
|
|
# Every save of a workflow keeps a snapshot, named after when it was saved.
|
|
if userdata and userdata[0] == "save" and getattr(response, "status", 0) == 200:
|
|
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(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:
|
|
await hub.path_moved(userdata[1], userdata[2])
|
|
if userdata and userdata[0] == "delete" and getattr(response, "status", 0) in (200, 204):
|
|
await hub.drop_rooms("file:" + userdata[1], display(request))
|
|
await hub.files_changed(userdata[1])
|
|
# Hide API templates: rewrite the template index ComfyUI serves.
|
|
if (store.config.get("hide_api_templates") and request.method == "GET"
|
|
and path.startswith("/templates/index") and path.endswith(".json")):
|
|
try:
|
|
file_path = getattr(response, "_path", None)
|
|
if file_path:
|
|
with open(file_path, "r", encoding="utf-8") as f:
|
|
index = json.load(f)
|
|
else:
|
|
index = json.loads(response.body)
|
|
if isinstance(index, list):
|
|
response = web.json_response(filter_templates(index), headers={"Cache-Control": "no-store"})
|
|
except Exception as e:
|
|
log.debug(f"[RigShare] could not filter templates: {e}")
|
|
# The ComfyUI page: never cached (an ended session must reach the redirect
|
|
# above), and watched, so a session that ends while it is open goes back
|
|
# to the login screen instead of showing API errors.
|
|
if request.method == "GET" and request.path in ("/", "/index.html") and store.config.get("require_login"):
|
|
response = watch_login(response)
|
|
# Sliding session: renew the cookie JWT when it is getting old.
|
|
cookie = request.cookies.get(COOKIE)
|
|
if cookie and isinstance(response, web.Response) and not response.prepared:
|
|
fresh = store.refresh_token(cookie)
|
|
if fresh:
|
|
set_cookie(request, response, fresh)
|
|
return response
|
|
|
|
try:
|
|
server.app.middlewares.append(protect)
|
|
except RuntimeError as e:
|
|
log.error(f"[ComfyUI-RigShare] Could not install API protection middleware: {e}")
|
|
|
|
# ----- relay execution events to everyone ----------------------------
|
|
|
|
original_send_sync = server.send_sync
|
|
|
|
def send_sync(event, data, sid=None):
|
|
if event in ("execution_start", "execution_success", "execution_error", "execution_interrupted"):
|
|
try:
|
|
hub.run_event(event, data) # Performance: time the run
|
|
except Exception as e:
|
|
log.debug(f"[RigShare] could not time the run: {e}")
|
|
if sid is not None and store.config.get("broadcast_execution", True):
|
|
if event in BROADCAST_EVENTS or isinstance(event, int):
|
|
sid = None
|
|
return original_send_sync(event, data, sid)
|
|
|
|
server.send_sync = send_sync
|