Big refactoring

Author: codeSamuraii

app.py

@@ -1,8 +1,7 @@
 import os
-import redis
 import logging
-import redis.asyncio
 import sentry_sdk
+import redis.asyncio
 from fastapi import FastAPI
 from fastapi.responses import FileResponse
 from fastapi.staticfiles import StaticFiles
@@ -22,13 +21,13 @@
 # Redis
 redis_client = redis.asyncio.from_url(os.getenv("REDIS_URL", "redis://localhost:6379"))
 
+# FastAPI
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     setup_logging()
     yield
     await redis_client.close()
 
-# FastAPI
 app = FastAPI(
     debug=True,
     title="Transit.sh",

lib/callbacks.py

@@ -0,0 +1,28 @@
+from fastapi import HTTPException
+from starlette.requests import Request
+from starlette.websockets import WebSocket, WebSocketState
+from typing import Awaitable, Callable
+
+
+def send_error_and_close(websocket: WebSocket) -> Callable[[Exception | str], Awaitable[None]]:
+    """Callback to send an error message and close the WebSocket connection."""
+
+    async def _send_error_and_close(error: Exception | str) -> None:
+        message = str(error) if isinstance(error, Exception) else error
+        if websocket.client_state == WebSocketState.CONNECTED:
+            await websocket.send_text(f"Error: {message}")
+            await websocket.close(code=1011, reason=message)
+
+    return _send_error_and_close
+
+
+def raise_http_exception(request: Request) -> Callable[[Exception | str], Awaitable[None]]:
+    """Callback to raise an HTTPException with a specific status code."""
+
+    async def _raise_http_exception(error: Exception | str) -> None:
+        message = str(error) if isinstance(error, Exception) else error
+        code = error.status_code if isinstance(error, HTTPException) else 400
+        if not await request.is_disconnected():
+            raise HTTPException(status_code=code, detail=message)
+
+    return _raise_http_exception

lib/metadata.py

@@ -0,0 +1,56 @@
+import json
+from dataclasses import dataclass, asdict
+from starlette.datastructures import Headers
+from typing import Optional, Self
+
+
+@dataclass(frozen=True)
+class FileMetadata:
+    size: int
+    name: str
+    content_type: Optional[str] = None
+
+    def to_json(self) -> str:
+        return json.dumps(asdict(self), skipkeys=True)
+
+    @classmethod
+    def from_json(cls, data: str) -> Self:
+        return cls(**json.loads(data))
+
+    @classmethod
+    def get_from_http_headers(cls, headers: Headers, filename: str) -> Self:
+        return cls(
+            name=cls.escape_filename(filename),
+            size=cls.process_length(headers.get('content-length', '0')),
+            content_type=headers.get('content-type', '')
+        )
+
+    @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) -> str:
+        """Escape special characters in the filename."""
+        return str(filename).encode('latin-1', 'ignore').decode('utf-8', 'ignore')
+
+    @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 __str__(self):
+        return f"{self.name} ({self.size/(1024**2):.1f} MiB - {self.content_type})"
+
+    def __repr__(self):
+        return f"FileMetadata(name={self.name!r}, size={self.size/(1024**2):.1f}, content_type={self.content_type!r})"

lib/store.py

