Adding websocket support & refactoring

Amend: requirements.txt

Author: codeSamuraii

lib/__init__.py

@@ -0,0 +1 @@
+from .duplex import Duplex

lib/classes.py renamed to lib/duplex.py

@@ -2,7 +2,6 @@
 import weakref
 from fastapi import Request
 from dataclasses import dataclass
-from fastapi.responses import StreamingResponse
 
 
 @dataclass
@@ -16,52 +15,71 @@ class Duplex:
 
     instances = weakref.WeakValueDictionary()
 
-    def __init__(self, stream, identifier: str, file: File, wait_for_client: bool):
-        self.stream = stream
+    def __init__(self, identifier: str, file: File, stream=None):
         self.identifier = identifier
         self.file = file
-        self.queue = asyncio.Queue(1 if wait_for_client else 0)
+        self.queue = asyncio.Queue(1)
         self.client_connected = asyncio.Event()
 
     @staticmethod
-    def get_upload_details(request: Request):
-        stream = request.stream
-        identifier = request.path_params.get('identifier')
+    def get_file_from_request(request: Request):
         file = File(
             name=request.path_params.get('file_name'),
             size=int(request.headers.get('content-length')),
             content_type=request.headers.get('content-type')
         )
-        return stream, identifier, file
+        return file
+    
+    @staticmethod
+    def get_file_from_header(header: dict):
+        file = File(
+            name=header['file_name'],
+            size=int(header['file_size']),
+            content_type=header['file_type']
+        )
+        return file
+    
+    def get_file_info(self):
+        return self.file.name, self.file.size, self.file.content_type
 
     @classmethod
-    def from_upload(cls, request: Request):
-        stream, identifier, file = cls.get_upload_details(request)
-        duplex = cls(stream, identifier, file, wait_for_client=True)
+    def create_duplex(cls, identifier: str, file: File):
+        duplex = cls(identifier, file)
+        cls.instances[identifier] = duplex
+        return duplex
+    
+    @classmethod
+    def create_duplex_ws(cls, identifier: str, name: str, size: int, type: str):
+        file = File(name=name, size=size, content_type=type)
+        duplex = cls(identifier, file)
         cls.instances[identifier] = duplex
         return duplex
 
     @classmethod
-    def from_identifer(cls, identifier: str):
+    def get(cls, identifier: str):
         if duplex := cls.instances.get(identifier):
             return duplex
         else:
             raise KeyError(f"Duplex '{identifier}' not found.")
 
     def get_file_info(self):
         return self.file.name, self.file.size, self.file.content_type
+
+    async def wait_for_empty_queue(self, seconds=600):
+        while not self.queue.empty() and seconds > 0:
+            await asyncio.sleep(1)
+            seconds -= 1
 
-    async def transfer(self):
+    async def transfer(self, stream):
         bytes_read = 0
 
-        async for chunk in self.stream():
+        async for chunk in stream:
             bytes_read += len(chunk)
             await self.queue.put(chunk)
 
         await self.queue.put(None)
-
-        while not self.queue.empty():
-            await asyncio.sleep(0.5)
+        await self.wait_for_empty_queue()
+        return bytes_read
 
     async def receive(self):
         while True:

requirements.txt

@@ -14,3 +14,4 @@ pydantic-core>=2.10.1
 sniffio>=1.3.0
 starlette>=0.27.0
 typing-extensions>=4.8.0
+websockets>=12.0

views/__init__.py

@@ -0,0 +1,2 @@
+from .http import router as http_router
+from .websockets import router as ws_router

views/http.py

@@ -0,0 +1,49 @@
+import asyncio
+from fastapi import Request, APIRouter
+from fastapi.responses import Response, StreamingResponse, PlainTextResponse
+
+from lib import Duplex
+
+router = APIRouter()
+
+
+@router.put("/{identifier}/{file_name}")
+async def http_upload(request: Request, identifier: str, file_name: str):
+    uid = identifier
+    print(f"{uid} - HTTP transfer request." )
+
+    file = Duplex.get_file_from_request(request)
+    duplex = Duplex.create_duplex(identifier, file)
+
+    print(f"{uid} - Waiting for client to connect...")
+    await duplex.client_connected.wait()
+
+    print(f"{uid} - Client connected. Transfering...")
+    await duplex.transfer(request.stream())
+
+    print(f"{uid} - Transfer complete.")
+    return Response(status_code=200)
+
+
+@router.get("/{identifier}")
+async def http_download(identifier: str):
+    uid = identifier
+    print(f"{uid} - HTTP download request." )
+
+    try:
+        duplex = Duplex.get(identifier)
+    except KeyError:
+        return PlainTextResponse("File not found", status_code=404)
+    
+    print(f"{uid} - Notifying client is connected.")
+    duplex.client_connected.set()
+    await asyncio.sleep(0.5)
+
+    file_name, file_size, file_type = duplex.get_file_info()
+
+    print(f"{uid} - Starting download.")
+    return StreamingResponse(
+        duplex.receive(),
+        media_type=file_type,
+        headers={"Content-Disposition": f"attachment; filename={file_name}", "Content-Length": str(file_size)}
+    )

