595 lines
24 KiB
Python
595 lines
24 KiB
Python
import os
|
|
import sys
|
|
import heapq
|
|
import copy
|
|
import contextvars
|
|
import mimetypes
|
|
from aiohttp import web
|
|
from typing import Optional
|
|
|
|
import folder_paths
|
|
from server import PromptServer
|
|
from execution import PromptQueue, MAXIMUM_HISTORY_SIZE
|
|
|
|
from .users_db import UsersDB
|
|
|
|
|
|
class AccessControl:
|
|
def __init__(self, users_db: UsersDB, server: PromptServer):
|
|
self.users_db = users_db
|
|
self.server = server
|
|
|
|
self._current_user = contextvars.ContextVar("user_id", default=None)
|
|
self.__current_user_id = None
|
|
|
|
self.__get_output_directory = folder_paths.get_output_directory
|
|
self.__get_temp_directory = folder_paths.get_temp_directory
|
|
self.__get_input_directory = folder_paths.get_input_directory
|
|
self.__get_save_image_path = folder_paths.get_save_image_path
|
|
|
|
self.__prompt_queue = self.server.prompt_queue
|
|
self.__prompt_queue_put = self.__prompt_queue.put
|
|
self.__prompt_queue_get_flags = self.__prompt_queue.get_flags
|
|
self.__user_manager_get_request_user_id = getattr(
|
|
self.server.user_manager, "get_request_user_id", None
|
|
)
|
|
|
|
@property
|
|
def folder_paths(self) -> tuple:
|
|
return (
|
|
self.__get_output_directory(),
|
|
self.__get_temp_directory(),
|
|
self.__get_input_directory(),
|
|
)
|
|
|
|
def set_current_user_id(self, user_id: str, set_fallback: bool = False) -> None:
|
|
"""Set the current authenticated user."""
|
|
self._current_user.set(user_id)
|
|
|
|
if set_fallback:
|
|
self.__current_user_id = user_id
|
|
|
|
def get_current_user_id(self) -> str:
|
|
"""Retrieve the active user, falling back to the latest prompt submitter."""
|
|
if self._current_user.get():
|
|
return self._current_user.get()
|
|
|
|
return self.__current_user_id
|
|
|
|
def get_context_user_id(self) -> str:
|
|
"""Retrieve only the current request/execution context user."""
|
|
return self._current_user.get()
|
|
|
|
def is_admin_user(self, user_id: str) -> bool:
|
|
if not user_id:
|
|
return False
|
|
_, user = self.users_db.get_user(user_id=user_id)
|
|
return bool(user and user.get("admin"))
|
|
|
|
def _can_access_owner(self, owner_id: str, current_user_id: str = None) -> bool:
|
|
current_user_id = current_user_id or self.get_current_user_id()
|
|
if self.is_admin_user(current_user_id):
|
|
return True
|
|
return bool(owner_id and owner_id == current_user_id)
|
|
|
|
def _with_user_prefix(self, filename_prefix: str) -> str:
|
|
user_id = self.get_current_user_id()
|
|
if not user_id or not filename_prefix:
|
|
return filename_prefix
|
|
|
|
normalized = os.path.normpath(str(filename_prefix)).replace("\\", "/")
|
|
if normalized == user_id or normalized.startswith(f"{user_id}/"):
|
|
return filename_prefix
|
|
|
|
return f"{user_id}/{filename_prefix}"
|
|
|
|
def get_user_output_directory(self) -> str:
|
|
"""Get the user-specific output directory."""
|
|
return os.path.join(
|
|
self.__get_output_directory(),
|
|
self.get_current_user_id() or "public",
|
|
)
|
|
|
|
def get_user_temp_directory(self) -> str:
|
|
"""Get the user-specific temp directory."""
|
|
return os.path.join(
|
|
self.__get_temp_directory(),
|
|
self.get_current_user_id() or "public",
|
|
)
|
|
|
|
def get_user_input_directory(self) -> str:
|
|
"""Get the user-specific input directory."""
|
|
input_directory = os.path.join(
|
|
self.__get_input_directory(),
|
|
self.get_current_user_id() or "public",
|
|
)
|
|
|
|
os.makedirs(input_directory, exist_ok=True)
|
|
|
|
return input_directory
|
|
|
|
def add_user_specific_folder_paths(self, json_data) -> None:
|
|
"""Add user-specific output prefixes to prompt JSON data."""
|
|
if isinstance(json_data, dict):
|
|
for key, value in json_data.items():
|
|
if key == "filename_prefix" and isinstance(value, str):
|
|
json_data[key] = self._with_user_prefix(value)
|
|
else:
|
|
self.add_user_specific_folder_paths(value)
|
|
elif isinstance(json_data, list):
|
|
for item in json_data:
|
|
self.add_user_specific_folder_paths(item)
|
|
|
|
return json_data
|
|
|
|
def get_user_save_image_path(
|
|
self,
|
|
filename_prefix: str,
|
|
output_dir: str,
|
|
image_width=0,
|
|
image_height=0,
|
|
) -> tuple[str, str, int, str, str]:
|
|
"""Ensure standard ComfyUI save helpers write under the current user."""
|
|
return self.__get_save_image_path(
|
|
self._with_user_prefix(filename_prefix),
|
|
output_dir,
|
|
image_width,
|
|
image_height,
|
|
)
|
|
|
|
def patch_folder_paths(self) -> None:
|
|
"""Patch folder paths and save helpers with user-specific behavior."""
|
|
folder_paths.get_temp_directory = self.get_user_temp_directory
|
|
folder_paths.get_input_directory = self.get_user_input_directory
|
|
folder_paths.get_save_image_path = self.get_user_save_image_path
|
|
|
|
self.server.add_on_prompt_handler(self.add_user_specific_folder_paths)
|
|
|
|
def sync_comfy_users(self) -> None:
|
|
"""Mirror account-manager users into ComfyUI's in-memory user map."""
|
|
user_manager = getattr(self.server, "user_manager", None)
|
|
if not user_manager or not hasattr(user_manager, "users"):
|
|
return
|
|
|
|
for user_id, user in self.users_db.load_users().items():
|
|
username = user.get("username") or user_id
|
|
user_manager.users[user_id] = username
|
|
|
|
def patch_comfy_user_manager(self) -> None:
|
|
"""Let ComfyUI APIs resolve authenticated account-manager users."""
|
|
user_manager = getattr(self.server, "user_manager", None)
|
|
if not user_manager or self.__user_manager_get_request_user_id is None:
|
|
return
|
|
|
|
def get_request_user_id(request):
|
|
user_id = request.get("user_id")
|
|
if user_id and user_id in self.users_db.load_users():
|
|
return user_id
|
|
|
|
header_user_id = request.headers.get("comfy-user")
|
|
if header_user_id and header_user_id in self.users_db.load_users():
|
|
return header_user_id
|
|
|
|
return self.__user_manager_get_request_user_id(request)
|
|
|
|
user_manager.get_request_user_id = get_request_user_id
|
|
self.sync_comfy_users()
|
|
|
|
def create_view_access_control_middleware(self) -> web.middleware:
|
|
"""Guard /view and /api/view output access by top-level user folder."""
|
|
|
|
@web.middleware
|
|
async def view_access_control_middleware(
|
|
request: web.Request, handler
|
|
) -> web.Response:
|
|
if request.path not in ("/view", "/api/view"):
|
|
return await handler(request)
|
|
|
|
view_type = request.rel_url.query.get("type", "output")
|
|
if view_type != "output":
|
|
return await handler(request)
|
|
|
|
user_id = request.get("user_id")
|
|
if self.is_admin_user(user_id):
|
|
return await handler(request)
|
|
|
|
subfolder = request.rel_url.query.get("subfolder", "")
|
|
owner_id = os.path.normpath(subfolder).replace("\\", "/").split("/", 1)[0]
|
|
|
|
if owner_id and owner_id == user_id:
|
|
return await handler(request)
|
|
|
|
return web.HTTPForbidden(reason="You do not have access to this output.")
|
|
|
|
return view_access_control_middleware
|
|
|
|
def create_folder_access_control_middleware(
|
|
self, folder_paths: tuple = ()
|
|
) -> web.middleware:
|
|
"""Create middleware for direct folder access control."""
|
|
|
|
folder_paths = folder_paths or self.folder_paths
|
|
|
|
@web.middleware
|
|
async def folder_access_control_middleware(
|
|
request: web.Request, handler
|
|
) -> web.Response:
|
|
if not request.path.startswith(folder_paths):
|
|
return await handler(request)
|
|
|
|
current_user_id = request.get("user_id")
|
|
|
|
try:
|
|
path_parts = request.path.strip("/").split("/")
|
|
folder_user_id = path_parts[1]
|
|
except Exception:
|
|
return web.HTTPNotFound(reason="Folder not found.")
|
|
|
|
if not self._can_access_owner(folder_user_id, current_user_id):
|
|
return web.HTTPForbidden(
|
|
reason="You do not have access to this folder."
|
|
)
|
|
|
|
return await handler(request)
|
|
|
|
return folder_access_control_middleware
|
|
|
|
def _wrap_queue_item(self, item):
|
|
if isinstance(item, tuple) and len(item) >= 7:
|
|
return item
|
|
return item + (self.get_current_user_id(),)
|
|
|
|
@staticmethod
|
|
def _unwrap_queue_item(item):
|
|
if isinstance(item, tuple) and len(item) >= 7:
|
|
return item[:6]
|
|
return item
|
|
|
|
@staticmethod
|
|
def _queue_item_user_id(item):
|
|
if isinstance(item, tuple) and len(item) >= 7:
|
|
return item[6]
|
|
return None
|
|
|
|
def _visible_queue_items(self, items):
|
|
current_user_id = self.get_current_user_id()
|
|
visible = []
|
|
for item in items:
|
|
owner_id = self._queue_item_user_id(item)
|
|
if self._can_access_owner(owner_id, current_user_id):
|
|
visible.append(self._unwrap_queue_item(item))
|
|
return visible
|
|
|
|
def user_queue_put(self, item):
|
|
"""Put an item in the user-specific queue."""
|
|
wrapped_item = self._wrap_queue_item(item)
|
|
with self.__prompt_queue.mutex:
|
|
heapq.heappush(self.__prompt_queue.queue, wrapped_item)
|
|
self.server.queue_updated()
|
|
self.__prompt_queue.not_empty.notify()
|
|
|
|
def user_queue_get(self, timeout=None):
|
|
"""Get an item from the user-specific queue."""
|
|
user_queue = self.__prompt_queue.queue
|
|
with self.__prompt_queue.not_empty:
|
|
while len(user_queue) == 0:
|
|
self.__prompt_queue.not_empty.wait(timeout=timeout)
|
|
if timeout is not None and len(user_queue) == 0:
|
|
return None
|
|
|
|
item = heapq.heappop(user_queue)
|
|
user_id = self._queue_item_user_id(item)
|
|
self.set_current_user_id(user_id, True)
|
|
|
|
i = self.__prompt_queue.task_counter
|
|
self.__prompt_queue.currently_running[i] = copy.deepcopy(item)
|
|
self.__prompt_queue.task_counter += 1
|
|
self.server.queue_updated()
|
|
return (self._unwrap_queue_item(item), i)
|
|
|
|
def user_queue_task_done(
|
|
self,
|
|
item_id,
|
|
history_result,
|
|
status: Optional["PromptQueue.ExecutionStatus"],
|
|
process_item=None,
|
|
):
|
|
"""Mark a user-specific queue task as done."""
|
|
with self.__prompt_queue.mutex:
|
|
wrapped_prompt = self.__prompt_queue.currently_running.pop(item_id)
|
|
prompt = self._unwrap_queue_item(wrapped_prompt)
|
|
user_id = self._queue_item_user_id(wrapped_prompt)
|
|
|
|
if len(self.__prompt_queue.history) > MAXIMUM_HISTORY_SIZE:
|
|
self.__prompt_queue.history.pop(next(iter(self.__prompt_queue.history)))
|
|
|
|
status_dict: Optional[dict] = None
|
|
if status is not None:
|
|
status_dict = copy.deepcopy(status._asdict())
|
|
|
|
if process_item is not None:
|
|
prompt = process_item(prompt)
|
|
|
|
self.__prompt_queue.history[prompt[1]] = {
|
|
"prompt": prompt,
|
|
"outputs": {},
|
|
"status": status_dict,
|
|
"user_id": user_id,
|
|
}
|
|
self.__prompt_queue.history[prompt[1]].update(history_result)
|
|
self.server.queue_updated()
|
|
|
|
def user_queue_get_current_queue(self):
|
|
"""Get the current user-specific queue."""
|
|
with self.__prompt_queue.mutex:
|
|
running = self._visible_queue_items(self.__prompt_queue.currently_running.values())
|
|
queued = self._visible_queue_items(copy.deepcopy(self.__prompt_queue.queue))
|
|
return (running, queued)
|
|
|
|
def user_queue_get_current_queue_volatile(self):
|
|
"""Get the current user-specific queue without deep copying items."""
|
|
with self.__prompt_queue.mutex:
|
|
running = self._visible_queue_items(self.__prompt_queue.currently_running.values())
|
|
queued = self._visible_queue_items(copy.copy(self.__prompt_queue.queue))
|
|
return (running, queued)
|
|
|
|
def user_queue_get_tasks_remaining(self):
|
|
with self.__prompt_queue.mutex:
|
|
return (
|
|
len(self._visible_queue_items(self.__prompt_queue.queue))
|
|
+ len(self._visible_queue_items(self.__prompt_queue.currently_running.values()))
|
|
)
|
|
|
|
def user_queue_wipe_queue(self):
|
|
"""Wipe the current user's queue, or all queue items for admins."""
|
|
with self.__prompt_queue.mutex:
|
|
current_user_id = self.get_current_user_id()
|
|
if self.is_admin_user(current_user_id):
|
|
self.__prompt_queue.queue = []
|
|
else:
|
|
self.__prompt_queue.queue = [
|
|
item
|
|
for item in self.__prompt_queue.queue
|
|
if self._queue_item_user_id(item) != current_user_id
|
|
]
|
|
self.server.queue_updated()
|
|
|
|
def user_queue_delete_queue_item(self, function):
|
|
"""Delete an item from the current user's queue."""
|
|
with self.__prompt_queue.mutex:
|
|
for x in range(len(self.__prompt_queue.queue)):
|
|
item = self.__prompt_queue.queue[x]
|
|
unwrapped_item = self._unwrap_queue_item(item)
|
|
if (
|
|
function(unwrapped_item)
|
|
and self._can_access_owner(self._queue_item_user_id(item))
|
|
):
|
|
self.__prompt_queue.queue.pop(x)
|
|
heapq.heapify(self.__prompt_queue.queue)
|
|
self.server.queue_updated()
|
|
return True
|
|
return False
|
|
|
|
def _visible_history(self):
|
|
current_user_id = self.get_current_user_id()
|
|
if self.is_admin_user(current_user_id):
|
|
return self.__prompt_queue.history
|
|
return {
|
|
k: v
|
|
for k, v in self.__prompt_queue.history.items()
|
|
if v.get("user_id") == current_user_id
|
|
}
|
|
|
|
def user_queue_get_history(self, prompt_id=None, max_items=None, offset=-1, map_function=None):
|
|
"""Get the current user's queue history."""
|
|
with self.__prompt_queue.mutex:
|
|
user_history = self._visible_history()
|
|
if prompt_id is None:
|
|
out = {}
|
|
i = 0
|
|
if offset < 0 and max_items is not None:
|
|
offset = len(user_history) - max_items
|
|
for k in user_history:
|
|
if i >= offset:
|
|
p = user_history[k]
|
|
if map_function is None:
|
|
p = copy.deepcopy(p)
|
|
else:
|
|
p = map_function(p)
|
|
out[k] = p
|
|
if max_items is not None and len(out) >= max_items:
|
|
break
|
|
i += 1
|
|
return out
|
|
if prompt_id in user_history:
|
|
p = user_history[prompt_id]
|
|
if map_function is None:
|
|
p = copy.deepcopy(p)
|
|
else:
|
|
p = map_function(p)
|
|
return {prompt_id: p}
|
|
return {}
|
|
|
|
def user_queue_wipe_history(self):
|
|
"""Wipe the current user's history, or all history for admins."""
|
|
with self.__prompt_queue.mutex:
|
|
current_user_id = self.get_current_user_id()
|
|
if self.is_admin_user(current_user_id):
|
|
self.__prompt_queue.history = {}
|
|
else:
|
|
self.__prompt_queue.history = {
|
|
k: v
|
|
for k, v in self.__prompt_queue.history.items()
|
|
if v.get("user_id") != current_user_id
|
|
}
|
|
|
|
def user_queue_delete_history_item(self, id_to_delete):
|
|
with self.__prompt_queue.mutex:
|
|
history_item = self.__prompt_queue.history.get(id_to_delete)
|
|
if history_item and self._can_access_owner(history_item.get("user_id")):
|
|
self.__prompt_queue.history.pop(id_to_delete, None)
|
|
|
|
def patch_prompt_queue(self):
|
|
"""Patch the prompt queue with user-specific methods."""
|
|
self.__prompt_queue.put = self.user_queue_put
|
|
self.__prompt_queue.get = self.user_queue_get
|
|
self.__prompt_queue.task_done = self.user_queue_task_done
|
|
self.__prompt_queue.get_current_queue = self.user_queue_get_current_queue
|
|
self.__prompt_queue.get_current_queue_volatile = self.user_queue_get_current_queue_volatile
|
|
self.__prompt_queue.get_tasks_remaining = self.user_queue_get_tasks_remaining
|
|
self.__prompt_queue.wipe_queue = self.user_queue_wipe_queue
|
|
self.__prompt_queue.delete_queue_item = self.user_queue_delete_queue_item
|
|
self.__prompt_queue.get_history = self.user_queue_get_history
|
|
self.__prompt_queue.wipe_history = self.user_queue_wipe_history
|
|
self.__prompt_queue.delete_history_item = self.user_queue_delete_history_item
|
|
|
|
def _default_owner_id(self) -> str:
|
|
return self.get_context_user_id() or ""
|
|
|
|
def patch_assets(self) -> None:
|
|
"""Patch ComfyUI asset ownership and visibility without editing ComfyUI."""
|
|
try:
|
|
import app.assets.database.queries as queries
|
|
import app.assets.database.queries.common as common_queries
|
|
import app.assets.database.queries.asset_reference as asset_reference_queries
|
|
import app.assets.database.queries.tags as tag_queries
|
|
import app.assets.services as asset_services
|
|
import app.assets.services.ingest as ingest
|
|
import app.assets.services.asset_management as asset_management
|
|
import app.assets.services.tagging as tagging
|
|
from app.assets.database.models import AssetReference
|
|
from app.database.db import create_session
|
|
except Exception:
|
|
return
|
|
|
|
def build_visible_owner_clause(owner_id: str):
|
|
owner_id = (owner_id or "").strip()
|
|
if self.is_admin_user(owner_id):
|
|
return asset_reference_queries.sa.true()
|
|
if owner_id == "":
|
|
return AssetReference.owner_id == ""
|
|
return AssetReference.owner_id == owner_id
|
|
|
|
def get_reference_with_owner_check(session, reference_id: str, owner_id: str):
|
|
ref = asset_reference_queries.get_reference_by_id(
|
|
session, reference_id=reference_id
|
|
)
|
|
if not ref or ref.deleted_at is not None:
|
|
raise ValueError(f"AssetReference {reference_id} not found")
|
|
if self.is_admin_user(owner_id):
|
|
return ref
|
|
if not ref.owner_id or ref.owner_id != owner_id:
|
|
raise PermissionError("not owner")
|
|
return ref
|
|
|
|
common_queries.build_visible_owner_clause = build_visible_owner_clause
|
|
asset_reference_queries.build_visible_owner_clause = build_visible_owner_clause
|
|
tag_queries.build_visible_owner_clause = build_visible_owner_clause
|
|
queries.build_visible_owner_clause = build_visible_owner_clause
|
|
|
|
asset_reference_queries.get_reference_with_owner_check = get_reference_with_owner_check
|
|
queries.get_reference_with_owner_check = get_reference_with_owner_check
|
|
asset_management.get_reference_with_owner_check = get_reference_with_owner_check
|
|
tagging.get_reference_with_owner_check = get_reference_with_owner_check
|
|
|
|
original_ingest_existing_file = ingest.ingest_existing_file
|
|
original_upload_from_temp_path = ingest.upload_from_temp_path
|
|
original_register_file_in_place = ingest.register_file_in_place
|
|
original_create_from_hash = ingest.create_from_hash
|
|
|
|
def ingest_existing_file(*args, owner_id="", **kwargs):
|
|
owner_id = owner_id or self._default_owner_id()
|
|
return original_ingest_existing_file(*args, owner_id=owner_id, **kwargs)
|
|
|
|
def upload_from_temp_path(*args, owner_id="", **kwargs):
|
|
owner_id = owner_id or self._default_owner_id()
|
|
return original_upload_from_temp_path(*args, owner_id=owner_id, **kwargs)
|
|
|
|
def register_file_in_place(*args, owner_id="", **kwargs):
|
|
owner_id = owner_id or self._default_owner_id()
|
|
return original_register_file_in_place(*args, owner_id=owner_id, **kwargs)
|
|
|
|
def create_from_hash(*args, owner_id="", **kwargs):
|
|
owner_id = owner_id or self._default_owner_id()
|
|
return original_create_from_hash(*args, owner_id=owner_id, **kwargs)
|
|
|
|
ingest.ingest_existing_file = ingest_existing_file
|
|
ingest.upload_from_temp_path = upload_from_temp_path
|
|
ingest.register_file_in_place = register_file_in_place
|
|
ingest.create_from_hash = create_from_hash
|
|
|
|
asset_services.upload_from_temp_path = upload_from_temp_path
|
|
asset_services.register_file_in_place = register_file_in_place
|
|
asset_services.create_from_hash = create_from_hash
|
|
|
|
def resolve_hash_to_path(asset_hash: str, owner_id: str = ""):
|
|
with create_session() as session:
|
|
asset = asset_management.queries_get_asset_by_hash(session, asset_hash)
|
|
if not asset:
|
|
return None
|
|
refs = asset_management.list_references_by_asset_id(
|
|
session, asset_id=asset.id
|
|
)
|
|
if self.is_admin_user(owner_id):
|
|
visible = refs
|
|
else:
|
|
visible = [r for r in refs if r.owner_id == owner_id]
|
|
abs_path = asset_management.select_best_live_path(visible)
|
|
if not abs_path:
|
|
return None
|
|
display_name = os.path.basename(abs_path)
|
|
for ref in visible:
|
|
if ref.file_path == abs_path and ref.name:
|
|
display_name = ref.name
|
|
break
|
|
content_type = (
|
|
asset.mime_type
|
|
or mimetypes.guess_type(display_name)[0]
|
|
or "application/octet-stream"
|
|
)
|
|
return asset_management.DownloadResolutionResult(
|
|
abs_path=abs_path,
|
|
content_type=content_type,
|
|
download_name=display_name,
|
|
)
|
|
|
|
asset_management.resolve_hash_to_path = resolve_hash_to_path
|
|
asset_services.resolve_hash_to_path = resolve_hash_to_path
|
|
|
|
route_module = sys.modules.get("app.assets.api.routes")
|
|
if route_module:
|
|
route_module.upload_from_temp_path = upload_from_temp_path
|
|
route_module.create_from_hash = create_from_hash
|
|
|
|
for module_name in ("server", "__main__", "main"):
|
|
module = sys.modules.get(module_name)
|
|
if not module:
|
|
continue
|
|
if hasattr(module, "register_file_in_place"):
|
|
module.register_file_in_place = register_file_in_place
|
|
if hasattr(module, "resolve_hash_to_path"):
|
|
module.resolve_hash_to_path = resolve_hash_to_path
|
|
|
|
def create_manager_access_control_middleware(
|
|
self, manager_directory: str = "/extensions/comfyui-manager", manager_routes: tuple = ()
|
|
) -> web.middleware:
|
|
"""Create middleware for manager access control."""
|
|
|
|
@web.middleware
|
|
async def manager_access_control_middleware(
|
|
request: web.Request, handler
|
|
) -> web.Response:
|
|
user_id = request.get("user_id")
|
|
|
|
if self.users_db.get_admin_user()[0] == user_id or (
|
|
not request.path.startswith(manager_routes)
|
|
and not request.path.lower().startswith(manager_directory)
|
|
):
|
|
return await handler(request)
|
|
|
|
return web.HTTPForbidden(
|
|
reason="You do not have access to comfyui manager."
|
|
)
|
|
|
|
return manager_access_control_middleware
|