mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-03-15 06:07:33 +08:00
Add owner_id check to resolve_hash_to_path
Filter asset references by owner visibility so the /view endpoint only resolves hashes for assets the requesting user can access. Adds table-driven tests for owner visibility cases. Amp-Thread-ID: https://ampcode.com/threads/T-019ce377-8bde-7048-bc28-a9df063409f9 Co-authored-by: Amp <amp@ampcode.com>
This commit is contained in:
parent
aa6bd95596
commit
6587a7394e
@ -282,9 +282,13 @@ def list_assets_page(
|
|||||||
|
|
||||||
def resolve_hash_to_path(
|
def resolve_hash_to_path(
|
||||||
asset_hash: str,
|
asset_hash: str,
|
||||||
|
owner_id: str = "",
|
||||||
) -> DownloadResolutionResult | None:
|
) -> DownloadResolutionResult | None:
|
||||||
"""Resolve a blake3 hash to an on-disk file path.
|
"""Resolve a blake3 hash to an on-disk file path.
|
||||||
|
|
||||||
|
Only references visible to *owner_id* are considered (owner-less
|
||||||
|
references are always visible).
|
||||||
|
|
||||||
Returns a DownloadResolutionResult with abs_path, content_type, and
|
Returns a DownloadResolutionResult with abs_path, content_type, and
|
||||||
download_name, or None if no asset or live path is found.
|
download_name, or None if no asset or live path is found.
|
||||||
"""
|
"""
|
||||||
@ -293,11 +297,15 @@ def resolve_hash_to_path(
|
|||||||
if not asset:
|
if not asset:
|
||||||
return None
|
return None
|
||||||
refs = list_references_by_asset_id(session, asset_id=asset.id)
|
refs = list_references_by_asset_id(session, asset_id=asset.id)
|
||||||
abs_path = select_best_live_path(refs)
|
visible = [
|
||||||
|
r for r in refs
|
||||||
|
if r.owner_id == "" or r.owner_id == owner_id
|
||||||
|
]
|
||||||
|
abs_path = select_best_live_path(visible)
|
||||||
if not abs_path:
|
if not abs_path:
|
||||||
return None
|
return None
|
||||||
display_name = os.path.basename(abs_path)
|
display_name = os.path.basename(abs_path)
|
||||||
for ref in refs:
|
for ref in visible:
|
||||||
if ref.file_path == abs_path and ref.name:
|
if ref.file_path == abs_path and ref.name:
|
||||||
display_name = ref.name
|
display_name = ref.name
|
||||||
break
|
break
|
||||||
|
|||||||
@ -504,7 +504,8 @@ class PromptServer():
|
|||||||
# node preview, it constructs /view?filename=<asset_hash>, so this
|
# node preview, it constructs /view?filename=<asset_hash>, so this
|
||||||
# endpoint must resolve blake3 hashes to their on-disk file paths.
|
# endpoint must resolve blake3 hashes to their on-disk file paths.
|
||||||
if filename.startswith("blake3:"):
|
if filename.startswith("blake3:"):
|
||||||
result = resolve_hash_to_path(filename)
|
owner_id = self.user_manager.get_request_user_id(request)
|
||||||
|
result = resolve_hash_to_path(filename, owner_id=owner_id)
|
||||||
if result is None:
|
if result is None:
|
||||||
return web.Response(status=404)
|
return web.Response(status=404)
|
||||||
file, filename = result.abs_path, result.download_name
|
file, filename = result.abs_path, result.download_name
|
||||||
|
|||||||
@ -11,6 +11,7 @@ from app.assets.services import (
|
|||||||
delete_asset_reference,
|
delete_asset_reference,
|
||||||
set_asset_preview,
|
set_asset_preview,
|
||||||
)
|
)
|
||||||
|
from app.assets.services.asset_management import resolve_hash_to_path
|
||||||
|
|
||||||
|
|
||||||
def _make_asset(session: Session, hash_val: str = "blake3:test", size: int = 1024) -> Asset:
|
def _make_asset(session: Session, hash_val: str = "blake3:test", size: int = 1024) -> Asset:
|
||||||
@ -266,3 +267,42 @@ class TestSetAssetPreview:
|
|||||||
preview_asset_id=None,
|
preview_asset_id=None,
|
||||||
owner_id="user2",
|
owner_id="user2",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class TestResolveHashToPath:
|
||||||
|
def test_returns_none_for_unknown_hash(self, mock_create_session):
|
||||||
|
result = resolve_hash_to_path("blake3:" + "a" * 64)
|
||||||
|
assert result is None
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
"ref_owner, query_owner, expect_found",
|
||||||
|
[
|
||||||
|
("user1", "user1", True),
|
||||||
|
("user1", "user2", False),
|
||||||
|
("", "anyone", True),
|
||||||
|
("", "", True),
|
||||||
|
],
|
||||||
|
ids=[
|
||||||
|
"owner_sees_own_ref",
|
||||||
|
"other_owner_blocked",
|
||||||
|
"ownerless_visible_to_anyone",
|
||||||
|
"ownerless_visible_to_empty",
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_owner_visibility(
|
||||||
|
self, ref_owner, query_owner, expect_found,
|
||||||
|
mock_create_session, session: Session, temp_dir,
|
||||||
|
):
|
||||||
|
f = temp_dir / "file.bin"
|
||||||
|
f.write_bytes(b"data")
|
||||||
|
asset = _make_asset(session, hash_val="blake3:" + "b" * 64)
|
||||||
|
ref = _make_reference(session, asset, name="file.bin", owner_id=ref_owner)
|
||||||
|
ref.file_path = str(f)
|
||||||
|
session.commit()
|
||||||
|
|
||||||
|
result = resolve_hash_to_path(asset.hash, owner_id=query_owner)
|
||||||
|
if expect_found:
|
||||||
|
assert result is not None
|
||||||
|
assert result.abs_path == str(f)
|
||||||
|
else:
|
||||||
|
assert result is None
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user