Enhance session management and Dockerfile for improved performance and security

- Added session timeout management in broker.py to automatically clean up inactive sessions.
- Updated Dockerfile to create a non-root user for enhanced security and adjusted file permissions.
- Improved logging in server.py to handle exceptions and connection closures more effectively.
- Enhanced hypercorn.toml with security and performance configurations, including timeouts and worker class settings.

Dockerfile

@@ -8,27 +8,33 @@ RUN apt-get update && apt-get install -y \
     build-essential \
     && rm -rf /var/lib/apt/lists/*
 
+# Crear usuario no root
+RUN adduser --disabled-password --gecos '' appuser && \
+    mkdir -p /code/static && \
+    chown -R appuser:appuser /code
+
 # Copiar solo requirements.txt primero para aprovechar la caché de Docker
-COPY requirements.txt .
+COPY --chown=appuser:appuser requirements.txt .
 
 # Instalar dependencias de Python con pip limpio
 RUN pip install --no-cache-dir --upgrade pip && \
     pip install --no-cache-dir -r requirements.txt
 
-# Crear estructura de directorios
-RUN mkdir -p /code/static
-
 # Copiar archivos estáticos primero
-COPY static/swagger.html /code/static/
-COPY openapi.yaml /code/
+COPY --chown=appuser:appuser static/swagger.html /code/static/
+COPY --chown=appuser:appuser openapi.yaml /code/
 
 # Copiar el resto de los archivos
-COPY . .
+COPY --chown=appuser:appuser . .
 
 # Variables de entorno para Hugging Face Spaces
 ENV PYTHONUNBUFFERED=1
 ENV PYTHONPATH=/code
 ENV PORT=7860
+ENV PYTHONDONTWRITEBYTECODE=1
+
+# Cambiar al usuario no root
+USER appuser
 
 # Exponer el puerto que Hugging Face Spaces espera
 EXPOSE 7860

broker.py

@@ -1,5 +1,6 @@
 import asyncio
 import uuid
+import time
 from typing import AsyncGenerator, Dict, Tuple, Any
 from dataclasses import dataclass
 
@@ -26,14 +27,37 @@ class ClientResponse:
     data: Any
 
 class SessionBroker:
-    def __init__(self):
+    def __init__(self, session_timeout: int = 3600):  # 1 hora por defecto
         self.sessions: Dict[str, asyncio.Queue] = {}
         self.pending_responses: Dict[Tuple[str, str], asyncio.Future] = {}
+        self.session_last_active: Dict[str, float] = {}
+        self.session_timeout = session_timeout
+        asyncio.create_task(self._cleanup_inactive_sessions())
+
+    async def _cleanup_inactive_sessions(self):
+        while True:
+            current_time = time.time()
+            inactive_sessions = [
+                session_id for session_id, last_active in self.session_last_active.items()
+                if current_time - last_active > self.session_timeout
+            ]
+            
+            for session_id in inactive_sessions:
+                if session_id in self.sessions:
+                    del self.sessions[session_id]
+                    del self.session_last_active[session_id]
+                    self.pending_responses = {
+                        k: v for k, v in self.pending_responses.items() 
+                        if k[0] != session_id
+                    }
+            
+            await asyncio.sleep(60)  # Verificar cada minuto
 
     async def send_request(self, session_id: str, data: Any, timeout: int = 60) -> Any:
         if session_id not in self.sessions:
             raise SessionDoesNotExist()
 
+        self.session_last_active[session_id] = time.time()
         request_id = str(uuid.uuid4())
         future = asyncio.get_event_loop().create_future()
         self.pending_responses[(session_id, request_id)] = future
@@ -49,6 +73,7 @@ class SessionBroker:
                 del self.pending_responses[(session_id, request_id)]
             
     async def receive_response(self, session_id: str, response: ClientResponse) -> None:
+        self.session_last_active[session_id] = time.time()
         if (session_id, response.request_id) in self.pending_responses:
             future = self.pending_responses.pop((session_id, response.request_id))
             if not future.done():
@@ -63,11 +88,18 @@ class SessionBroker:
 
         queue = asyncio.Queue()
         self.sessions[session_id] = queue
+        self.session_last_active[session_id] = time.time()
 
         try:
             while True:
                 yield await queue.get()
+                self.session_last_active[session_id] = time.time()
         finally:
             if session_id in self.sessions:
                 del self.sessions[session_id]
-                self.pending_responses = {k: v for k, v in self.pending_responses.items() if k[0] != session_id}
+                if session_id in self.session_last_active:
+                    del self.session_last_active[session_id]
+                self.pending_responses = {
+                    k: v for k, v in self.pending_responses.items() 
+                    if k[0] != session_id
+                }

hypercorn.toml

@@ -12,3 +12,12 @@ websocket_max_message_size = 16777216
 accesslog = "-"
 errorlog = "-"
 loglevel = "INFO"
+
+# Configuraciones de seguridad
+h11_max_incomplete_size = 16384
+keep_alive_timeout = 5
+graceful_timeout = 10
+
+# Configuraciones de rendimiento
+worker_class = "asyncio"
+backlog = 2048

server.py

@@ -36,7 +36,8 @@ with open(OPENAPI_PATH) as f:
     openapi_spec = yaml.safe_load(f)
 
 # Configurar Quart Schema con la especificación OpenAPI
-QuartSchema(app, openapi_spec)
+schema = QuartSchema(app)
+schema.update_openapi(openapi_spec)
 
 app.asgi_app = ProxyHeadersMiddleware(app.asgi_app, trusted_hosts=TRUSTED_HOSTS)
 app.logger.setLevel(LOG_LEVEL)
@@ -103,26 +104,44 @@ async def status() -> Status:
 
 @app.websocket('/session')
 async def session_handler():
-    session_id = secrets.token_hex()
-    app.logger.info(f"{websocket.remote_addr} - NEW SESSION - {session_id}")
-    await websocket.send_as(Session(session_id=session_id), Session)
-
-    task = None
     try:
-        task = asyncio.ensure_future(_receive(session_id))
-        async for request in broker.subscribe(session_id):
-            app.logger.info(f"{websocket.remote_addr} - REQUEST - {session_id} - {json.dumps(asdict(request))}")
-            await websocket.send_as(request, ClientRequest)
-    finally:
-        if task is not None:
-            task.cancel()
-            await task
+        session_id = secrets.token_hex()
+        app.logger.info(f"{websocket.remote_addr} - NEW SESSION - {session_id}")
+        await websocket.send_as(Session(session_id=session_id), Session)
+
+        task = None
+        try:
+            task = asyncio.ensure_future(_receive(session_id))
+            async for request in broker.subscribe(session_id):
+                app.logger.info(f"{websocket.remote_addr} - REQUEST - {session_id} - {json.dumps(asdict(request))}")
+                await websocket.send_as(request, ClientRequest)
+        except websockets.exceptions.ConnectionClosed:
+            app.logger.warning(f"{websocket.remote_addr} - CONNECTION CLOSED - {session_id}")
+        except Exception as e:
+            app.logger.error(f"{websocket.remote_addr} - ERROR - {session_id} - {str(e)}")
+            raise
+        finally:
+            if task is not None:
+                task.cancel()
+                try:
+                    await task
+                except asyncio.CancelledError:
+                    pass
+    except Exception as e:
+        app.logger.error(f"Error in session handler: {str(e)}")
+        raise
 
 async def _receive(session_id: str) -> None:
-    while True:
-        response = await websocket.receive_as(ClientResponse)
-        app.logger.info(f"{websocket.remote_addr} - RESPONSE - {session_id} - {json.dumps(asdict(response))}")
-        await broker.receive_response(session_id, response)
+    try:
+        while True:
+            response = await websocket.receive_as(ClientResponse)
+            app.logger.info(f"{websocket.remote_addr} - RESPONSE - {session_id} - {json.dumps(asdict(response))}")
+            await broker.receive_response(session_id, response)
+    except websockets.exceptions.ConnectionClosed:
+        app.logger.warning(f"{websocket.remote_addr} - RECEIVE CONNECTION CLOSED - {session_id}")
+    except Exception as e:
+        app.logger.error(f"{websocket.remote_addr} - RECEIVE ERROR - {session_id} - {str(e)}")
+        raise
 
 @app.post('/command')
 @validate_request(Command)