diff --git a/core/request_utils.py b/core/request_utils.py index 49691ea..06a75b0 100644 --- a/core/request_utils.py +++ b/core/request_utils.py @@ -1,7 +1,24 @@ """Helpers for validating and reading common request payload fields.""" +from collections.abc import Mapping from typing import Any, Optional, Tuple +from .type_utils import normalize_sha256 + + +def extract_request_sha256( + payload: Mapping[str, Any], + *, + keys: Tuple[str, ...], +) -> str: + """Normalize the first non-empty SHA256 alias found in a request payload.""" + value = "" + for key in keys: + value = payload.get(key) + if value: + break + return normalize_sha256(value) + async def read_optional_object_payload(request: Any) -> dict[str, Any]: """Read an optional JSON object, returning an empty dict on invalid input.""" diff --git a/core/services/civitai_search_service.py b/core/services/civitai_search_service.py index 4f57e51..973b7f4 100644 --- a/core/services/civitai_search_service.py +++ b/core/services/civitai_search_service.py @@ -1,6 +1,7 @@ """CivitAI model search service.""" from .. import path_utils +from ..request_utils import extract_request_sha256 from ..routes.context import RouteContext from ..sources.common import collect_download_urls from .model_utils import CivitAISearchDependencies, ModelServiceDependencies @@ -61,7 +62,6 @@ class CivitAISearchService(ModelServiceDependencies): is_path_in_configured_model_roots = self.is_path_in_configured_model_roots looks_like_model_file = self.looks_like_model_file normalize_category_to_model_type = self.normalize_category_to_model_type - normalize_sha256 = self.normalize_sha256 read_json_safe = self.read_json_safe request_public_url = self.request_public_url resolve_civarchive_by_hash = self.resolve_civarchive_by_hash @@ -77,11 +77,9 @@ class CivitAISearchService(ModelServiceDependencies): force_refresh = to_bool( data.get("force_refresh") or data.get("force"), False ) - provided_hash = normalize_sha256( - data.get("sha256") - or data.get("hash") - or data.get("file_hash") - or "" + provided_hash = extract_request_sha256( + data, + keys=("sha256", "hash", "file_hash"), ) hf_token = data.get("hf_token", "") brave_search_api_key = data.get("brave_search_api_key", "") diff --git a/core/services/hash_service.py b/core/services/hash_service.py index 93e41fe..d1a6295 100644 --- a/core/services/hash_service.py +++ b/core/services/hash_service.py @@ -6,7 +6,7 @@ import uuid from .. import path_utils from ..local_hash_matches import collect_local_hash_matches_for_result -from ..request_utils import validate_workflow_payload +from ..request_utils import extract_request_sha256, validate_workflow_payload from ..routes.context import RouteContext @@ -254,11 +254,9 @@ class HashService: async def local_matches_by_hash(self, request): """Search local model metadata sidecars for a remote SHA256.""" data = await request.json() - sha256 = self.normalize_sha256( - data.get("sha256") - or data.get("hash") - or data.get("SHA256") - or "" + sha256 = extract_request_sha256( + data, + keys=("sha256", "hash", "SHA256"), ) if not sha256: return self.web.json_response( diff --git a/core/services/search_orchestrator.py b/core/services/search_orchestrator.py index 7b51d92..086cff3 100644 --- a/core/services/search_orchestrator.py +++ b/core/services/search_orchestrator.py @@ -1,8 +1,8 @@ """Request orchestration for model source searches.""" from ..local_hash_matches import collect_local_hash_matches_for_result +from ..request_utils import extract_request_sha256 from ..routes.context import RouteContext -from ..type_utils import normalize_sha256 from .search_cache import SearchResultCache from .search_dependencies import SearchDependencies from .search_providers import ( @@ -138,8 +138,9 @@ class SearchOrchestrator: try: data = await request.json() filename = str(data.get("filename", "") or "").strip() - sha256 = normalize_sha256( - data.get("sha256") or data.get("hash") or data.get("file_hash") + sha256 = extract_request_sha256( + data, + keys=("sha256", "hash", "file_hash"), ) category = data.get("category", "") base_model_context = data.get("base_model_context", "") diff --git a/tests/test_request_utils.py b/tests/test_request_utils.py index ce46ebb..d71ce28 100644 --- a/tests/test_request_utils.py +++ b/tests/test_request_utils.py @@ -3,7 +3,29 @@ from unittest.mock import AsyncMock import pytest -from core.request_utils import read_optional_object_payload +from core.request_utils import extract_request_sha256, read_optional_object_payload + + +def test_extract_request_sha256_preserves_alias_precedence_and_normalization(): + sha256 = "a" * 64 + + assert extract_request_sha256( + { + "sha256": "", + "hash": f"sha256:{sha256.upper()}", + "file_hash": "b" * 64, + }, + keys=("sha256", "hash", "file_hash"), + ) == sha256 + assert extract_request_sha256( + {"sha256": "invalid", "hash": sha256}, + keys=("sha256", "hash"), + ) == "" + assert extract_request_sha256( + {"SHA256": sha256}, + keys=("sha256", "hash", "SHA256"), + ) == sha256 + assert extract_request_sha256({}, keys=("sha256", "hash")) == "" @pytest.mark.asyncio