views/websockets.py

@@ -0,0 +1,64 @@
+import asyncio
+from fastapi import WebSocket, APIRouter
+
+from lib import Duplex
+
+
+router = APIRouter()
+
+
+@router.websocket("/send/{identifier}")
+async def websocket_upload(websocket: WebSocket, identifier: str):
+    uid = identifier
+    await websocket.accept()
+    print(f"{uid} - Websocket transfer request." )
+
+    header = await websocket.receive_json()
+
+    try:
+        file = Duplex.get_file_from_header(header)
+    except KeyError:
+        print(f"{uid} - Invalid header: {header}")
+        return
+    
+    duplex = Duplex.create_duplex(uid, file)
+
+    await duplex.client_connected.wait()
+    await websocket.send_text(f"Go for file chunks")
+
+    print(f"{uid} - Starting transfer...")
+    await duplex.transfer(websocket.iter_bytes())
+
+    print(f"{uid} - Transfer complete.")
+
+
+@router.websocket("/receive/{identifier}")
+async def websocket_download(websocket: WebSocket, identifier: str):
+    uid = identifier
+    await websocket.accept()
+    print(f"{uid} - Websocket download request." )
+
+    try:
+        duplex = Duplex.get(identifier)
+    except KeyError:
+        print(f"{uid} - File not found.")
+        await websocket.send_text("File not found")
+        return
+
+    file_name, file_size, file_type = duplex.get_file_info()
+    await websocket.send_json({'file_name': file_name, 'file_size': file_size, 'file_type': file_type})
+
+    print(f"{uid} - Waiting for go-ahead...")
+    while (msg := await websocket.receive_text()) != "Go for file chunks":
+        print(f"{uid} - Unexpected message: {msg}")
+    
+    print(f"{uid} - Notifying client is connected.")
+    duplex.client_connected.set()
+    await asyncio.sleep(0.5)
+
+    print(f"{uid} - Starting transfer...")
+    async for chunk in duplex.receive():
+        await websocket.send_bytes(chunk)
+    await websocket.send_bytes(b'')
+    print(f"{uid} - Transfer complete.")
+

webapp.py

@@ -1,13 +1,13 @@
-import asyncio
 from pathlib import Path
-from fastapi import FastAPI, Request
+from fastapi import FastAPI
 from fastapi.staticfiles import StaticFiles
-from fastapi.responses import StreamingResponse, FileResponse, PlainTextResponse, Response
-
-from lib.classes import Duplex
+from fastapi.responses import FileResponse
 
+from views import http_router, ws_router
 
 app = FastAPI()
+app.include_router(http_router)
+app.include_router(ws_router)
 
 
 @app.get('/')
@@ -25,43 +25,10 @@ async def get_health():
     return {"status": "ok"}
 
 
-@app.put("/{identifier}/{file_name}")
-async def upload_file(request: Request, identifier: str, file_name: str):
-    duplex = Duplex.from_upload(request)
-    id_ = duplex.identifier
-
-    print(f"[{id_}] Waiting for client to connect...")
-    await duplex.client_connected.wait()
-
-    print(f"[{id_}] Client connected. Transfering...")
-    await duplex.transfer()
-
-    print(f"[{id_}] Transfer complete.")
-    return Response(status_code=200)
-
-
-@app.get("/{identifier}")
-async def get_file(identifier: str):
-    try:
-        duplex = Duplex.from_identifer(identifier)
-        file_name, file_size, file_type = duplex.get_file_info()
-    except KeyError:
-        return PlainTextResponse("File not found", status_code=404)
-    
-    duplex.client_connected.set()
-    await asyncio.sleep(0.5)
-
-    return StreamingResponse(
-        duplex.receive(),
-        media_type=file_type,
-        headers={"Content-Disposition": f"attachment; filename={file_name}", "Content-Length": str(file_size)}
-    )
-
-
-# Mount local static directory for HTML
+# Mount local static for HTML
 app.mount('/static', StaticFiles(directory='static', html=True), name='static')
 
-# Mount remote disk if present or local static for CSS
+# Mount remote if present or local static for CSS
 if Path('/extra').exists():
     app.mount('/css', StaticFiles(directory='/extra'), name='css')
 else: