From bcb02c1c9069ad1fffdde4d1ea16fb4dca660b21 Mon Sep 17 00:00:00 2001 From: niqzart Date: Mon, 16 Feb 2026 03:18:34 +0300 Subject: [PATCH 1/4] fix: use aiofiles for saving files --- app/storage_v2/models/files_db.py | 10 +++--- app/storage_v2/routers/files_rst.py | 2 +- app/users/routes/avatar_rst.py | 5 +-- poetry.lock | 49 ++++++++++++++++++----------- pyproject.toml | 6 ++-- 5 files changed, 43 insertions(+), 29 deletions(-) diff --git a/app/storage_v2/models/files_db.py b/app/storage_v2/models/files_db.py index be7ff471..fd2cea91 100644 --- a/app/storage_v2/models/files_db.py +++ b/app/storage_v2/models/files_db.py @@ -1,9 +1,9 @@ from enum import StrEnum from pathlib import Path -from shutil import copyfileobj -from typing import BinaryIO, Literal, Self +from typing import Literal, Self from uuid import UUID, uuid4 +import aiofiles from pydantic_marshals.sqlalchemy import MappedModel from sqlalchemy import Enum from sqlalchemy.orm import Mapped, mapped_column @@ -57,7 +57,7 @@ def content_disposition(self) -> ContentDisposition: @classmethod async def create_with_content( cls, - content: BinaryIO, + content: bytes, filename: str | None, file_kind: FileKind, ) -> Self: @@ -65,8 +65,8 @@ async def create_with_content( name=filename or "upload", kind=file_kind, ) - with file.path.open("wb") as f: - copyfileobj(content, f) # TODO maybe convert to async + async with aiofiles.open(file.path, "wb") as f: + await f.write(content) return file async def delete(self) -> None: diff --git a/app/storage_v2/routers/files_rst.py b/app/storage_v2/routers/files_rst.py index 48b1db61..2752f367 100644 --- a/app/storage_v2/routers/files_rst.py +++ b/app/storage_v2/routers/files_rst.py @@ -36,7 +36,7 @@ async def upload_file( raise StorageTokenResponses.INVALID_STORAGE_TOKEN file = await File.create_with_content( - content=upload.file, + content=await upload.read(), filename=upload.filename, file_kind=file_kind, ) diff --git a/app/users/routes/avatar_rst.py b/app/users/routes/avatar_rst.py index f1b12f7d..955a04b0 100644 --- a/app/users/routes/avatar_rst.py +++ b/app/users/routes/avatar_rst.py @@ -1,5 +1,6 @@ from typing import Annotated +import aiofiles import filetype # type: ignore[import-untyped] from fastapi import File, UploadFile from filetype.types.image import Webp # type: ignore[import-untyped] @@ -28,8 +29,8 @@ async def update_or_create_avatar( if not filetype.match(avatar.file, [Webp()]): raise AvatarResponses.WRONG_FORMAT - with user.avatar_path.open("wb") as file: - file.write(await avatar.read()) + async with aiofiles.open(user.avatar_path, "wb") as file: + await file.write(await avatar.read()) @router.delete( diff --git a/poetry.lock b/poetry.lock index ab0b99b0..38d962f6 100644 --- a/poetry.lock +++ b/poetry.lock @@ -2,46 +2,45 @@ [[package]] name = "aiofiles" -version = "24.1.0" +version = "25.1.0" description = "File support for asyncio." optional = false -python-versions = ">=3.8" +python-versions = ">=3.9" groups = ["main"] files = [ - {file = "aiofiles-24.1.0-py3-none-any.whl", hash = "sha256:b4ec55f4195e3eb5d7abd1bf7e061763e864dd4954231fb8539a0ef8bb8260e5"}, - {file = "aiofiles-24.1.0.tar.gz", hash = "sha256:22a075c9e5a3810f0c2e48f3008c94d68c65d763b9b03857924c99e57355166c"}, + {file = "aiofiles-25.1.0-py3-none-any.whl", hash = "sha256:abe311e527c862958650f9438e859c1fa7568a141b22abcd015e120e86a85695"}, + {file = "aiofiles-25.1.0.tar.gz", hash = "sha256:a8d728f0a29de45dc521f18f07297428d56992a742f0cd2701ba86e44d23d5b2"}, ] [[package]] name = "aiogram" -version = "3.20.0.post0" +version = "3.25.0" description = "Modern and fully asynchronous framework for Telegram Bot API" optional = false -python-versions = ">=3.9" +python-versions = "<3.15,>=3.10" groups = ["main"] files = [ - {file = "aiogram-3.20.0.post0-py3-none-any.whl", hash = "sha256:c8f5a68b0729e74efa15a7fc285bd49fa3d0603de5e424404219a822f8e5f4d1"}, - {file = "aiogram-3.20.0.post0.tar.gz", hash = "sha256:2443799b4514ac251fcf2d561603e250f8763808542222fa1136901058eda9a3"}, + {file = "aiogram-3.25.0-py3-none-any.whl", hash = "sha256:0243966e93fbde14e90c0dfd0b3776c637ebf7ddcca2c7ee81ecbd68d9490cce"}, + {file = "aiogram-3.25.0.tar.gz", hash = "sha256:8a8b0c34f8c4ca8a6501b954abb0eeba26743449e35e20b70c0d810347354c3c"}, ] [package.dependencies] -aiofiles = ">=23.2.1,<24.2" -aiohttp = ">=3.9.0,<3.12" +aiofiles = ">=23.2.1,<26.0" +aiohttp = ">=3.9.0,<3.14" certifi = ">=2023.7.22" magic-filter = ">=1.0.12,<1.1" -pydantic = ">=2.4.1,<2.12" +pydantic = ">=2.4.1,<2.13" typing-extensions = ">=4.7.0,<=5.0" [package.extras] cli = ["aiogram-cli (>=1.1.0,<2.0.0)"] -dev = ["black (>=24.4.2,<24.5.0)", "isort (>=5.13.2,<5.14.0)", "motor-types (>=1.0.0b4,<1.1.0)", "mypy (>=1.10.0,<1.11.0)", "packaging (>=24.1,<25.0)", "pre-commit (>=3.5,<4.0)", "ruff (>=0.5.1,<0.6.0)", "toml (>=0.10.2,<0.11.0)"] docs = ["furo (>=2024.8.6,<2024.9.0)", "markdown-include (>=0.8.1,<0.9.0)", "pygments (>=2.18.0,<2.19.0)", "pymdown-extensions (>=10.3,<11.0)", "sphinx (>=8.0.2,<8.1.0)", "sphinx-autobuild (>=2024.9.3,<2024.10.0)", "sphinx-copybutton (>=0.5.2,<0.6.0)", "sphinx-intl (>=2.2.0,<2.3.0)", "sphinx-substitution-extensions (>=2024.8.6,<2024.9.0)", "sphinxcontrib-towncrier (>=0.4.0a0,<0.5.0)", "towncrier (>=24.8.0,<24.9.0)"] fast = ["aiodns (>=3.0.0)", "uvloop (>=0.17.0) ; (sys_platform == \"darwin\" or sys_platform == \"linux\") and platform_python_implementation != \"PyPy\" and python_version < \"3.13\"", "uvloop (>=0.21.0) ; (sys_platform == \"darwin\" or sys_platform == \"linux\") and platform_python_implementation != \"PyPy\" and python_version >= \"3.13\""] -i18n = ["babel (>=2.13.0,<2.14.0)"] -mongo = ["motor (>=3.3.2,<3.7.0)"] -proxy = ["aiohttp-socks (>=0.8.3,<0.9.0)"] -redis = ["redis[hiredis] (>=5.0.1,<5.3.0)"] -test = ["aresponses (>=2.1.6,<2.2.0)", "pycryptodomex (>=3.19.0,<3.20.0)", "pytest (>=7.4.2,<7.5.0)", "pytest-aiohttp (>=1.0.5,<1.1.0)", "pytest-asyncio (>=0.21.1,<0.22.0)", "pytest-cov (>=4.1.0,<4.2.0)", "pytest-html (>=4.0.2,<4.1.0)", "pytest-lazy-fixture (>=0.6.3,<0.7.0)", "pytest-mock (>=3.12.0,<3.13.0)", "pytest-mypy (>=0.10.3,<0.11.0)", "pytz (>=2023.3,<2024.0)"] +i18n = ["babel (>=2.13.0,<3)"] +mongo = ["motor (>=3.3.2,<3.8)", "pymongo (>4.5,<4.16)"] +proxy = ["aiohttp-socks (>=0.10.1,<0.11.0)"] +redis = ["redis[hiredis] (>=6.2.0,<8)"] +signature = ["cryptography (>=46.0.0)"] [[package]] name = "aiohappyeyeballs" @@ -3700,6 +3699,18 @@ asgiref = ">=3.8.1,<4.0.0" pydantic = ">=2.7.0,<3.0.0" python-socketio = ">=5.11.2,<6.0.0" +[[package]] +name = "types-aiofiles" +version = "25.1.0.20251011" +description = "Typing stubs for aiofiles" +optional = false +python-versions = ">=3.9" +groups = ["types"] +files = [ + {file = "types_aiofiles-25.1.0.20251011-py3-none-any.whl", hash = "sha256:8ff8de7f9d42739d8f0dadcceeb781ce27cd8d8c4152d4a7c52f6b20edb8149c"}, + {file = "types_aiofiles-25.1.0.20251011.tar.gz", hash = "sha256:1c2b8ab260cb3cd40c15f9d10efdc05a6e1e6b02899304d80dfa0410e028d3ff"}, +] + [[package]] name = "types-cffi" version = "1.17.0.20250326" @@ -4255,5 +4266,5 @@ propcache = ">=0.2.0" [metadata] lock-version = "2.1" -python-versions = "~=3.12,<4.0" -content-hash = "ce2caa8369ff8b651eff2569fa2f3a03b009d4c2bec31e3a5ead53b778af01d1" +python-versions = "~=3.12,<3.15" +content-hash = "c4470c4c66ef2a9814680cfd1a0c66728c7b443cdaee5f109b7241660bff6fb5" diff --git a/pyproject.toml b/pyproject.toml index 00ec9785..72176715 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ authors = [ {name = "xi.team", email = "xieffect@yandex.ru"}, ] readme = "README.md" -requires-python = "~=3.12,<4.0" +requires-python = "~=3.12,<3.15" dynamic = ["dependencies"] [build-system] @@ -34,8 +34,9 @@ livekit-api = "1.0.5" passlib = "^1.7.4" cryptography = "^42.0.5" discord-webhook = {extras = ["async"], version = "^1.3.1"} -aiogram = "^3.4.1" +aiogram = "^3.25.0" aiosmtplib = "^3.0.2" +aiofiles = "^25.1.0" itsdangerous = "^2.2.0" faststream = {extras = ["redis"], version = "^0.6.2"} sentry-sdk = {extras = ["asyncio", "fastapi", "sqlalchemy", "redis", "httpx"], version = "2.44.0"} @@ -43,6 +44,7 @@ sentry-sdk = {extras = ["asyncio", "fastapi", "sqlalchemy", "redis", "httpx"], v [tool.poetry.group.types.dependencies] types-passlib = "^1.7.7.13" types-redis = "^4.6.0" +types-aiofiles = "^25.1.0.20251011" [tool.poetry.group.dev.dependencies] watchfiles = "^0.21.0" From d4a283fa9fed6a6a077f25f5fa8b1eea42fa31fb Mon Sep 17 00:00:00 2001 From: niqzart Date: Mon, 16 Feb 2026 03:45:33 +0300 Subject: [PATCH 2/4] feat: allow more image formats and add conversion & resizing for avatars --- .flake8 | 2 +- app/common/filetype_ext.py | 26 ++++++++++++ app/users/models/users_db.py | 4 +- app/users/routes/avatar_rst.py | 26 ++++++++---- poetry.lock | 4 +- pyproject.toml | 2 +- tests/common/faker_ext.py | 2 + tests/users/functional/test_avatars.py | 58 +++++++++++++++++++------- 8 files changed, 97 insertions(+), 27 deletions(-) create mode 100644 app/common/filetype_ext.py diff --git a/.flake8 b/.flake8 index 6ab4761e..7b9f818b 100644 --- a/.flake8 +++ b/.flake8 @@ -26,7 +26,7 @@ extend-ignore = VNE003 WPS115 # # weird - PIE803 C101 FNE007 FNE008 N812 ANN101 ANN102 WPS110 WPS111 WPS114 WPS338 WPS407 WPS414 WPS440 VNE001 VNE002 CM001 + PIE803 C101 FNE007 FNE008 N812 ANN101 ANN102 PT004 WPS110 WPS111 WPS114 WPS338 WPS407 WPS414 WPS440 VNE001 VNE002 CM001 # too many WPS200 WPS201 WPS202 WPS203 WPS204 WPS210 WPS211 WPS212 WPS213 WPS214 WPS217 WPS218 WPS221 WPS224 WPS230 WPS231 WPS234 WPS235 WPS238 # "vague" imports diff --git a/app/common/filetype_ext.py b/app/common/filetype_ext.py new file mode 100644 index 00000000..0cc68297 --- /dev/null +++ b/app/common/filetype_ext.py @@ -0,0 +1,26 @@ +from typing import Final + +import filetype # type: ignore[import-untyped] +from filetype.types import image # type: ignore[import-untyped] + +FILE_HEADER_SIZE: Final[int] = 8192 + +SUPPORTED_IMAGE_FORMATS: list[filetype.Type] = [ + image.Avif(), + image.Bmp(), + image.Gif(), + image.Ico(), + image.Jpeg(), + image.Jpx(), + image.Png(), + image.Tiff(), + image.Webp(), +] + + +def match_filetype(obj: bytes, matchers: list[filetype.Type]) -> filetype.Type | None: + return filetype.match(obj, matchers) + + +def match_image_filetype(obj: bytes) -> filetype.Type | None: + return match_filetype(obj, SUPPORTED_IMAGE_FORMATS) diff --git a/app/users/models/users_db.py b/app/users/models/users_db.py index 29349369..9a29a854 100644 --- a/app/users/models/users_db.py +++ b/app/users/models/users_db.py @@ -2,7 +2,7 @@ from datetime import datetime, timedelta from enum import StrEnum from pathlib import Path -from typing import Annotated, Self +from typing import Annotated, ClassVar, Self from passlib.handlers.pbkdf2 import pbkdf2_sha256 from pydantic import AfterValidator, AwareDatetime, StringConstraints @@ -30,6 +30,8 @@ class OnboardingStage(StrEnum): class User(Base): __tablename__ = "users" + avatar_shape: ClassVar[tuple[int, int]] = 128, 128 + @staticmethod def generate_hash(password: str) -> str: return pbkdf2_sha256.hash(password) diff --git a/app/users/routes/avatar_rst.py b/app/users/routes/avatar_rst.py index 955a04b0..f2a5b2b8 100644 --- a/app/users/routes/avatar_rst.py +++ b/app/users/routes/avatar_rst.py @@ -1,13 +1,14 @@ -from typing import Annotated +from io import BytesIO import aiofiles -import filetype # type: ignore[import-untyped] -from fastapi import File, UploadFile -from filetype.types.image import Webp # type: ignore[import-untyped] +from fastapi import UploadFile +from PIL import Image from starlette import status from app.common.fastapi_ext import APIRouterExt, Responses +from app.common.filetype_ext import FILE_HEADER_SIZE, match_image_filetype from app.users.dependencies.users_dep import AuthorizedUser +from app.users.models.users_db import User router = APIRouterExt(tags=["current user avatar"]) @@ -24,13 +25,24 @@ class AvatarResponses(Responses): ) async def update_or_create_avatar( user: AuthorizedUser, - avatar: Annotated[UploadFile, File(description="image/webp")], + avatar: UploadFile, ) -> None: - if not filetype.match(avatar.file, [Webp()]): + avatar_header_data = await avatar.read(FILE_HEADER_SIZE) + + if match_image_filetype(avatar_header_data) is None: raise AvatarResponses.WRONG_FORMAT + await avatar.seek(0) + avatar_image: Image.Image = Image.open(BytesIO(await avatar.read())) + + avatar_image = avatar_image.resize(User.avatar_shape) + + processed_avatar = BytesIO() + avatar_image.save(processed_avatar, format="webp") + processed_avatar.seek(0) + async with aiofiles.open(user.avatar_path, "wb") as file: - await file.write(await avatar.read()) + await file.write(processed_avatar.read()) @router.delete( diff --git a/poetry.lock b/poetry.lock index 38d962f6..a9b6e4c6 100644 --- a/poetry.lock +++ b/poetry.lock @@ -2448,7 +2448,7 @@ version = "11.2.1" description = "Python Imaging Library (Fork)" optional = false python-versions = ">=3.9" -groups = ["tests"] +groups = ["main"] files = [ {file = "pillow-11.2.1-cp310-cp310-macosx_10_10_x86_64.whl", hash = "sha256:d57a75d53922fc20c165016a20d9c44f73305e67c351bbc60d1adaf662e74047"}, {file = "pillow-11.2.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:127bf6ac4a5b58b3d32fc8289656f77f80567d65660bc46f72c0d77e6600cc95"}, @@ -4267,4 +4267,4 @@ propcache = ">=0.2.0" [metadata] lock-version = "2.1" python-versions = "~=3.12,<3.15" -content-hash = "c4470c4c66ef2a9814680cfd1a0c66728c7b443cdaee5f109b7241660bff6fb5" +content-hash = "f4ed3dd8ca073154f866b1ffca749b0f7aae54fdd63570cc4ec2694978b0137a" diff --git a/pyproject.toml b/pyproject.toml index 72176715..3aba2ec9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,6 +40,7 @@ aiofiles = "^25.1.0" itsdangerous = "^2.2.0" faststream = {extras = ["redis"], version = "^0.6.2"} sentry-sdk = {extras = ["asyncio", "fastapi", "sqlalchemy", "redis", "httpx"], version = "2.44.0"} +pillow = "^11.2.1" [tool.poetry.group.types.dependencies] types-passlib = "^1.7.7.13" @@ -88,7 +89,6 @@ polyfactory = "^2.21.0" faker = "^37.1.0" faker-file = "^0.18.4" rstr = "^3.2.2" -pillow = "^11.2.1" [tool.isort] profile = "black" diff --git a/tests/common/faker_ext.py b/tests/common/faker_ext.py index 0b7cd6c0..b1b9a5e0 100644 --- a/tests/common/faker_ext.py +++ b/tests/common/faker_ext.py @@ -4,6 +4,7 @@ from faker_file.providers import ( # type: ignore[import-untyped] bin_file, pdf_file, + png_file, webp_file, ) @@ -13,6 +14,7 @@ def _setup_faker(faker: Faker) -> None: faker.add_provider(internet) faker.add_provider(bin_file.BinFileProvider) faker.add_provider(webp_file.GraphicWebpFileProvider) + faker.add_provider(png_file.GraphicPngFileProvider) faker.add_provider(pdf_file.PdfFileProvider) diff --git a/tests/users/functional/test_avatars.py b/tests/users/functional/test_avatars.py index 8fd3247b..e37d7444 100644 --- a/tests/users/functional/test_avatars.py +++ b/tests/users/functional/test_avatars.py @@ -1,48 +1,73 @@ from collections.abc import AsyncIterator +from io import BytesIO import pytest from faker import Faker +from PIL import Image +from pytest_lazy_fixtures import lfc from starlette import status from starlette.testclient import TestClient from app.users.models.users_db import User from tests.common.assert_contains_ext import assert_nodata_response, assert_response +from tests.common.types import PytestRequest pytestmark = pytest.mark.anyio +@pytest.fixture( + params=[ + pytest.param(lfc(lambda faker: faker.graphic_webp_file(raw=True)), id="webp"), + pytest.param(lfc(lambda faker: faker.graphic_png_file(raw=True)), id="png"), + ] +) +def image_content(request: PytestRequest[bytes]) -> bytes: + return request.param + + @pytest.fixture() -async def image(faker: Faker) -> bytes: - return faker.graphic_webp_file(raw=True) # type: ignore[no-any-return] +def processed_image_content(image_content: bytes) -> bytes: + image = Image.open(BytesIO(image_content)) + image.resize(User.avatar_shape) + + processed_image_buffer = BytesIO() + image.save(processed_image_buffer, format="webp") + + processed_image_buffer.seek(0) + return processed_image_buffer.read() @pytest.fixture() -async def _create_avatar(user: User, image: bytes) -> AsyncIterator[None]: +async def create_avatar(faker: Faker, user: User) -> AsyncIterator[None]: with user.avatar_path.open("wb") as f: - f.write(image) + f.write(faker.graphic_webp_file(raw=True)) yield user.avatar_path.unlink(missing_ok=True) async def test_avatar_uploading( - authorized_client: TestClient, user: User, image: bytes + authorized_client: TestClient, + user: User, + image_content: bytes, + processed_image_content: bytes, ) -> None: assert_nodata_response( authorized_client.put( "/api/protected/user-service/users/current/avatar/", - files={"avatar": ("avatar.webp", image, "image/webp")}, + files={"avatar": ("avatar.webp", image_content, "image/webp")}, ) ) assert user.avatar_path.is_file() with user.avatar_path.open("rb") as f: - assert f.read() == image + assert f.read() == processed_image_content user.avatar_path.unlink() async def test_avatar_uploading_wrong_format( - authorized_client: TestClient, faker: Faker + authorized_client: TestClient, + faker: Faker, ) -> None: assert_response( authorized_client.put( @@ -54,24 +79,27 @@ async def test_avatar_uploading_wrong_format( ) -@pytest.mark.usefixtures("_create_avatar") +@pytest.mark.usefixtures("create_avatar") async def test_avatar_replacing( - authorized_client: TestClient, user: User, faker: Faker + authorized_client: TestClient, + user: User, + faker: Faker, + image_content: bytes, + processed_image_content: bytes, ) -> None: - image_2 = faker.graphic_webp_file(raw=True) assert_nodata_response( authorized_client.put( "/api/protected/user-service/users/current/avatar/", - files={"avatar": ("avatar.webp", image_2, "image/webp")}, + files={"avatar": ("avatar.webp", image_content, "image/webp")}, ) ) assert user.avatar_path.is_file() with user.avatar_path.open("rb") as f: - assert f.read() == image_2 + assert f.read() == processed_image_content -@pytest.mark.usefixtures("_create_avatar") +@pytest.mark.usefixtures("create_avatar") async def test_avatar_deletion(authorized_client: TestClient, user: User) -> None: assert_nodata_response( authorized_client.delete("/api/protected/user-service/users/current/avatar/") @@ -80,7 +108,7 @@ async def test_avatar_deletion(authorized_client: TestClient, user: User) -> Non assert not user.avatar_path.is_file() -@pytest.mark.usefixtures("_create_avatar") +@pytest.mark.usefixtures("create_avatar") async def test_mub_user_deletion_with_avatar( mub_client: TestClient, user: User ) -> None: From b2585d0bdc43861fd0fb0afa47143ff8b51915a0 Mon Sep 17 00:00:00 2001 From: niqzart Date: Mon, 16 Feb 2026 04:25:00 +0300 Subject: [PATCH 3/4] feat: allow more image formats & add conversion in storage-2 --- .../dependencies/storage_token_dep.py | 22 ++++++ app/storage_v2/dependencies/uploads_dep.py | 19 ++++-- app/storage_v2/models/files_db.py | 2 +- app/storage_v2/routers/files_rst.py | 37 +++++----- tests/storage_v2/conftest.py | 51 +++++++++++--- .../functional/test_file_reads_rst.py | 2 +- .../functional/test_file_uploads_rst.py | 68 +++++++++++++++++-- 7 files changed, 164 insertions(+), 37 deletions(-) diff --git a/app/storage_v2/dependencies/storage_token_dep.py b/app/storage_v2/dependencies/storage_token_dep.py index 63858acf..c2fe8865 100644 --- a/app/storage_v2/dependencies/storage_token_dep.py +++ b/app/storage_v2/dependencies/storage_token_dep.py @@ -7,6 +7,7 @@ from app.common.dependencies.authorization_dep import AuthorizationData from app.common.fastapi_ext import Responses, with_responses from app.common.schemas.storage_sch import StorageTokenPayloadSchema +from app.storage_v2.models.access_groups_db import AccessGroup class StorageTokenResponses(Responses): @@ -36,3 +37,24 @@ def validate_and_deserialize_storage_token( StorageTokenPayload = Annotated[ StorageTokenPayloadSchema, Depends(validate_and_deserialize_storage_token) ] + + +@with_responses(StorageTokenResponses) +async def validate_upload_permissions( + storage_token_payload: StorageTokenPayload, +) -> StorageTokenPayloadSchema: + if not storage_token_payload.can_upload_files: + raise StorageTokenResponses.INVALID_STORAGE_TOKEN + + access_group = await AccessGroup.find_first_by_id( + storage_token_payload.access_group_id + ) + if access_group is None: + raise StorageTokenResponses.INVALID_STORAGE_TOKEN + + return storage_token_payload + + +UploadAllowedStorageTokenPayload = Annotated[ + StorageTokenPayloadSchema, Depends(validate_upload_permissions) +] diff --git a/app/storage_v2/dependencies/uploads_dep.py b/app/storage_v2/dependencies/uploads_dep.py index a342b525..a5c2c521 100644 --- a/app/storage_v2/dependencies/uploads_dep.py +++ b/app/storage_v2/dependencies/uploads_dep.py @@ -1,21 +1,32 @@ from typing import Annotated from fastapi import Depends, UploadFile -from filetype import filetype # type: ignore[import-untyped] -from filetype.types.image import Webp # type: ignore[import-untyped] from starlette import status from app.common.fastapi_ext import Responses, with_responses +from app.common.filetype_ext import FILE_HEADER_SIZE, match_image_filetype class FileFormatResponses(Responses): WRONG_FORMAT = status.HTTP_415_UNSUPPORTED_MEDIA_TYPE, "Invalid file format" + CONTENT_TYPE_MISMATCH = ( + status.HTTP_415_UNSUPPORTED_MEDIA_TYPE, + "File content doesn't match the content-type header", + ) @with_responses(FileFormatResponses) -def validate_image_upload(upload: UploadFile) -> UploadFile: - if not filetype.match(upload.file, [Webp()]): +async def validate_image_upload(upload: UploadFile) -> UploadFile: + upload_header_data = await upload.read(FILE_HEADER_SIZE) + image_type = match_image_filetype(upload_header_data) + + if image_type is None: raise FileFormatResponses.WRONG_FORMAT + + if image_type.mime != upload.content_type: + raise FileFormatResponses.CONTENT_TYPE_MISMATCH + + await upload.seek(0) return upload diff --git a/app/storage_v2/models/files_db.py b/app/storage_v2/models/files_db.py index fd2cea91..118abed4 100644 --- a/app/storage_v2/models/files_db.py +++ b/app/storage_v2/models/files_db.py @@ -58,7 +58,7 @@ def content_disposition(self) -> ContentDisposition: async def create_with_content( cls, content: bytes, - filename: str | None, + filename: str, file_kind: FileKind, ) -> Self: file = await cls.create( diff --git a/app/storage_v2/routers/files_rst.py b/app/storage_v2/routers/files_rst.py index 2752f367..0efb2103 100644 --- a/app/storage_v2/routers/files_rst.py +++ b/app/storage_v2/routers/files_rst.py @@ -1,8 +1,10 @@ from datetime import datetime +from io import BytesIO from os import stat from typing import Annotated from fastapi import Header, UploadFile +from PIL import Image from starlette import status from starlette.responses import FileResponse, Response from starlette.staticfiles import NotModifiedResponse @@ -13,9 +15,10 @@ from app.storage_v2.dependencies.storage_token_dep import ( StorageTokenPayload, StorageTokenResponses, + UploadAllowedStorageTokenPayload, ) from app.storage_v2.dependencies.uploads_dep import ValidatedImageUpload -from app.storage_v2.models.access_groups_db import AccessGroup, AccessGroupFile +from app.storage_v2.models.access_groups_db import AccessGroupFile from app.storage_v2.models.files_db import File, FileKind router = APIRouterExt(tags=["files"]) @@ -23,21 +26,13 @@ async def upload_file( storage_token_payload: StorageTokenPayloadSchema, - upload: UploadFile, + upload_content: bytes, + upload_filename: str | None, file_kind: FileKind, ) -> File: - if not storage_token_payload.can_upload_files: - raise StorageTokenResponses.INVALID_STORAGE_TOKEN - - access_group = await AccessGroup.find_first_by_id( - storage_token_payload.access_group_id - ) - if access_group is None: - raise StorageTokenResponses.INVALID_STORAGE_TOKEN - file = await File.create_with_content( - content=await upload.read(), - filename=upload.filename, + content=upload_content, + filename=upload_filename or "upload", file_kind=file_kind, ) @@ -56,12 +51,13 @@ async def upload_file( summary="Upload a new uncategorized file", ) async def upload_uncategorized_file( - storage_token_payload: StorageTokenPayload, + storage_token_payload: UploadAllowedStorageTokenPayload, upload: UploadFile, ) -> File: return await upload_file( storage_token_payload=storage_token_payload, - upload=upload, + upload_content=await upload.read(), + upload_filename=upload.filename, file_kind=FileKind.UNCATEGORIZED, ) @@ -73,12 +69,19 @@ async def upload_uncategorized_file( summary="Upload a new image file", ) async def upload_image_file( - storage_token_payload: StorageTokenPayload, + storage_token_payload: UploadAllowedStorageTokenPayload, upload: ValidatedImageUpload, ) -> File: + image = Image.open(BytesIO(await upload.read())) + processed_image = BytesIO() + image.save(processed_image, format="webp") + processed_image.seek(0) + upload_content = processed_image.read() + return await upload_file( storage_token_payload=storage_token_payload, - upload=upload, + upload_content=upload_content, + upload_filename=upload.filename, file_kind=FileKind.IMAGE, ) diff --git a/tests/storage_v2/conftest.py b/tests/storage_v2/conftest.py index 54c6cd89..9c77791e 100644 --- a/tests/storage_v2/conftest.py +++ b/tests/storage_v2/conftest.py @@ -1,10 +1,12 @@ from dataclasses import dataclass +from io import BytesIO from os import stat from typing import Any, Protocol from uuid import UUID, uuid4 import pytest from faker import Faker +from PIL import Image from pytest_lazy_fixtures import lf from starlette.responses import FileResponse from starlette.testclient import TestClient @@ -104,45 +106,78 @@ def uncategorized_file_content(faker: Faker) -> bytes: return faker.bin_file(raw=True) # type: ignore[no-any-return] +def process_image_content(image_content: bytes) -> bytes: + image = Image.open(BytesIO(image_content)) + processed_image_buffer = BytesIO() + image.save(processed_image_buffer, format="webp") + processed_image_buffer.seek(0) + return processed_image_buffer.read() + + @pytest.fixture() -def image_file_content(faker: Faker) -> bytes: +def webp_image_file_content(faker: Faker) -> bytes: return faker.graphic_webp_file(raw=True) # type: ignore[no-any-return] +@pytest.fixture() +def png_image_file_content(faker: Faker) -> bytes: + return faker.graphic_png_file(raw=True) # type: ignore[no-any-return] + + @dataclass class FileInputData: kind: FileKind name: str - content: bytes content_type: str + input_content: bytes + processed_content: bytes @pytest.fixture() def uncategorized_file_input_data( - faker: Faker, uncategorized_file_content: bytes + faker: Faker, + uncategorized_file_content: bytes, ) -> FileInputData: return FileInputData( kind=FileKind.UNCATEGORIZED, name=faker.file_name(), - content=uncategorized_file_content, + input_content=uncategorized_file_content, + processed_content=uncategorized_file_content, content_type=faker.mime_type(), ) @pytest.fixture() -def image_file_input_data(faker: Faker, image_file_content: bytes) -> FileInputData: +def webp_image_file_input_data( + faker: Faker, webp_image_file_content: bytes +) -> FileInputData: return FileInputData( kind=FileKind.IMAGE, name=faker.file_name(extension="webp"), - content=image_file_content, + input_content=webp_image_file_content, + processed_content=process_image_content(webp_image_file_content), content_type="image/webp", ) +@pytest.fixture() +def png_image_file_input_data( + faker: Faker, png_image_file_content: bytes +) -> FileInputData: + return FileInputData( + kind=FileKind.IMAGE, + name=faker.file_name(extension="png"), + input_content=png_image_file_content, + processed_content=process_image_content(png_image_file_content), + content_type="image/png", + ) + + @pytest.fixture( params=[ pytest.param(lf("uncategorized_file_input_data"), id="uncategorized"), - pytest.param(lf("image_file_input_data"), id="image"), + pytest.param(lf("webp_image_file_input_data"), id="webp_image"), + pytest.param(lf("png_image_file_input_data"), id="png_image"), ], ) def parametrized_file_input_data( @@ -163,7 +198,7 @@ async def file( ) with file.path.open("wb") as f: - f.write(parametrized_file_input_data.content) + f.write(parametrized_file_input_data.processed_content) return file diff --git a/tests/storage_v2/functional/test_file_reads_rst.py b/tests/storage_v2/functional/test_file_reads_rst.py index d4b355b7..204ce3d1 100644 --- a/tests/storage_v2/functional/test_file_reads_rst.py +++ b/tests/storage_v2/functional/test_file_reads_rst.py @@ -79,7 +79,7 @@ async def test_file_reading( }, expected_json=None, ) - assert response.content == parametrized_file_input_data.content + assert response.content == parametrized_file_input_data.processed_content async def test_file_reading_not_modified_by_etag( diff --git a/tests/storage_v2/functional/test_file_uploads_rst.py b/tests/storage_v2/functional/test_file_uploads_rst.py index 0240058b..84458aaa 100644 --- a/tests/storage_v2/functional/test_file_uploads_rst.py +++ b/tests/storage_v2/functional/test_file_uploads_rst.py @@ -1,6 +1,8 @@ +import random from uuid import UUID import pytest +from faker import Faker from pytest_lazy_fixtures import lf, lfc from starlette import status from starlette.testclient import TestClient @@ -47,7 +49,7 @@ async def test_file_uploading( files={ "upload": ( parametrized_file_input_data.name, - parametrized_file_input_data.content, + parametrized_file_input_data.input_content, parametrized_file_input_data.content_type, ) }, @@ -73,15 +75,69 @@ async def test_file_uploading( assert file.path.is_file() with file.path.open("rb") as f: - assert f.read() == parametrized_file_input_data.content + assert f.read() == parametrized_file_input_data.processed_content await file.delete() +CONTENT_TYPES_AND_FILE_EXTENSIONS: list[tuple[str, str]] = [ + ("image/avif", "avif"), + ("image/bmp", "bmp"), + ("image/gif", "gif"), + ("image/x-icon", "ico"), + ("image/jpeg", "jpe"), + ("image/jpeg", "jpeg"), + ("image/jpeg", "jpg"), + ("image/jpx", "jpx"), + ("image/png", "png"), + ("image/tiff", "tif"), + ("image/tiff", "tiff"), + ("image/webp", "webp"), +] + + +@pytest.mark.parametrize( + "file_input_data", + [ + pytest.param(lf("webp_image_file_input_data"), id="webp"), + pytest.param(lf("png_image_file_input_data"), id="png"), + ], +) +async def test_image_file_uploading_content_type_mismatch( + faker: Faker, + authorized_client: TestClient, + file_upload_storage_token: str, + file_input_data: FileInputData, +) -> None: + content_type, file_extension = random.choice( + [ + (content_type, file_extension) + for content_type, file_extension in CONTENT_TYPES_AND_FILE_EXTENSIONS + if content_type != file_input_data.content_type + ] + ) + + assert_response( + authorized_client.post( + "/api/protected/storage-service/v2/file-kinds/image/files/", + headers={"X-Storage-Token": file_upload_storage_token}, + files={ + "upload": ( + faker.file_name(extension=file_extension), + file_input_data.input_content, + content_type, + ) + }, + ), + expected_code=status.HTTP_415_UNSUPPORTED_MEDIA_TYPE, + expected_json={"detail": "File content doesn't match the content-type header"}, + ) + + async def test_image_file_uploading_wrong_content_format( + faker: Faker, authorized_client: TestClient, uncategorized_file_content: bytes, - image_file_input_data: FileInputData, file_upload_storage_token: str, ) -> None: assert_response( @@ -90,9 +146,9 @@ async def test_image_file_uploading_wrong_content_format( headers={"X-Storage-Token": file_upload_storage_token}, files={ "upload": ( - image_file_input_data.name, + faker.file_name(extension="webp"), uncategorized_file_content, - image_file_input_data.content_type, + "image/webp", ) }, ), @@ -150,7 +206,7 @@ async def test_file_uploading_invalid_token( files={ "upload": ( parametrized_file_input_data.name, - parametrized_file_input_data.content, + parametrized_file_input_data.input_content, parametrized_file_input_data.content_type, ) }, From 3ac7b57a916135507bce3b5c10ec5e314f6a7b79 Mon Sep 17 00:00:00 2001 From: niqzart Date: Mon, 16 Feb 2026 05:26:30 +0300 Subject: [PATCH 4/4] fix: include image-related asserts in testing for avatars & storage --- tests/storage_v2/conftest.py | 56 +++++++++++++++---- .../functional/test_file_uploads_rst.py | 25 ++++++++- tests/users/functional/test_avatars.py | 51 +++++++++++++++-- 3 files changed, 116 insertions(+), 16 deletions(-) diff --git a/tests/storage_v2/conftest.py b/tests/storage_v2/conftest.py index 9c77791e..ab4f4bd3 100644 --- a/tests/storage_v2/conftest.py +++ b/tests/storage_v2/conftest.py @@ -1,3 +1,4 @@ +from collections.abc import AsyncIterator from dataclasses import dataclass from io import BytesIO from os import stat @@ -74,15 +75,29 @@ def outsider_internal_client( @pytest.fixture() -async def ydoc(faker: Faker, active_session: ActiveSession) -> YDoc: +async def ydoc(faker: Faker, active_session: ActiveSession) -> AsyncIterator[YDoc]: async with active_session(): - return await YDoc.create(content=faker.binary(length=64)) + ydoc = await YDoc.create(content=faker.binary(length=64)) + + yield ydoc + + async with active_session() as session: + session.add(ydoc) + await ydoc.delete() @pytest.fixture() -async def other_ydoc(faker: Faker, active_session: ActiveSession) -> YDoc: +async def other_ydoc( + faker: Faker, active_session: ActiveSession +) -> AsyncIterator[YDoc]: async with active_session(): - return await YDoc.create(content=faker.binary(length=64)) + ydoc = await YDoc.create(content=faker.binary(length=64)) + + yield ydoc + + async with active_session() as session: + session.add(ydoc) + await ydoc.delete() @pytest.fixture() @@ -91,9 +106,18 @@ def missing_ydoc_id() -> UUID: @pytest.fixture() -async def access_group(active_session: ActiveSession, ydoc: YDoc) -> AccessGroup: +async def access_group( + active_session: ActiveSession, + ydoc: YDoc, +) -> AsyncIterator[AccessGroup]: async with active_session(): - return await AccessGroup.create(main_ydoc_id=ydoc.id) + access_group = await AccessGroup.create(main_ydoc_id=ydoc.id) + + yield access_group + + async with active_session() as session: + session.add(access_group) + await access_group.delete() @pytest.fixture() @@ -143,7 +167,7 @@ def uncategorized_file_input_data( name=faker.file_name(), input_content=uncategorized_file_content, processed_content=uncategorized_file_content, - content_type=faker.mime_type(), + content_type=faker.mime_type(category="application"), ) @@ -190,7 +214,7 @@ def parametrized_file_input_data( async def file( active_session: ActiveSession, parametrized_file_input_data: FileInputData, -) -> File: +) -> AsyncIterator[File]: async with active_session(): file = await File.create( name=parametrized_file_input_data.name, @@ -200,7 +224,11 @@ async def file( with file.path.open("wb") as f: f.write(parametrized_file_input_data.processed_content) - return file + yield file + + async with active_session() as session: + session.add(file) + await file.delete() @pytest.fixture() @@ -215,13 +243,19 @@ async def access_group_file( active_session: ActiveSession, access_group: AccessGroup, file: File, -) -> AccessGroupFile: +) -> AsyncIterator[AccessGroupFile]: async with active_session(): - return await AccessGroupFile.create( + access_group_file = await AccessGroupFile.create( access_group_id=access_group.id, file_id=file.id, ) + yield access_group_file + + async with active_session() as session: + session.add(access_group_file) + await access_group_file.delete() + @pytest.fixture() def file_response(file: File) -> FileResponse: diff --git a/tests/storage_v2/functional/test_file_uploads_rst.py b/tests/storage_v2/functional/test_file_uploads_rst.py index 84458aaa..7173678b 100644 --- a/tests/storage_v2/functional/test_file_uploads_rst.py +++ b/tests/storage_v2/functional/test_file_uploads_rst.py @@ -1,8 +1,11 @@ import random +from io import BytesIO from uuid import UUID import pytest from faker import Faker +from PIL import Image +from pydantic_marshals.contains import assert_contains from pytest_lazy_fixtures import lf, lfc from starlette import status from starlette.testclient import TestClient @@ -75,7 +78,27 @@ async def test_file_uploading( assert file.path.is_file() with file.path.open("rb") as f: - assert f.read() == parametrized_file_input_data.processed_content + real_file_content = f.read() + + if parametrized_file_input_data.content_type.startswith("image/"): + image_result = Image.open(BytesIO(real_file_content)) + try: + image_result.verify() + except Exception as e: + raise AssertionError("Invalid resulting image") from e + + assert_contains( + { + "image_format": image_result.format, + "image_content": real_file_content, + }, + { + "image_format": "WEBP", + "image_content": parametrized_file_input_data.processed_content, + }, + ) + else: + assert real_file_content == parametrized_file_input_data.processed_content await file.delete() diff --git a/tests/users/functional/test_avatars.py b/tests/users/functional/test_avatars.py index e37d7444..652e2519 100644 --- a/tests/users/functional/test_avatars.py +++ b/tests/users/functional/test_avatars.py @@ -4,6 +4,7 @@ import pytest from faker import Faker from PIL import Image +from pydantic_marshals.contains import assert_contains from pytest_lazy_fixtures import lfc from starlette import status from starlette.testclient import TestClient @@ -27,8 +28,8 @@ def image_content(request: PytestRequest[bytes]) -> bytes: @pytest.fixture() def processed_image_content(image_content: bytes) -> bytes: - image = Image.open(BytesIO(image_content)) - image.resize(User.avatar_shape) + image: Image.Image = Image.open(BytesIO(image_content)) + image = image.resize(User.avatar_shape) processed_image_buffer = BytesIO() image.save(processed_image_buffer, format="webp") @@ -60,7 +61,28 @@ async def test_avatar_uploading( assert user.avatar_path.is_file() with user.avatar_path.open("rb") as f: - assert f.read() == processed_image_content + real_image_content = f.read() + + image_result = Image.open(BytesIO(real_image_content)) + try: + image_result.verify() + except Exception as e: + raise AssertionError("Invalid resulting image") from e + + assert_contains( + { + "image_format": image_result.format, + "image_width": image_result.width, + "image_height": image_result.height, + "image_content": real_image_content, + }, + { + "image_format": "WEBP", + "image_width": User.avatar_shape[0], + "image_height": User.avatar_shape[1], + "image_content": processed_image_content, + }, + ) user.avatar_path.unlink() @@ -96,7 +118,28 @@ async def test_avatar_replacing( assert user.avatar_path.is_file() with user.avatar_path.open("rb") as f: - assert f.read() == processed_image_content + real_image_content = f.read() + + image_result = Image.open(BytesIO(real_image_content)) + try: + image_result.verify() + except Exception as e: + raise AssertionError("Invalid resulting image") from e + + assert_contains( + { + "image_format": image_result.format, + "image_width": image_result.width, + "image_height": image_result.height, + "image_content": real_image_content, + }, + { + "image_format": "WEBP", + "image_width": User.avatar_shape[0], + "image_height": User.avatar_shape[1], + "image_content": processed_image_content, + }, + ) @pytest.mark.usefixtures("create_avatar")