Refactor Dockerfile, hypercorn.toml, and server.py for improved security, static file handling, and session management

- Updated Dockerfile to enhance security by creating a non-root user and adjusting file permissions.
- Improved static file handling by copying necessary files and updating the structure for better organization.
- Enhanced hypercorn.toml with server binding options and logging configurations.
- Refactored server.py to streamline session management, including custom exceptions and improved logging.
- Added new WebSocket endpoints for session handling and refined API response structures for better clarity.

Dockerfile

@@ -1,6 +1,7 @@
 # Dockerfile optimizado para Hugging Face Spaces
 FROM python:3.11-slim
 
+# Establecer el directorio de trabajo
 WORKDIR /code
 
 # Instalar dependencias del sistema
@@ -8,7 +9,7 @@ RUN apt-get update && apt-get install -y \
     build-essential \
     && rm -rf /var/lib/apt/lists/*
 
-# Crear usuario no root
+# Crear un usuario no root para mejorar la seguridad
 RUN adduser --disabled-password --gecos '' appuser && \
     mkdir -p /code/static && \
     chown -R appuser:appuser /code
@@ -20,12 +21,16 @@ COPY --chown=appuser:appuser requirements.txt .
 RUN pip install --no-cache-dir --upgrade pip && \
     pip install --no-cache-dir -r requirements.txt
 
-# Copiar archivos estáticos primero
+# Copiar archivos estáticos y de configuración primero
 COPY --chown=appuser:appuser static/swagger.html /code/static/
 COPY --chown=appuser:appuser openapi.yaml /code/
 
-# Copiar el resto de los archivos
-COPY --chown=appuser:appuser . .
+# Copiar el resto de los archivos del proyecto
+COPY --chown=appuser:appuser server.py ./
+COPY --chown=appuser:appuser hypercorn.toml ./
+# Si tienes otros archivos necesarios, cópialos aquí
+# Por ejemplo:
+# COPY --chown=appuser:appuser other_module.py ./ 
 
 # Variables de entorno para Hugging Face Spaces
 ENV PYTHONUNBUFFERED=1
@@ -33,11 +38,11 @@ ENV PYTHONPATH=/code
 ENV PORT=7860
 ENV PYTHONDONTWRITEBYTECODE=1
 
-# Cambiar al usuario no root
+# Cambiar al usuario no root para evitar privilegios elevados
 USER appuser
 
 # Exponer el puerto que Hugging Face Spaces espera
 EXPOSE 7860
 
-# Comando para ejecutar la aplicación
+# Comando para ejecutar la aplicación usando Hypercorn
 CMD ["hypercorn", "--config", "hypercorn.toml", "server:app"]

hypercorn.toml

@@ -1,14 +1,14 @@
-# Hypercorn configuration file
+# hypercorn.toml
 
-# Server binding options
+# Opciones de binding del servidor
 bind = "0.0.0.0:7860"
 workers = 1
 
-# Websocket settings
+# Configuración de WebSockets
 websocket_ping_interval = 20
 websocket_max_message_size = 16777216
 
-# Logging configuration
+# Configuración de logging
 accesslog = "-"
 errorlog = "-"
 loglevel = "INFO"
@@ -20,4 +20,4 @@ graceful_timeout = 10
 
 # Configuraciones de rendimiento
 worker_class = "asyncio"
-backlog = 2048
+backlog = 2048

poetry.toml

@@ -1,2 +1 @@
-[virtualenvs]
 in-project = true

server.py

@@ -1,50 +1,58 @@
 # server.py
-from dataclasses import dataclass, asdict
-import secrets
-import logging
+
 import asyncio
+import importlib.metadata
 import json
-import yaml
-from pathlib import Path
-from typing import Tuple
+import logging
+import secrets
+import uuid
+from dataclasses import dataclass, asdict
+from typing import Any, AsyncGenerator, Dict, Tuple, Union
 
-from quart import Quart, websocket, request, send_from_directory, redirect
+from quart import Quart, websocket, request, send_from_directory
 from quart_schema import QuartSchema, validate_request, validate_response
-from quart_cors import cors
-from uvicorn.middleware.proxy_headers import ProxyHeadersMiddleware
-
-from broker import SessionBroker, SessionDoesNotExist, ClientRequest, ClientResponse, ClientError
+from hypercorn.middleware.proxy_headers import ProxyHeadersMiddleware
 
-# Configuración
-VERSION = "1.0.0"  # Versión de la API
+# Configuraciones
 TIMEOUT: int = 40
 LOG_LEVEL: int = logging.DEBUG
-TRUSTED_HOSTS: list[str] = ["127.0.0.1", "10.16.38.136", "10.16.3.13", "10.16.13.73"]
+TRUSTED_HOSTS: list[str] = ["127.0.0.1", "172.18.0.3"]
 
-# Create app
+# Inicialización de la aplicación Quart
 app = Quart(__name__)
-app = cors(app, 
-    allow_origin=["https://*.hf.space", "https://*.huggingface.co"],
-    allow_methods=["GET", "POST", "OPTIONS"],
-    allow_headers=["Content-Type"],
-    max_age=3600
-)
-
-# Cargar OpenAPI spec
-OPENAPI_PATH = Path(__file__).parent / "openapi.yaml"
-with open(OPENAPI_PATH) as f:
-    openapi_spec = yaml.safe_load(f)
-
-# Configurar Quart Schema con la especificación OpenAPI
-schema = QuartSchema(app)
-schema.update_openapi(openapi_spec)
-
+QuartSchema(app)
 app.asgi_app = ProxyHeadersMiddleware(app.asgi_app, trusted_hosts=TRUSTED_HOSTS)
 app.logger.setLevel(LOG_LEVEL)
 
-broker = SessionBroker()
+# Excepciones personalizadas
+
+class SessionDoesNotExist(Exception):
+    """Error al solicitar un ID de sesión que no existe."""
+    pass
+
+class SessionAlreadyExists(Exception):
+    """Error al crear una sesión con un ID que ya existe."""
+    pass
+
+class ClientError(Exception):
+    """Error cuando el cliente devuelve un error."""
+    def __init__(self, message: str):
+        super().__init__(message)
+        self.message = message
+
+# Clases de datos para solicitudes y respuestas
+
+@dataclass
+class ClientRequest:
+    request_id: str
+    data: Any 
+
+@dataclass
+class ClientResponse:
+    request_id: str
+    error: bool
+    data: Any
 
-# Modelos de datos
 @dataclass
 class Status:
     status: str
@@ -88,124 +96,179 @@ class WriteResponse:
 class ErrorResponse:
     error: str
 
-# Rutas API
-@app.route('/')
-async def root():
-    return await send_from_directory('static', 'swagger.html')
+# Broker para manejar sesiones y comunicaciones
 
-@app.route('/docs')
-async def docs():
-    return redirect('/')
+class SessionBroker:
+    def __init__(self):
+        """Diccionario de session_id -> cola de mensajes pendientes por enviar al cliente"""
+        self.sessions: Dict[str, asyncio.Queue] = {}
+        """Diccionario de (session_id, request_id) -> futuro de la respuesta del cliente"""
+        self.pending_responses: Dict[Tuple[str, str], asyncio.Future] = {}
 
-@app.get("/status")
-@validate_response(Status)
-async def status() -> Status:
-    return Status(status="OK", version=VERSION)
+    async def send_request(self, session_id: str, data: Any, timeout: int = 60) -> Any:
+        if session_id not in self.sessions:
+            raise SessionDoesNotExist()
 
-@app.websocket('/session')
-async def session_handler():
-    try:
-        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)
+        request_id = str(uuid.uuid4())
+        loop = asyncio.get_event_loop()
+        future = loop.create_future()
+        self.pending_responses[(session_id, request_id)] = future
+
+        await self.sessions[session_id].put(ClientRequest(request_id=request_id, data=data))
 
-        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)}")
+            return await asyncio.wait_for(future, timeout)
+        except asyncio.TimeoutError:
             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
+            self.pending_responses.pop((session_id, request_id), None)
+
+    async def receive_response(self, session_id: str, response: ClientResponse) -> None:
+        key = (session_id, response.request_id)
+        future = self.pending_responses.pop(key, None)
+        if future and not future.done():
+            if response.error:
+                future.set_exception(ClientError(message=response.data))
+            else:
+                future.set_result(response.data)
+
+    async def subscribe(self, session_id: str) -> AsyncGenerator[ClientRequest, None]:
+        if session_id in self.sessions:
+            raise SessionAlreadyExists()
+
+        queue = asyncio.Queue()
+        self.sessions[session_id] = queue
+
+        try:
+            while True:
+                yield await queue.get()
+        finally:
+            del self.sessions[session_id]
+            # Eliminar todas las respuestas pendientes de esta sesión
+            keys_to_remove = [key for key in self.pending_responses if key[0] == session_id]
+            for key in keys_to_remove:
+                future = self.pending_responses.pop(key)
+                if not future.done():
+                    future.set_exception(SessionDoesNotExist())
+
+# Instanciación del broker
+broker = SessionBroker()
+
+# Funciones y rutas de la API
 
 async def _receive(session_id: str) -> None:
+    while True:
+        try:
+            message = await websocket.receive()
+            response_data = json.loads(message)
+            client_response = ClientResponse(**response_data)
+            app.logger.info(f"{websocket.remote_addr} - RESPONSE - {session_id} - {json.dumps(asdict(client_response))}")
+            await broker.receive_response(session_id, client_response)
+        except Exception as e:
+            app.logger.error(f"Error al recibir respuesta: {e}")
+            break
+
+@app.get("/wss/status")
+@validate_response(Status)
+async def status() -> Status:
     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')
+        version = importlib.metadata.version('serverwitch-api')
+    except importlib.metadata.PackageNotFoundError:
+        version = "unknown"
+    return Status(status="OK", version=version)
+
+@app.websocket('/wss/session')
+async def session_handler():
+    session_id = secrets.token_hex()
+    app.logger.info(f"{websocket.remote_addr} - NEW SESSION - {session_id}")
+    session_message = Session(session_id=session_id)
+    await websocket.send(json.dumps(asdict(session_message)))
+
+    task = asyncio.create_task(_receive(session_id))
+    try:
+        async for request_data in broker.subscribe(session_id):
+            app.logger.info(f"{websocket.remote_addr} - REQUEST - {session_id} - {json.dumps(asdict(request_data))}")
+            await websocket.send(json.dumps(asdict(request_data)))
+    except SessionAlreadyExists:
+        error_response = ErrorResponse(error="Session already exists.")
+        await websocket.send(json.dumps(asdict(error_response)))
+    finally:
+        task.cancel()
+        try:
+            await task
+        except asyncio.CancelledError:
+            pass
+
+@app.post('/wss/command')
 @validate_request(Command)
 @validate_response(CommandResponse, 200)
 @validate_response(ErrorResponse, 500)
-async def command(data: Command) -> Tuple[CommandResponse | ErrorResponse, int]:
+async def command(data: Command) -> Tuple[Union[CommandResponse, ErrorResponse], int]:
     try:
-        response = CommandResponse(**await broker.send_request(
+        response_data = await broker.send_request(
             data.session_id,
             {'action': 'command', 'command': data.command},
-            timeout=TIMEOUT))
+            timeout=TIMEOUT
+        )
+        response = CommandResponse(**response_data)
         return response, 200
     except SessionDoesNotExist:
         app.logger.warning(f"{request.remote_addr} - INVALID SESSION ID - {repr(data.session_id)}")
-        return ErrorResponse('Session does not exist.'), 500
+        return ErrorResponse(error='Session does not exist.'), 500
     except ClientError as e:
-        return ErrorResponse(e.message), 500
+        return ErrorResponse(error=e.message), 500
     except asyncio.TimeoutError:
-        return ErrorResponse('Timeout when waiting for client.'), 500
+        return ErrorResponse(error='Timeout when waiting for client.'), 500
 
-@app.post('/read')
+@app.post('/wss/read')
 @validate_request(Read)
 @validate_response(ReadResponse, 200)
 @validate_response(ErrorResponse, 500)
-async def read(data: Read) -> Tuple[ReadResponse | ErrorResponse, int]:
+async def read(data: Read) -> Tuple[Union[ReadResponse, ErrorResponse], int]:
     try:
-        response = ReadResponse(**await broker.send_request(
+        response_data = await broker.send_request(
             data.session_id,
             {'action': 'read', 'path': data.path},
-            timeout=TIMEOUT))
+            timeout=TIMEOUT
+        )
+        response = ReadResponse(**response_data)
         return response, 200
     except SessionDoesNotExist:
         app.logger.warning(f"{request.remote_addr} - INVALID SESSION ID - {repr(data.session_id)}")
-        return ErrorResponse('Session does not exist.'), 500
+        return ErrorResponse(error='Session does not exist.'), 500
     except ClientError as e:
-        return ErrorResponse(e.message), 500
+        return ErrorResponse(error=e.message), 500
     except asyncio.TimeoutError:
-        return ErrorResponse('Timeout when waiting for client.'), 500
+        return ErrorResponse(error='Timeout when waiting for client.'), 500
 
-@app.post('/write')
+@app.post('/wss/write')
 @validate_request(Write)
 @validate_response(WriteResponse, 200)
 @validate_response(ErrorResponse, 500)
-async def write(data: Write) -> Tuple[WriteResponse | ErrorResponse, int]:
+async def write(data: Write) -> Tuple[Union[WriteResponse, ErrorResponse], int]:
     try:
-        response = WriteResponse(**await broker.send_request(
+        response_data = await broker.send_request(
             data.session_id,
             {'action': 'write', 'path': data.path, 'content': data.content},
-            timeout=TIMEOUT))
+            timeout=TIMEOUT
+        )
+        response = WriteResponse(**response_data)
         return response, 200
     except SessionDoesNotExist:
         app.logger.warning(f"{request.remote_addr} - INVALID SESSION ID - {repr(data.session_id)}")
-        return ErrorResponse('Session does not exist.'), 500
+        return ErrorResponse(error='Session does not exist.'), 500
     except ClientError as e:
-        return ErrorResponse(e.message), 500
+        return ErrorResponse(error=e.message), 500
     except asyncio.TimeoutError:
-        return ErrorResponse('Timeout when waiting for client.'), 500
+        return ErrorResponse(error='Timeout when waiting for client.'), 500
 
-@app.route("/health")
-async def health_check():
-    return {"status": "healthy"}
+# Rutas para servir archivos estáticos y OpenAPI
+@app.route('/static/<path:path>')
+async def send_static(path):
+    return await send_from_directory('static', path)
 
-def run():
-    app.run(host='0.0.0.0', port=7860)
+@app.route('/openapi.yaml')
+async def openapi_spec():
+    return await send_from_directory('.', 'openapi.yaml')
 
-if __name__ == "__main__":
-    run()
+# No se necesita ejecutar nada aquí, ya que Hypercorn se encargará de iniciar la aplicación