Scope gallery image endpoints by owner (#3001)

tests/test_gallery_endpoint_matching.py

@@ -1,34 +1,11 @@
-import ast
-from pathlib import Path
-
 def test_gallery_url_normalization_bug():
-    # Read and parse the actual source file
-    source_path = Path("routes/gallery_routes.py")
-    assert source_path.exists(), "gallery_routes.py could not be found"
-    
-    source = source_path.read_text(encoding="utf-8")
-    tree = ast.parse(source)
-    
-    # Locate the comparison node within harmonize_image that references ep.base_url and base
-    compare_node = None
-    for node in ast.walk(tree):
-        if isinstance(node, ast.Compare):
-            segment = ast.get_source_segment(source, node) or ""
-            if "ep.base_url" in segment and "base" in segment and "_norm_url" not in segment:
-                compare_node = node
-                break
-                
-    assert compare_node is not None, "Could not find the ep.base_url vs base comparison inside gallery_routes.py"
-    
-    # Compile the compare node into an expression
-    expr = ast.Expression(body=compare_node)
-    compiled_code = compile(expr, "<string>", "eval")
-    
+    from routes.gallery_routes import _normalize_image_endpoint_base
+
     def check_match(ep_url: str, base_url: str) -> bool:
-        class MockEP:
-            def __init__(self, url):
-                self.base_url = url
-        return eval(compiled_code, {}, {"ep": MockEP(ep_url), "base": base_url})
+        return (
+            _normalize_image_endpoint_base(ep_url)
+            == _normalize_image_endpoint_base(base_url)
+        )
 
     # Test cases that SHOULD NOT match under a correct implementation
     # (Buggy rstrip('/v1') logic incorrectly treats these as equal)

tests/test_gallery_image_endpoint_owner_scope.py

@@ -0,0 +1,126 @@
+"""Owner-scope regression for gallery image endpoint selection.
+
+The image editor/upscale proxies select ``ModelEndpoint`` rows and may copy the
+row's stored ``api_key`` for OpenAI-compatible image endpoints. That lookup must
+only consider endpoints visible to the caller, otherwise users sharing the same
+base URL can borrow another account's private image API key.
+"""
+
+from types import SimpleNamespace
+
+import routes.gallery_routes as gallery_routes
+
+
+class _Predicate:
+    def __init__(self, check):
+        self._check = check
+
+    def __call__(self, row):
+        return self._check(row)
+
+    def __or__(self, other):
+        return _Predicate(lambda row: self(row) or other(row))
+
+
+class _Column:
+    def __init__(self, name):
+        self.name = name
+
+    def __eq__(self, value):
+        return _Predicate(lambda row: getattr(row, self.name) == value)
+
+
+class _ModelEndpoint:
+    base_url = _Column("base_url")
+    model_type = _Column("model_type")
+    is_enabled = _Column("is_enabled")
+    owner = _Column("owner")
+
+
+class _Query:
+    def __init__(self, rows):
+        self._rows = list(rows)
+
+    def filter(self, *predicates):
+        self._rows = [row for row in self._rows if all(pred(row) for pred in predicates)]
+        return self
+
+    def all(self):
+        return list(self._rows)
+
+
+class _DB:
+    def __init__(self, rows):
+        self._rows = rows
+
+    def query(self, model):
+        assert model is _ModelEndpoint
+        return _Query(self._rows)
+
+
+def _ep(base_url, owner, *, enabled=True, model_type="image", api_key="sk-secret"):
+    return SimpleNamespace(
+        base_url=base_url,
+        owner=owner,
+        is_enabled=enabled,
+        model_type=model_type,
+        api_key=api_key,
+    )
+
+
+def _patch_model(monkeypatch):
+    monkeypatch.setattr(gallery_routes, "ModelEndpoint", _ModelEndpoint)
+
+
+URL = "https://api.example.com/v1"
+
+
+def test_first_visible_image_endpoint_rejects_another_owner(monkeypatch):
+    _patch_model(monkeypatch)
+    rows = [_ep(URL, "bob")]
+
+    assert gallery_routes._first_visible_image_endpoint(_DB(rows), "alice") is None
+
+
+def test_first_visible_image_endpoint_prefers_callers_own_row(monkeypatch):
+    _patch_model(monkeypatch)
+    rows = [_ep(URL, None, api_key="shared"), _ep(URL, "alice", api_key="own")]
+
+    ep = gallery_routes._first_visible_image_endpoint(_DB(rows), "alice")
+
+    assert ep is not None
+    assert ep.owner == "alice"
+    assert ep.api_key == "own"
+
+
+def test_visible_image_endpoint_for_base_rejects_same_url_other_owner(monkeypatch):
+    _patch_model(monkeypatch)
+    rows = [_ep(URL, "bob")]
+
+    assert gallery_routes._visible_image_endpoint_for_base(_DB(rows), URL, "alice") is None
+
+
+def test_visible_image_endpoint_for_base_allows_shared_or_own(monkeypatch):
+    _patch_model(monkeypatch)
+    rows = [
+        _ep("https://other.example/v1", "alice"),
+        _ep(URL, None, api_key="shared"),
+        _ep(URL, "alice", api_key="own"),
+    ]
+
+    ep = gallery_routes._visible_image_endpoint_for_base(_DB(rows), "https://api.example.com", "alice")
+
+    assert ep is not None
+    assert ep.owner == "alice"
+    assert ep.api_key == "own"
+    assert ep.base_url == URL
+
+
+def test_image_endpoint_owner_filter_is_noop_in_single_user_mode(monkeypatch):
+    _patch_model(monkeypatch)
+    rows = [_ep(URL, "bob")]
+
+    ep = gallery_routes._visible_image_endpoint_for_base(_DB(rows), URL, None)
+
+    assert ep is not None
+    assert ep.owner == "bob"