Use Pydantic for FileMetadata validation

Author: codeSamuraii

lib/metadata.py

@@ -1,81 +1,51 @@
-import json
-import math
-from dataclasses import dataclass, asdict
 from starlette.datastructures import Headers
-from typing import Optional, Self
+from pydantic import BaseModel, Field, field_validator, ByteSize, StrictStr, ConfigDict, AliasChoices
+from typing import Optional, Self, Annotated
 
 
-@dataclass(frozen=True)
-class FileMetadata:
-    size: int
-    name: str
-    content_type: Optional[str] = None
+class FileMetadata(BaseModel):
+    name: StrictStr = Field(description="File name", min_length=2, max_length=255, validation_alias=AliasChoices('name', 'file_name'))
+    size: ByteSize = Field(description="Size in bytes", gt=0,validation_alias=AliasChoices('size', 'file_size'))
+    type: StrictStr = Field(description="MIME type", default='application/octet-stream',validation_alias=AliasChoices('type', 'file_type', 'content_type'))
 
-    def to_json(self) -> str:
-        return json.dumps(asdict(self), skipkeys=True)
+    model_config = ConfigDict(validate_by_name=True, populate_by_name=True)
 
-    def to_readable_dict(self) -> dict:
-        return dict(
-            file_name=self.name,
-            file_size=self.format_size(self.size),
-            file_type=self.format_type(self.content_type)
-        )
+    @field_validator('name')
+    @classmethod
+    def validate_name(cls, v: str) -> str:
+        safe_filename = str(v).translate(str.maketrans(':;|*@/\\', '       ')).strip()
+        return safe_filename.encode('latin-1', 'ignore').decode('utf-8', 'ignore')
 
     @classmethod
     def from_json(cls, data: str) -> Self:
-        return cls(**json.loads(data))
+        return cls.model_validate_json(data)
+
+    def to_json(self) -> str:
+        return self.model_dump_json()
 
     @classmethod
     def get_from_http_headers(cls, headers: Headers, filename: str) -> Self:
+        """Create metadata from headers of an HTTP upload request."""
         return cls(
-            name=cls.escape_filename(filename),
-            size=cls.process_length(headers.get('content-length', '0')),
-            content_type=headers.get('content-type', '')
+            name=filename,
+            size=headers.get('content-length', '0'),
+            type=headers.get('content-type', '') or None
         )
 
     @classmethod
     def get_from_json(cls, header: dict) -> Self:
-        return cls(
-            name=cls.escape_filename(header['file_name']),
-            size=cls.process_length(header['file_size']),
-            content_type=header['file_type']
-        )
-
-    @staticmethod
-    def escape_filename(filename: str | int) -> str:
-        """Escape special characters in the filename."""
-        safe_filename = str(filename).translate(str.maketrans('', '', ':;|*@/\\'))
-        return safe_filename.encode('latin-1', 'ignore').decode('utf-8', 'ignore')
-
-    @staticmethod
-    def format_size(size_bytes: int) -> str:
-        """Return human-readable file size."""
-        if size_bytes == 0:
-            return "0 B"
-        units = ("B", "KiB", "MiB", "GiB", "TiB", "PiB", "EiB", "ZiB")
-        i = math.floor(math.log(size_bytes, 1024))
-        p = math.pow(1024, i)
-        s = round(size_bytes / p, 1)
-        return f"{s} {units[i]}"
-
-    @staticmethod
-    def format_type(content_type: Optional[str]) -> str:
-        """Return human-readable file type."""
-        return content_type or "unknown"
+        """Create metadata from a JSON dictionary."""
+        return cls(**header)
 
-    @staticmethod
-    def process_length(length: str | int) -> int:
-        """Convert size string to bytes."""
-        try:
-            size = int(str(length).strip().replace(' ', ''))
-        except ValueError:
-            raise ValueError(f"Invalid size format: {length}")
-        if size <= 0:
-            raise ValueError("File size has to be positive.")
-        return size
+    def to_readable_dict(self) -> dict:
+        return dict(
+            file_name=self.name,
+            file_size=self.size.human_readable(),
+            file_type=self.type,
+        )
 
     def __str__(self):
-        return f"{self.name} ({self.size/(1024**2):.1f} MiB - {self.content_type})"
+        return f"{self.name} ({self.size.human_readable()} - {self.type})"
 
     def __repr__(self):
-        return f"FileMetadata(name={self.name!r}, size={self.size/(1024**2):.1f}, content_type={self.content_type!r})"
+        return f"FileMetadata(name={self.name!r}, size={self.size.human_readable()}, type={self.type!r})"

lib/transfer.py

@@ -44,7 +44,7 @@ def _format_uid(uid: str):
         return str(uid).strip().encode('ascii', 'ignore').decode()
 
     def get_file_info(self):
-        return self.file.name, self.file.size, self.file.content_type
+        return self.file.name, self.file.size, self.file.type
 
     async def wait_for_event(self, event_name: str, timeout: float = 300.0):
         await self.store.wait_for_event(event_name, timeout)

