From a1bc773fb522ef7700522204f5474e9e0f72e3f6 Mon Sep 17 00:00:00 2001
From: InfinityPacer <160988576+InfinityPacer@users.noreply.github.com>
Date: Tue, 11 Feb 2025 17:10:50 +0800
Subject: [PATCH] feat(security): add AVIF support

---
 app/api/endpoints/system.py | 4 +++-
 app/chain/recommend.py      | 1 +
 app/core/config.py          | 2 +-
 app/utils/security.py       | 2 +-
 requirements.in             | 1 +
 5 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/app/api/endpoints/system.py b/app/api/endpoints/system.py
index 954261bd..376089b8 100644
--- a/app/api/endpoints/system.py
+++ b/app/api/endpoints/system.py
@@ -8,6 +8,7 @@ from pathlib import Path
 from typing import Optional, Union
 
 import aiofiles
+import pillow_avif  # noqa 用于自动注册AVIF支持
 from PIL import Image
 from fastapi import APIRouter, Depends, HTTPException, Header, Request, Response
 from fastapi.responses import StreamingResponse
@@ -91,7 +92,8 @@ def fetch_image(
     # 请求远程图片
     referer = "https://movie.douban.com/" if "doubanio.com" in url else None
     proxies = settings.PROXY if proxy else None
-    response = RequestUtils(ua=settings.USER_AGENT, proxies=proxies, referer=referer).get_res(url=url)
+    response = RequestUtils(ua=settings.USER_AGENT, proxies=proxies, referer=referer,
+                            accept_type="image/avif,image/webp,image/apng,*/*").get_res(url=url)
     if not response:
         raise HTTPException(status_code=502, detail="Failed to fetch the image from the remote server")
 
diff --git a/app/chain/recommend.py b/app/chain/recommend.py
index c8d877af..d461045c 100644
--- a/app/chain/recommend.py
+++ b/app/chain/recommend.py
@@ -3,6 +3,7 @@ import tempfile
 from pathlib import Path
 from typing import List
 
+import pillow_avif  # noqa 用于自动注册AVIF支持
 from PIL import Image
 
 from app.chain import ChainBase
diff --git a/app/core/config.py b/app/core/config.py
index a4acf5c1..41ba137d 100644
--- a/app/core/config.py
+++ b/app/core/config.py
@@ -247,7 +247,7 @@ class ConfigModel(BaseModel):
     )
     # 允许的图片文件后缀格式
     SECURITY_IMAGE_SUFFIXES: List[str] = Field(
-        default_factory=lambda: [".jpg", ".jpeg", ".png", ".webp", ".gif", ".svg"]
+        default_factory=lambda: [".jpg", ".jpeg", ".png", ".webp", ".gif", ".svg", ".avif"]
     )
     # 重命名时支持的S0别名
     RENAME_FORMAT_S0_NAMES: List[str] = Field(
diff --git a/app/utils/security.py b/app/utils/security.py
index f7367915..8190d928 100644
--- a/app/utils/security.py
+++ b/app/utils/security.py
@@ -42,7 +42,7 @@ class SecurityUtils:
     @staticmethod
     def is_safe_url(url: str, allowed_domains: Union[Set[str], List[str]], strict: bool = False) -> bool:
         """
-        验证URL是否在允许的域名列表中，包括带有端口的域名。
+        验证URL是否在允许的域名列表中，包括带有端口的域名
 
         :param url: 需要验证的 URL
         :param allowed_domains: 允许的域名集合，域名可以包含端口
diff --git a/requirements.in b/requirements.in
index 85aaa5bd..0a0a5ef6 100644
--- a/requirements.in
+++ b/requirements.in
@@ -32,6 +32,7 @@ func_timeout==4.3.5
 bs4~=0.0.1
 beautifulsoup4~=4.12.2
 pillow~=10.4.0
+pillow-avif-plugin~=1.4.6
 pyTelegramBotAPI~=4.12.0
 playwright~=1.37.0
 cf-clearance~=0.31.0