From bb00814d7adba9c6738732db0d5a7b3418d5c4d4 Mon Sep 17 00:00:00 2001 From: InfinityPacer <160988576+InfinityPacer@users.noreply.github.com> Date: Mon, 27 Jul 2026 12:13:38 +0800 Subject: [PATCH] =?UTF-8?q?feat(login):=20=E6=8F=90=E4=BE=9B=E6=97=A0?= =?UTF-8?q?=E7=8A=B6=E6=80=81=E7=AD=BE=E5=90=8D=E5=90=8C=E6=BA=90=E5=A3=81?= =?UTF-8?q?=E7=BA=B8=E4=BB=A3=E7=90=86=20(#6200)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- app/api/endpoints/login.py | 147 +++++++++- app/helper/image.py | 151 +++++++++- app/utils/http.py | 33 +++ tests/test_login_wallpaper_proxy.py | 409 ++++++++++++++++++++++++++++ 4 files changed, 733 insertions(+), 7 deletions(-) create mode 100644 tests/test_login_wallpaper_proxy.py diff --git a/app/api/endpoints/login.py b/app/api/endpoints/login.py index 1b1098a7..24dd9a02 100644 --- a/app/api/endpoints/login.py +++ b/app/api/endpoints/login.py @@ -1,7 +1,8 @@ from datetime import timedelta -from typing import Any, List, Annotated +from typing import Annotated, Any, List +from urllib.parse import quote, urlparse, urlunparse -from fastapi import APIRouter, Depends, Form, HTTPException, Request, Response +from fastapi import APIRouter, Depends, Form, Header, HTTPException, Request, Response from fastapi.security import OAuth2PasswordRequestForm from fastapi.responses import JSONResponse @@ -11,11 +12,19 @@ from app.core import security from app.core.config import settings from app.db.systemconfig_oper import SystemConfigOper from app.helper.sites import SitesHelper # noqa -from app.helper.image import WallpaperHelper +from app.helper.image import ImageHelper, WallpaperHelper from app.schemas.types import SystemConfigKey +from app.utils.crypto import HashUtils +from app.utils.http import RequestUtils +from app.utils.security import SecurityUtils +from app.utils.url import UrlUtils router = APIRouter() +_LOGIN_WALLPAPER_MAX_BYTES = 32 * 1024 * 1024 +_LOGIN_WALLPAPER_PUBLIC_PURPOSE = "login-wallpaper-public" +_LOGIN_WALLPAPER_MEDIA_PURPOSE = "login-wallpaper-media" + @router.post("/access-token", summary="获取token", response_model=schemas.Token) def login_access_token( @@ -95,8 +104,134 @@ def wallpaper() -> Any: @router.get("/wallpapers", summary="登录页面电影海报列表", response_model=List[str]) -def wallpapers() -> Any: +def wallpapers(same_origin: bool = False) -> Any: """ - 获取登录页面电影海报 + 获取登录页面电影海报。 + + 默认保持外链列表合同；同源模式只对绝对 HTTP(S) 地址做一对一签名转换，不改变 + 来源数量、顺序、重复项或相对地址。 """ - return WallpaperHelper().get_wallpapers() + wallpaper_urls = WallpaperHelper().get_wallpapers() + if not same_origin: + return wallpaper_urls + + purpose = ( + _LOGIN_WALLPAPER_MEDIA_PURPOSE + if settings.WALLPAPER == "mediaserver" + else _LOGIN_WALLPAPER_PUBLIC_PURPOSE + ) + return [_login_wallpaper_proxy_url(url, purpose) for url in wallpaper_urls] + + +def _login_wallpaper_proxy_url(url: str, purpose: str) -> str: + """ + 将可代理的绝对壁纸地址转换为登录页专用同源签名地址。 + + `//cdn.example/one.jpg` 这类网络路径引用会被浏览器解析成跨源地址，因此按 + 当前页面协议补全后同样走代理；只有不带 netloc 的相对地址才原样返回。 + """ + parsed = urlparse(url) + if not parsed.netloc: + return url + if not parsed.scheme: + url = urlunparse(parsed._replace(scheme="https")) + elif parsed.scheme not in {"http", "https"}: + return url + signed_url = SecurityUtils.sign_url(url, purpose=purpose) + return ( + f"{settings.API_V1_STR}/login/wallpapers/image" + f"?url={quote(signed_url, safe='')}" + ) + + +def _url_origin(url: str) -> tuple[str, str, int | None] | None: + """返回 HTTP(S) 地址的规范 origin，拒绝缺少主机或携带用户信息的地址。""" + try: + parsed = urlparse(url) + port = parsed.port + except ValueError: + return None + if ( + parsed.scheme not in {"http", "https"} + or not parsed.hostname + or parsed.username + or parsed.password + ): + return None + if port is None: + port = 443 if parsed.scheme == "https" else 80 + return parsed.scheme, parsed.hostname.lower(), port + + +@router.get("/wallpapers/image", summary="登录页面同源壁纸") +async def wallpaper_image( + url: str, + if_none_match: Annotated[str | None, Header()] = None, +) -> Response: + """ + 读取后端壁纸来源签发的图片；客户端无法修改目标后继续复用签名。 + + 响应带内容 ETag，缓存过期后条件请求命中时只返回 304。 + """ + source_url = SecurityUtils.verify_signed_url( + url, purpose=_LOGIN_WALLPAPER_PUBLIC_PURPOSE + ) + media_source = False + if not source_url: + source_url = SecurityUtils.verify_signed_url( + url, purpose=_LOGIN_WALLPAPER_MEDIA_PURPOSE + ) + media_source = bool(source_url) + source_origin = _url_origin(source_url or "") + if not source_url or not source_origin: + raise HTTPException(status_code=404, detail="Wallpaper not found") + + async def is_safe_public_target(target_url: str) -> bool: + """自定义公共来源可跨域跳转，但每个目标都必须通过 DNS/私网校验。""" + target_origin = _url_origin(target_url) + if not target_origin: + return False + return await SecurityUtils.is_safe_image_url_async( + target_url, + {target_origin[1]}, + allowed_private_ranges=settings.IMAGE_PROXY_ALLOWED_PRIVATE_RANGES, + ) + + async def is_safe_redirect(target_url: str) -> bool: + """媒体服务器签名只授权原 origin；其它跳转按公共图片目标重新校验。""" + target_origin = _url_origin(target_url) + if media_source and target_origin == source_origin: + return True + return await is_safe_public_target(target_url) + + if not media_source and not await is_safe_public_target(source_url): + raise HTTPException(status_code=404, detail="Wallpaper not found") + + # 媒体服务器来源额外继承原 origin 的私网授权，与公共来源的重定向授权范围不同， + # 因此按来源类型声明策略标识，只让同策略的并发请求共享一次抓取。 + redirect_policy = ( + f"{_LOGIN_WALLPAPER_MEDIA_PURPOSE}:{source_origin[0]}://" + f"{source_origin[1]}:{source_origin[2]}" + if media_source + else _LOGIN_WALLPAPER_PUBLIC_PURPOSE + ) + content = await ImageHelper().async_fetch_image_guarded( + url=source_url, + redirect_validator=is_safe_redirect, + redirect_policy=redirect_policy, + max_bytes=_LOGIN_WALLPAPER_MAX_BYTES, + use_cache=True, + ) + if not content: + raise HTTPException(status_code=502, detail="Wallpaper unavailable") + + etag = f'"{HashUtils.md5(content)}"' + headers = RequestUtils.generate_cache_headers(etag, max_age=86400) + if RequestUtils.if_none_match_matches(if_none_match, etag): + return Response(status_code=304, headers=headers) + + return Response( + content=content, + media_type=UrlUtils.get_mime_type(source_url, "image/jpeg"), + headers=headers, + ) diff --git a/app/helper/image.py b/app/helper/image.py index f76b3088..70a4fdea 100644 --- a/app/helper/image.py +++ b/app/helper/image.py @@ -1,6 +1,9 @@ +import asyncio import io +import threading from pathlib import Path -from typing import Optional, List +from typing import Awaitable, Callable, Optional, List +from urllib.parse import urljoin from PIL import Image @@ -177,6 +180,8 @@ class ImageHelper(metaclass=Singleton): _ttl = settings.GLOBAL_IMAGE_CACHE_DAYS * 24 * 3600 self.file_cache = FileCache(base=_base_path, ttl=_ttl) self.async_file_cache = AsyncFileCache(base=_base_path, ttl=_ttl) + self._guarded_fetch_tasks: dict[tuple, asyncio.Task[Optional[bytes]]] = {} + self._guarded_fetch_tasks_lock = threading.Lock() @staticmethod def _prepare_cache_path(url: str) -> str: @@ -286,3 +291,147 @@ class ImageHelper(metaclass=Singleton): # 保存缓存 await self.async_file_cache.set(cache_path, content, region="images") return content + + async def async_fetch_image_guarded( + self, + url: str, + *, + redirect_validator: Callable[[str], Awaitable[bool]], + redirect_policy: str, + max_bytes: int, + proxy: Optional[bool] = None, + use_cache: bool = True, + max_redirects: int = 3, + ) -> Optional[bytes]: + """ + 以有界流式请求抓取需要逐跳校验的图片。 + + 每个重定向目标必须重新通过调用方的安全校验，字节上限和图片有效性检查在 + 写入共享缓存前生效。只有抓取策略完全等价的并发调用才共享一次远端抓取： + 合并键包含缓存键、`redirect_policy`、字节上限、代理与重定向上限，避免某 + 次调用收到超出自身上限的图片，或沿用他人的重定向授权。 + + :param redirect_validator: 逐跳校验重定向目标的协程 + :param redirect_policy: 描述该校验授权范围的稳定标识；`redirect_validator` + 通常是每次请求新建的闭包，无法按对象身份判断等价，由调用方显式声明 + """ + if not url or max_bytes <= 0: + return None + + cache_path = self._prepare_cache_path(url) + if use_cache: + content = await self.async_file_cache.get(cache_path, region="images") + if content: + if len(content) <= max_bytes and self._validate_image(content): + return content + await self.async_file_cache.delete(cache_path, region="images") + + task_key = ( + cache_path, + redirect_policy, + max_bytes, + proxy, + use_cache, + max_redirects, + ) + loop = asyncio.get_running_loop() + with self._guarded_fetch_tasks_lock: + task = self._guarded_fetch_tasks.get(task_key) + if task is None or task.get_loop() is not loop: + task = loop.create_task( + self._download_guarded_image( + url=url, + cache_path=cache_path, + redirect_validator=redirect_validator, + max_bytes=max_bytes, + proxy=proxy, + use_cache=use_cache, + max_redirects=max_redirects, + ) + ) + self._guarded_fetch_tasks[task_key] = task + task.add_done_callback( + lambda completed, key=task_key: self._forget_guarded_fetch_task( + key, completed + ) + ) + + return await asyncio.shield(task) + + def _forget_guarded_fetch_task( + self, task_key: tuple, task: asyncio.Task[Optional[bytes]] + ) -> None: + """抓取完成后只移除仍指向该任务的合并键，避免旧任务清除后继任务。""" + with self._guarded_fetch_tasks_lock: + if self._guarded_fetch_tasks.get(task_key) is task: + self._guarded_fetch_tasks.pop(task_key, None) + + async def _download_guarded_image( + self, + *, + url: str, + cache_path: str, + redirect_validator: Callable[[str], Awaitable[bool]], + max_bytes: int, + proxy: Optional[bool], + use_cache: bool, + max_redirects: int, + ) -> Optional[bytes]: + """执行一次受保护图片抓取，并在任务内复查并写入共享缓存。""" + if use_cache: + content = await self.async_file_cache.get(cache_path, region="images") + if content: + if len(content) <= max_bytes and self._validate_image(content): + return content + await self.async_file_cache.delete(cache_path, region="images") + + current_url = url + redirects = 0 + + while True: + params = self._get_request_params(current_url, proxy, cookies=None) + request = AsyncRequestUtils(**params, follow_redirects=False) + + async with request.get_stream(current_url) as response: + if response is None: + return None + + if response.status_code in {301, 302, 303, 307, 308}: + location = response.headers.get("location") + if not location or redirects >= max_redirects: + return None + next_url = urljoin(current_url, location) + if not await redirect_validator(next_url): + return None + current_url = next_url + redirects += 1 + continue + + if response.status_code != 200: + logger.warning( + "登录壁纸抓取失败，状态码: %s", + response.status_code, + ) + return None + + content_length = response.headers.get("content-length") + if content_length: + try: + if int(content_length) > max_bytes: + return None + except ValueError: + pass + + payload = bytearray() + async for chunk in response.aiter_bytes(): + payload.extend(chunk) + if len(payload) > max_bytes: + return None + break + + content = bytes(payload) + if not self._validate_image(content): + return None + if use_cache: + await self.async_file_cache.set(cache_path, content, region="images") + return content diff --git a/app/utils/http.py b/app/utils/http.py index 12662d5e..a6609fe8 100644 --- a/app/utils/http.py +++ b/app/utils/http.py @@ -720,6 +720,39 @@ class RequestUtils: return cache_headers + @staticmethod + def if_none_match_matches(if_none_match: Optional[str], etag: Optional[str]) -> bool: + """ + 按 entity-tag 语义判断 If-None-Match 是否命中当前响应的 ETag + + 条件缓存的客户端会发送带引号的标签、弱标签、逗号分隔的标签列表或 `*`， + 直接做裸摘要字符串全等比较会让 304 永远无法命中。这里按 RFC 9110 + 的弱比较规则解析请求头：`*` 命中任何已存在的 ETag，标签列表逐项比较， + 比较前统一剥离 `W/` 前缀和外层引号。 + + :param if_none_match: 请求的 If-None-Match 头部，可为 None + :param etag: 本次响应生成的 ETag，可为裸摘要或带引号的标签 + :return: 命中（应返回 304）时为 True + """ + if not if_none_match or not etag: + return False + + def normalize(tag: str) -> str: + tag = tag.strip() + if tag[:2].upper() == "W/": + tag = tag[2:].strip() + if len(tag) >= 2 and tag.startswith('"') and tag.endswith('"'): + tag = tag[1:-1] + return tag + + if if_none_match.strip() == "*": + return True + + current = normalize(etag) + if not current: + return False + return any(normalize(candidate) == current for candidate in if_none_match.split(",")) + @staticmethod def detect_encoding_from_html_response( response: Response, diff --git a/tests/test_login_wallpaper_proxy.py b/tests/test_login_wallpaper_proxy.py new file mode 100644 index 00000000..b005eefa --- /dev/null +++ b/tests/test_login_wallpaper_proxy.py @@ -0,0 +1,409 @@ +import asyncio +import io +import threading +from unittest.mock import AsyncMock, Mock +from urllib.parse import parse_qs, urlparse + +import pytest +from fastapi import HTTPException +from PIL import Image + +from app.api.endpoints import login as login_endpoint +from app.helper import image as image_module +from app.helper.image import ImageHelper + + +def test_wallpapers_preserves_default_external_contract(monkeypatch): + urls = [ + "https://images.example/one.jpg", + "/relative/two.jpg", + "https://images.example/one.jpg", + ] + helper = Mock() + helper.get_wallpapers.return_value = urls + monkeypatch.setattr(login_endpoint, "WallpaperHelper", Mock(return_value=helper)) + + result = login_endpoint.wallpapers(same_origin=False) + + assert result == urls + helper.get_wallpapers.assert_called_once_with() + + +def test_same_origin_wallpapers_preserve_order_duplicates_and_relative_urls( + monkeypatch, +): + urls = [ + "https://images.example/one.jpg", + "/relative/two.jpg", + "https://images.example/one.jpg", + ] + helper = Mock() + helper.get_wallpapers.return_value = urls + monkeypatch.setattr(login_endpoint, "WallpaperHelper", Mock(return_value=helper)) + monkeypatch.setattr(login_endpoint.settings, "WALLPAPER", "customize") + + result = login_endpoint.wallpapers(same_origin=True) + + assert len(result) == len(urls) + assert result[1] == urls[1] + assert result[0] == result[2] + signed_source = parse_qs(urlparse(result[0]).query)["url"][0] + assert ( + login_endpoint.SecurityUtils.verify_signed_url( + signed_source, + purpose=login_endpoint._LOGIN_WALLPAPER_PUBLIC_PURPOSE, + ) + == urls[0] + ) + + +def test_same_origin_wallpapers_do_not_truncate_custom_sources(monkeypatch): + urls = [f"https://images.example/{index}.jpg" for index in range(12)] + helper = Mock() + helper.get_wallpapers.return_value = urls + monkeypatch.setattr(login_endpoint, "WallpaperHelper", Mock(return_value=helper)) + monkeypatch.setattr(login_endpoint.settings, "WALLPAPER", "customize") + + result = login_endpoint.wallpapers(same_origin=True) + + assert len(result) == len(urls) + decoded_sources = [ + login_endpoint.SecurityUtils.verify_signed_url( + parse_qs(urlparse(item).query)["url"][0], + purpose=login_endpoint._LOGIN_WALLPAPER_PUBLIC_PURPOSE, + ) + for item in result + ] + assert decoded_sources == urls + + +def test_same_origin_wallpapers_convert_network_path_references(monkeypatch): + """`//host/path` 会被浏览器解析成跨源地址，同源模式必须一并代理。""" + urls = ["//cdn.example/wallpaper.jpg", "/relative/two.jpg"] + helper = Mock() + helper.get_wallpapers.return_value = urls + monkeypatch.setattr(login_endpoint, "WallpaperHelper", Mock(return_value=helper)) + monkeypatch.setattr(login_endpoint.settings, "WALLPAPER", "customize") + + result = login_endpoint.wallpapers(same_origin=True) + + assert result[1] == urls[1] + assert urlparse(result[0]).netloc == "" + signed_source = parse_qs(urlparse(result[0]).query)["url"][0] + assert ( + login_endpoint.SecurityUtils.verify_signed_url( + signed_source, + purpose=login_endpoint._LOGIN_WALLPAPER_PUBLIC_PURPOSE, + ) + == "https://cdn.example/wallpaper.jpg" + ) + + +def test_url_origin_rejects_malformed_port(): + assert login_endpoint._url_origin("https://images.example:invalid/one.jpg") is None + + +@pytest.mark.asyncio +async def test_wallpaper_image_uses_signed_public_source_without_credentials( + monkeypatch, +): + source_url = "https://images.example/one.jpg" + signed_url = login_endpoint.SecurityUtils.sign_url( + source_url, + purpose=login_endpoint._LOGIN_WALLPAPER_PUBLIC_PURPOSE, + ) + image_helper = Mock() + image_helper.async_fetch_image_guarded = AsyncMock(return_value=b"image-bytes") + safety_check = AsyncMock(return_value=True) + monkeypatch.setattr(login_endpoint, "ImageHelper", Mock(return_value=image_helper)) + monkeypatch.setattr( + login_endpoint.SecurityUtils, + "is_safe_image_url_async", + safety_check, + ) + + response = await login_endpoint.wallpaper_image(signed_url) + + assert response.status_code == 200 + assert response.body == b"image-bytes" + assert response.headers["cache-control"] == "public, max-age=86400" + fetch_options = image_helper.async_fetch_image_guarded.await_args.kwargs + assert fetch_options["url"] == source_url + assert fetch_options["use_cache"] is True + assert fetch_options["max_bytes"] == 32 * 1024 * 1024 + assert ( + fetch_options["redirect_policy"] + == login_endpoint._LOGIN_WALLPAPER_PUBLIC_PURPOSE + ) + assert "cookies" not in fetch_options + assert "max_pixels" not in fetch_options + assert await fetch_options["redirect_validator"]( + "https://cdn.example/redirected.jpg" + ) + assert safety_check.await_count == 2 + + +@pytest.mark.asyncio +async def test_wallpaper_image_serves_etag_and_conditional_requests(monkeypatch): + """条件请求命中时只返回 304，标准引号/弱标签/标签列表/`*` 都要匹配。""" + source_url = "https://images.example/one.jpg" + signed_url = login_endpoint.SecurityUtils.sign_url( + source_url, + purpose=login_endpoint._LOGIN_WALLPAPER_PUBLIC_PURPOSE, + ) + image_helper = Mock() + image_helper.async_fetch_image_guarded = AsyncMock(return_value=b"image-bytes") + monkeypatch.setattr(login_endpoint, "ImageHelper", Mock(return_value=image_helper)) + monkeypatch.setattr( + login_endpoint.SecurityUtils, + "is_safe_image_url_async", + AsyncMock(return_value=True), + ) + + response = await login_endpoint.wallpaper_image(signed_url) + etag = response.headers["etag"] + assert etag == f'"{login_endpoint.HashUtils.md5(b"image-bytes")}"' + + for candidate in (etag, f"W/{etag}", f'"other", {etag}', "*"): + conditional = await login_endpoint.wallpaper_image( + signed_url, if_none_match=candidate + ) + assert conditional.status_code == 304 + assert conditional.headers["etag"] == etag + assert not conditional.body + + stale = await login_endpoint.wallpaper_image( + signed_url, if_none_match='"stale-digest"' + ) + assert stale.status_code == 200 + assert stale.body == b"image-bytes" + + +@pytest.mark.asyncio +async def test_wallpaper_image_rejects_modified_or_unsafe_public_source(monkeypatch): + source_url = "https://images.example/one.jpg" + signed_url = login_endpoint.SecurityUtils.sign_url( + source_url, + purpose=login_endpoint._LOGIN_WALLPAPER_PUBLIC_PURPOSE, + ) + + with pytest.raises(HTTPException) as invalid_signature: + await login_endpoint.wallpaper_image(f"{signed_url}modified") + assert invalid_signature.value.status_code == 404 + + monkeypatch.setattr( + login_endpoint.SecurityUtils, + "is_safe_image_url_async", + AsyncMock(return_value=False), + ) + with pytest.raises(HTTPException) as unsafe_source: + await login_endpoint.wallpaper_image(signed_url) + assert unsafe_source.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_media_wallpaper_only_inherits_private_authority_on_same_origin( + monkeypatch, +): + source_url = "http://mediaserver.local:8096/image.jpg" + signed_url = login_endpoint.SecurityUtils.sign_url( + source_url, + purpose=login_endpoint._LOGIN_WALLPAPER_MEDIA_PURPOSE, + ) + image_helper = Mock() + image_helper.async_fetch_image_guarded = AsyncMock(return_value=b"image-bytes") + safety_check = AsyncMock(return_value=True) + monkeypatch.setattr(login_endpoint, "ImageHelper", Mock(return_value=image_helper)) + monkeypatch.setattr( + login_endpoint.SecurityUtils, + "is_safe_image_url_async", + safety_check, + ) + + await login_endpoint.wallpaper_image(signed_url) + + redirect_validator = ( + image_helper.async_fetch_image_guarded.await_args.kwargs["redirect_validator"] + ) + assert await redirect_validator( + "http://mediaserver.local:8096/redirected.jpg" + ) + assert safety_check.await_count == 0 + assert await redirect_validator("https://cdn.example/redirected.jpg") + assert safety_check.await_count == 1 + # 私网授权范围写入合并键，公共来源的并发请求不会共享这次抓取 + fetch_options = image_helper.async_fetch_image_guarded.await_args.kwargs + assert fetch_options["redirect_policy"] == ( + f"{login_endpoint._LOGIN_WALLPAPER_MEDIA_PURPOSE}" + ":http://mediaserver.local:8096" + ) + + +class _StreamResponse: + """提供受限图片抓取测试所需的最小流式响应合同。""" + + def __init__(self, status_code, *, headers=None, chunks=None, gate=None): + self.status_code = status_code + self.headers = headers or {} + self._chunks = chunks or [] + self._gate = gate + + async def aiter_bytes(self): + if self._gate: + await self._gate.wait() + for chunk in self._chunks: + yield chunk + + +class _StreamContext: + """模拟 AsyncRequestUtils.get_stream 返回的异步上下文。""" + + def __init__(self, response): + self._response = response + + async def __aenter__(self): + return self._response + + async def __aexit__(self, *_): + return False + + +def _guarded_image_helper(): + helper = object.__new__(ImageHelper) + helper.async_file_cache = Mock() + helper.async_file_cache.get = AsyncMock(return_value=None) + helper.async_file_cache.delete = AsyncMock() + helper.async_file_cache.set = AsyncMock() + helper._guarded_fetch_tasks = {} + helper._guarded_fetch_tasks_lock = threading.Lock() + return helper + + +def _png_bytes(): + content = io.BytesIO() + Image.new("RGB", (1, 1), color="black").save(content, format="PNG") + return content.getvalue() + + +@pytest.mark.asyncio +async def test_guarded_image_fetch_rejects_redirect_before_second_request( + monkeypatch, +): + helper = _guarded_image_helper() + request = Mock() + request.get_stream.return_value = _StreamContext( + _StreamResponse(302, headers={"location": "http://127.0.0.1/private.jpg"}) + ) + request_factory = Mock(return_value=request) + redirect_validator = AsyncMock(return_value=False) + monkeypatch.setattr(image_module, "AsyncRequestUtils", request_factory) + + content = await helper.async_fetch_image_guarded( + "https://images.example/one.jpg", + redirect_validator=redirect_validator, + redirect_policy="public", + max_bytes=1024, + use_cache=False, + ) + + assert content is None + redirect_validator.assert_awaited_once_with("http://127.0.0.1/private.jpg") + request_factory.assert_called_once() + + +@pytest.mark.asyncio +async def test_guarded_image_fetch_stops_chunked_response_over_byte_limit( + monkeypatch, +): + helper = _guarded_image_helper() + request = Mock() + request.get_stream.return_value = _StreamContext( + _StreamResponse(200, chunks=[b"1234", b"5678"]) + ) + monkeypatch.setattr(image_module, "AsyncRequestUtils", Mock(return_value=request)) + + content = await helper.async_fetch_image_guarded( + "https://images.example/one.jpg", + redirect_validator=AsyncMock(return_value=True), + redirect_policy="public", + max_bytes=6, + use_cache=False, + ) + + assert content is None + + +@pytest.mark.asyncio +async def test_guarded_image_fetch_coalesces_same_cache_key(monkeypatch): + helper = _guarded_image_helper() + gate = asyncio.Event() + request = Mock() + request.get_stream.return_value = _StreamContext( + _StreamResponse(200, chunks=[_png_bytes()], gate=gate) + ) + request_factory = Mock(return_value=request) + monkeypatch.setattr(image_module, "AsyncRequestUtils", request_factory) + + first = asyncio.create_task( + helper.async_fetch_image_guarded( + "https://images.example/one.png", + redirect_validator=AsyncMock(return_value=True), + redirect_policy="public", + max_bytes=1024, + use_cache=True, + ) + ) + second = asyncio.create_task( + helper.async_fetch_image_guarded( + "https://images.example/one.png", + redirect_validator=AsyncMock(return_value=True), + redirect_policy="public", + max_bytes=1024, + use_cache=True, + ) + ) + await asyncio.sleep(0) + gate.set() + + first_result, second_result = await asyncio.gather(first, second) + + assert first_result == second_result == _png_bytes() + assert request_factory.call_count == 1 + helper.async_file_cache.set.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_guarded_image_fetch_does_not_share_across_fetch_policies(monkeypatch): + """字节上限或重定向策略不同的调用必须各自抓取，不能继承他人的限制。""" + helper = _guarded_image_helper() + gate = asyncio.Event() + request = Mock() + request.get_stream.return_value = _StreamContext( + _StreamResponse(200, chunks=[_png_bytes()], gate=gate) + ) + request_factory = Mock(return_value=request) + monkeypatch.setattr(image_module, "AsyncRequestUtils", request_factory) + + def fetch(*, max_bytes, redirect_policy): + return asyncio.create_task( + helper.async_fetch_image_guarded( + "https://images.example/one.png", + redirect_validator=AsyncMock(return_value=True), + redirect_policy=redirect_policy, + max_bytes=max_bytes, + use_cache=True, + ) + ) + + larger_limit = fetch(max_bytes=4096, redirect_policy="public") + smaller_limit = fetch(max_bytes=8, redirect_policy="public") + other_policy = fetch(max_bytes=4096, redirect_policy="media:http://host:8096") + await asyncio.sleep(0) + gate.set() + + results = await asyncio.gather(larger_limit, smaller_limit, other_policy) + + assert request_factory.call_count == 3 + # 8 字节上限的调用不能收到超出自身上限的共享结果 + assert results[1] is None + assert results[0] == results[2] == _png_bytes()