requirements.txt

@@ -9,3 +9,4 @@ jinja2>=3.1.6
 httpx>=0.28.1
 rich>=14.0.0
 sentry-sdk[fastapi]>=2.32.0
+pydantic>=2.0.0

tests/conftest.py

@@ -28,7 +28,7 @@ def get_redis_override(*args, **kwargs) -> redis.Redis:
     app.state.redis = redis_client
 
     transport = httpx.ASGITransport(app=app)
-    async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
+    async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client:
         # Patch the `get_redis` method of the `Store` class
         with patch.object(Store, 'get_redis', new=get_redis_override):
             print("")
@@ -45,7 +45,7 @@ def get_redis_override(*args, **kwargs) -> redis.Redis:
 
     # Patch the `get_redis` method of the `Store` class
     with patch.object(Store, 'get_redis', new=get_redis_override):
-        with TestClient(app) as client:
+        with TestClient(app, base_url="http://testserver") as client:
             print("")
             yield client
 

tests/helpers.py

@@ -14,7 +14,7 @@ def generate_test_file(size_in_kb: int = 10) -> tuple[bytes, FileMetadata]:
     metadata = FileMetadata(
         name="test_file.bin",
         size=len(content),
-        content_type="application/octet-stream"
+        type="application/octet-stream"
     )
     return content, metadata
 

tests/test_endpoints.py

@@ -44,15 +44,15 @@ async def test_transfer_id_already_used(websocket_client):
         ws.send_json({
             'file_name': file_metadata.name,
             'file_size': file_metadata.size,
-            'file_type': file_metadata.content_type
+            'file_type': file_metadata.type
         })
 
         # Second attempt should fail with an error message
         with websocket_client.websocket_connect(f"/send/{uid}") as ws2:
             ws2.send_json({
                 'file_name': file_metadata.name,
                 'file_size': file_metadata.size,
-                'file_type': file_metadata.content_type
+                'file_type': file_metadata.type
             })
             response = ws2.receive_text()
             assert "Error: Transfer ID is already used." in response
@@ -76,7 +76,7 @@ async def mock_wait_for_client_connected(self):
         ws.send_json({
             'file_name': file_metadata.name,
             'file_size': file_metadata.size,
-            'file_type': file_metadata.content_type
+            'file_type': file_metadata.type
         })
         # This should timeout because we are not starting a receiver
         response = ws.receive_text()
@@ -97,7 +97,7 @@ async def sender():
                 ws.send_json({
                     'file_name': file_metadata.name,
                     'file_size': file_metadata.size,
-                    'file_type': file_metadata.content_type
+                    'file_type': file_metadata.type
                 })
                 await asyncio.sleep(1.0)  # Allow receiver to connect
 