@@ -12,7 +12,7 @@ class Store:
     Handles data queuing and event signaling for transfer coordination.
     """
 
-    _redis: Optional[redis.Redis] = None
+    redis_client: None | redis.Redis = None
 
     def __init__(self, transfer_id: str):
         self.transfer_id = transfer_id
@@ -22,58 +22,32 @@ def __init__(self, transfer_id: str):
         self._k_queue = self.key('queue')
         self._k_meta = self.key('metadata')
         self._k_cleanup = f'cleanup:{transfer_id}'
-        self._cleanup = False
 
     @classmethod
     def get_redis(cls) -> redis.Redis:
         """Get the Redis client instance."""
-        if cls._redis is None:
+        if cls.redis_client is None:
             from app import redis_client
-            cls._redis = redis_client
-        return cls._redis
+            cls.redis_client = redis_client
+        return cls.redis_client
 
-    def key(self, name: str):
-        """Get the Redis key for the provided name with this transfer."""
+    def key(self, name: str) -> str:
+        """Get the Redis key for this transfer with the provided name."""
         return f'transfer:{self.transfer_id}:{name}'
 
-    async def cleanup_started(self) -> bool:
-        """Check if cleanup has been initiated for this transfer."""
-        challenge = random.randbytes(8)
-        await self.redis.set(self._k_cleanup, challenge, ex=60, nx=True)
-        if await self.redis.get(self._k_cleanup) == challenge:
-            return False
-        return True
-
-    async def cleanup(self) -> int:
-        """Remove all keys related to this transfer."""
-        if await self.cleanup_started():
-            return 0
-
-        pattern = self.key('*')
-        keys_to_delete = set()
-
-        cursor = 0
-        while True:
-            cursor, keys = await self.redis.scan(cursor, match=pattern)
-            keys_to_delete |= set(keys)
-            if cursor == 0:
-                break
-
-        if keys_to_delete:
-            return await self.redis.delete(*keys_to_delete)
-
     ## Queue operations ##
 
     async def _wait_for_queue_space(self, maxsize: int) -> None:
         while await self.redis.llen(self._k_queue) >= maxsize:
             await asyncio.sleep(0.5)
 
-    async def put_in_queue(self, data: bytes, maxsize: int = 16, timeout: float = 30.0) -> None:
+    async def put_in_queue(self, data: bytes, maxsize: int = 16, timeout: float = 10.0) -> None:
         """Add data to the transfer queue with backpressure control."""
-        await asyncio.wait_for(self._wait_for_queue_space(maxsize), timeout=timeout)
+        async with asyncio.timeout(timeout):
+            await self._wait_for_queue_space(maxsize)
         await self.redis.lpush(self._k_queue, data)
 
-    async def get_from_queue(self, timeout: float = 30.0) -> bytes:
+    async def get_from_queue(self, timeout: float = 10.0) -> bytes:
         """Get data from the transfer queue with timeout."""
         result = await self.redis.brpop([self._k_queue], timeout=timeout)
         if not result:
@@ -87,15 +61,15 @@ async def get_from_queue(self, timeout: float = 30.0) -> bytes:
     async def set_event(self, event_name: str, expiry: float = 300.0) -> None:
         """Set an event flag for this transfer."""
         event_key = self.key(event_name)
-        event_marker_key = self.key(f"{event_name}:marker")
+        event_marker_key = f'{event_key}:marker'
 
         await self.redis.set(event_marker_key, '1', ex=int(expiry))
         await self.redis.publish(event_key, '1')
 
     async def wait_for_event(self, event_name: str, timeout: float = 300.0) -> None:
         """Wait for an event to be set for this transfer."""
         event_key = self.key(event_name)
-        event_marker_key = self.key(f"{event_name}:marker")
+        event_marker_key = f'{event_key}:marker'
         pubsub = self.redis.pubsub(ignore_subscribe_messages=True)
         await pubsub.subscribe(event_key)
 
@@ -135,10 +109,64 @@ async def _listen_for_message():
 
     ## Metadata operations ##
 
-    async def set_metadata(self, metadata: str, expiry: float = 3600.0) -> None:
+    async def set_metadata(self, metadata: str) -> None:
         """Store transfer metadata."""
-        await self.redis.set(self._k_meta, metadata, ex=int(expiry), nx=True)
+        if int (await self.redis.exists(self._k_meta)) > 0:
+            raise KeyError(f"Metadata for transfer '{self.transfer_id}' already exists.")
+        await self.redis.set(self._k_meta, metadata, nx=True)
 
-    async def get_metadata(self) -> Optional[str]:
+    async def get_metadata(self) -> str | None:
         """Retrieve transfer metadata."""
         return await self.redis.get(self._k_meta)
+
+    ## Transfer state operations ##
+
+    async def set_completed(self) -> None:
+        """Mark the transfer as completed."""
+        await self.redis.set(f'completed:{self.transfer_id}', '1', ex=300, nx=True)
+
+    async def is_completed(self) -> bool:
+        """Check if the transfer is marked as completed."""
+        return await self.redis.exists(f'completed:{self.transfer_id}') > 0
+
+    async def set_interrupted(self) -> None:
+        """Mark the transfer as interrupted."""
+        await self.redis.set(f'interrupt:{self.transfer_id}', '1', ex=300, nx=True)
+        await self.redis.ltrim(self._k_queue, 0, 0)
+
+    async def is_interrupted(self) -> bool:
+        """Check if the transfer was interrupted."""
+        return await self.redis.exists(f'interrupt:{self.transfer_id}') > 0
+
+    ## Cleanup operations ##
+
+    async def cleanup_started(self) -> bool:
+        """
+        Check if cleanup has already been initiated for this transfer.
+        This uses a set/get pattern with challenge to avoid race conditions.
+        """
+        challenge = random.randbytes(8)
+        await self.redis.set(self._k_cleanup, challenge, ex=60, nx=True)
+        if await self.redis.get(self._k_cleanup) == challenge:
+            return False
+        return True
+
+    async def cleanup(self) -> int:
+        """Remove all keys related to this transfer."""
+        if await self.cleanup_started():
+            return 0
+
+        pattern = self.key('*')
+        keys_to_delete = set()
+
+        cursor = 0
+        while True:
+            cursor, keys = await self.redis.scan(cursor, match=pattern)
+            keys_to_delete |= set(keys)
+            if cursor == 0:
+                break
+
+        if keys_to_delete:
+            self.log.debug(f"- Cleaning up {len(keys_to_delete)} keys")
+            return await self.redis.delete(*keys_to_delete)
+        return 0