@@ -142,7 +142,7 @@ async def test_prefetcher_request(test_client: httpx.AsyncClient, websocket_clie
         ws.send_json({
             'file_name': file_metadata.name,
             'file_size': file_metadata.size,
-            'file_type': file_metadata.content_type
+            'file_type': file_metadata.type
         })
         await asyncio.sleep(1.0)
 
@@ -168,7 +168,7 @@ async def test_browser_download_page(test_client: httpx.AsyncClient, websocket_c
         ws.send_json({
             'file_name': file_metadata.name,
             'file_size': file_metadata.size,
-            'file_type': file_metadata.content_type
+            'file_type': file_metadata.type
         })
         await asyncio.sleep(1.0)
 

tests/test_journeys.py

@@ -18,7 +18,7 @@ async def sender():
             ws.send_json({
                 'file_name': file_metadata.name,
                 'file_size': file_metadata.size,
-                'file_type': file_metadata.content_type
+                'file_type': file_metadata.type
             })
             await asyncio.sleep(1.0)
 
@@ -72,7 +72,7 @@ async def test_http_upload_http_download(test_client: httpx.AsyncClient):
 
     async def sender():
         headers = {
-            'Content-Type': file_metadata.content_type,
+            'Content-Type': file_metadata.type,
             'Content-Length': str(file_metadata.size)
         }
         async with test_client.stream("PUT", f"/{uid}/{file_metadata.name}", content=file_content, headers=headers) as response:

tests/test_unit.py

@@ -1,42 +1,80 @@
 import pytest
+from pydantic import ValidationError
 from lib.metadata import FileMetadata
 
 
-@pytest.mark.parametrize("size_bytes, expected", [
-    (0, "0 B"),
-    (1023, "1023.0 B"),
-    (1024, "1.0 KiB"),
-    (1536, "1.5 KiB"),
-    (1024 ** 2, "1.0 MiB"),
-    (int(1.5 * 1024 ** 2), "1.5 MiB"),
-    (1024 ** 3, "1.0 GiB"),
-])
-def test_format_size(size_bytes, expected):
-    assert FileMetadata.format_size(size_bytes) == expected
-
-
-@pytest.mark.parametrize("filename, expected", [
-    ("file:name.txt", "filename.txt"),
-    ("file|name.txt", "filename.txt"),
-    ("[email protected]", "filename.txt"),
-    ("file/name.txt", "filename.txt"),
-    ("file\\name.txt", "filename.txt"),
-    ("valid-name.zip", "valid-name.zip"),
-])
-def test_escape_filename(filename, expected):
-    assert FileMetadata.escape_filename(filename) == expected
-
-
-@pytest.mark.parametrize("length, expected", [
-    ("1024", 1024),
-    (2048, 2048),
-    (" 4096 ", 4096),
-])
-def test_process_length(length, expected):
-    assert FileMetadata.process_length(length) == expected
-
-
-@pytest.mark.parametrize("invalid_length", ["-100", "0", "abc", "1.5"])
-def test_process_length_invalid(invalid_length):
-    with pytest.raises(ValueError):
-        FileMetadata.process_length(invalid_length)
+def test_file_metadata_creation():
+    """Test that FileMetadata can be created with valid data."""
+    metadata = FileMetadata(
+        name="test.txt",
+        size=1024,
+        content_type="text/plain"
+    )
+    assert metadata.name == "test.txt"
+    assert metadata.size == 1024
+    assert metadata.type == "text/plain"
+
+
+def test_file_metadata_validation_invalid_size():
+    """Test that FileMetadata validates size field."""
+    with pytest.raises(ValidationError):
+        FileMetadata(name="test.txt", size=0)
+
+    with pytest.raises(ValidationError):
+        FileMetadata(name="test.txt", size=-1)
+
+
+def test_file_metadata_validation_invalid_name():
+    """Test that FileMetadata validates name field."""
+    with pytest.raises(ValidationError):
+        FileMetadata(name="", size=1024)
+
+
+def test_file_metadata_json_serialization():
+    """Test that FileMetadata can be serialized to and from JSON."""
+    metadata = FileMetadata(
+        name="test.txt",
+        size=1024,
+        content_type="text/plain"
+    )
+
+    json_str = metadata.to_json()
+    deserialized = FileMetadata.from_json(json_str)
+
+    assert deserialized.name == metadata.name
+    assert deserialized.size == metadata.size
+    assert deserialized.type == metadata.type
+
+
+def test_file_metadata_name_escaping():
+    """Test that FileMetadata properly escapes filenames during validation."""
+    metadata = FileMetadata(
+        name="file:name.txt",
+        size=1024
+    )
+    assert metadata.name == "file name.txt"
+
+
+def test_file_metadata_size_conversion():
+    """Test that FileMetadata properly converts size strings to integers."""
+    metadata = FileMetadata(
+        name="test.txt",
+        size="1024"
+    )
+    assert metadata.size == 1024
+    assert isinstance(metadata.size, int)
+
+
+def test_file_metadata_size_human_readable():
+    """Test that FileMetadata properly formats sizes using ByteSize's human_readable method."""
+    metadata = FileMetadata(
+        name="test.txt",
+        size=1024
+    )
+    assert metadata.size.human_readable() == "1.0KiB"
+
+    metadata = FileMetadata(
+        name="test.txt",
+        size=1048576
+    )
+    assert metadata.size.human_readable() == "1.0MiB"

views/http.py

@@ -5,6 +5,7 @@
 from starlette.background import BackgroundTask
 from fastapi.exceptions import HTTPException
 from fastapi.responses import StreamingResponse, PlainTextResponse
+from pydantic import ValidationError
 
 from lib.logging import get_logger
 from lib.callbacks import raise_http_exception
@@ -34,8 +35,8 @@ async def http_upload(request: Request, uid: str, filename: str):
     except KeyError as e:
         log.error("△ Cannot decode file metadata from HTTP headers.", exc_info=e)
         raise HTTPException(status_code=400, detail="Cannot decode file metadata from HTTP headers.")
-    except ValueError as e:
-        log.error("△ Invalid file size.", exc_info=e)
+    except ValidationError as e:
+        log.error("△ Invalid file metadata.", exc_info=e)
         raise HTTPException(status_code=400, detail="Invalid file metadata.")
 
     if file.size > 1024**3:
@@ -48,7 +49,7 @@ async def http_upload(request: Request, uid: str, filename: str):
     except KeyError as e:
         log.warning("△ Transfer ID is already used.")
         raise HTTPException(status_code=409, detail="Transfer ID is already used.")
-    except (ValueError, TypeError) as e:
+    except (TypeError, ValidationError) as e:
         log.error("△ Invalid transfer ID or file metadata.", exc_info=e)
         raise HTTPException(status_code=400, detail="Invalid transfer ID or file metadata.")
 
@@ -90,7 +91,7 @@ async def http_download(request: Request, uid: str):
         transfer = await FileTransfer.get(uid)
     except KeyError:
         raise HTTPException(status_code=404, detail="Transfer not found.")
-    except (ValueError, TypeError) as e:
+    except (TypeError, ValidationError) as e:
         log.error("▼ Invalid transfer ID.", exc_info=e)
         raise HTTPException(status_code=400, detail="Invalid transfer ID.")
     else: