Merge pull request #130 from PoCInnovation/80-distribox-slave

feat: Distribox Slave #80

Author: lg-epitech

.env.example

@@ -35,3 +35,13 @@ GUACD_PORT=4822
 # VNC_HOST: address guacd uses to reach VNC. 127.0.0.1 when guacd uses host
 # networking (default); change to host.docker.internal for bridged guacd.
 VNC_HOST=127.0.0.1
+
+# Master/Slave mode
+# Set to "master" (default) or "slave"
+DISTRIBOX_MODE=master
+# Slave-only: URL of the master node (e.g. http://192.168.1.10:8080)
+MASTER_URL=
+# Slave-only: API key generated when registering this slave on the master
+SLAVE_API_KEY=
+# Slave-only: port for the slave API (default 8081)
+SLAVE_PORT=8081

backend/Dockerfile

@@ -13,8 +13,8 @@ RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
 
 FROM builder as dev-stage
 COPY ./app /code/app
-CMD uvicorn app.main:app --host 0.0.0.0 --port ${BACKEND_PORT} --reload
+CMD ["sh", "-c", "exec uvicorn app.main:app --host 0.0.0.0 --port ${BACKEND_PORT} --reload"]
 
 FROM builder as production-stage
 COPY ./app /code/app
-CMD fastapi run app/main.py --port ${BACKEND_PORT}
+CMD ["sh", "-c", "exec fastapi run app/main.py --port ${BACKEND_PORT}"]

backend/app/core/config.py

@@ -72,6 +72,13 @@ def init_db():
                         "ALTER TABLE vms ADD COLUMN keyboard_layout VARCHAR"
                     )
                 )
+            if "slave_id" not in vm_columns:
+                conn.execute(
+                    text(
+                        "ALTER TABLE vms ADD COLUMN slave_id UUID "
+                        "REFERENCES slaves(id)"
+                    )
+                )
 
         if "events" in inspector.get_table_names():
             event_columns = {
@@ -116,3 +123,14 @@ def get_connection(cls):
 GUACD_HOST = get_env_or_default("GUACD_HOST", "host.docker.internal")
 GUACD_PORT = int(get_env_or_default("GUACD_PORT", "4822"))
 VNC_HOST = get_env_or_default("VNC_HOST", "127.0.0.1")
+VNC_LISTEN = get_env_or_default("VNC_LISTEN", "127.0.0.1")
+
+# Master/Slave mode: "master" (default) or "slave"
+DISTRIBOX_MODE = get_env_or_default("DISTRIBOX_MODE", "master")
+
+# Slave-specific config (only used when DISTRIBOX_MODE=slave)
+MASTER_URL = get_env_or_default("MASTER_URL", "")
+SLAVE_API_KEY = get_env_or_default("SLAVE_API_KEY", "")
+
+# Virtualization type: "kvm" (default, hardware accel) or "qemu" (software emulation)
+VIRT_TYPE = get_env_or_default("VIRT_TYPE", "kvm")

backend/app/core/policies.py

@@ -137,6 +137,22 @@
         "policy": "events:delete",
         "description": "Allows the user to delete events and their VMs.",
     },
+    {
+        "policy": "slaves:get",
+        "description": "Allows the user to list slave nodes.",
+    },
+    {
+        "policy": "slaves:getById",
+        "description": "Allows the user to fetch a slave node by id.",
+    },
+    {
+        "policy": "slaves:create",
+        "description": "Allows the user to register a new slave node.",
+    },
+    {
+        "policy": "slaves:delete",
+        "description": "Allows the user to unregister a slave node.",
+    },
 ]
 
 VALID_POLICIES = {entry["policy"] for entry in POLICIES}

backend/app/core/xml_builder.py

@@ -1,6 +1,7 @@
 from lxml import etree
 from app.models.vm import VmCreateXML
 from app.core.constants import VMS_DIR
+from app.core.config import VIRT_TYPE
 
 LAYOUT_TO_KEYMAP = {
     "en-us-qwerty": "en-us",
@@ -33,7 +34,7 @@
 
 def build_xml(vm_read: VmCreateXML):
 
-    domain = etree.Element("domain", type="kvm")
+    domain = etree.Element("domain", type=VIRT_TYPE)
 
     etree.SubElement(domain, "name").text = str(vm_read.id)
     # Frontend sends memory in GiB; libvirt XML expects MiB here.
@@ -50,7 +51,8 @@ def build_xml(vm_read: VmCreateXML):
     for feature in ["acpi", "apic", "pae"]:
         etree.SubElement(features, feature)
 
-    etree.SubElement(domain, "cpu", mode="host-passthrough")
+    if VIRT_TYPE == "kvm":
+        etree.SubElement(domain, "cpu", mode="host-passthrough")
 
     etree.SubElement(domain, "clock", offset="utc")
 
@@ -86,11 +88,12 @@ def build_xml(vm_read: VmCreateXML):
     etree.SubElement(iface, "source", network="default")
     etree.SubElement(iface, "model", type="virtio")
 
+    from app.core.config import VNC_LISTEN
     vnc_attrs = {
         "type": "vnc",
         "port": "-1",
         "autoport": "yes",
-        "listen": "127.0.0.1",
+        "listen": VNC_LISTEN,
     }
     if vm_read.keyboard_layout:
         keymap = LAYOUT_TO_KEYMAP.get(vm_read.keyboard_layout)

backend/app/main.py

@@ -7,13 +7,15 @@
 from fastapi.middleware.cors import CORSMiddleware
 from sqlmodel import Session, select
 from app.core.policies import DISTRIBOX_ADMIN_POLICY
-from app.routes import vm, image, host, auth, user_management, tunnel, event
+from app.routes import vm, image, host, auth, user_management, tunnel, event, slave
+from app.routes import slave_agent
 from app.orm.user import UserORM
 from app.orm.vm_credential import VmCredentialORM  # noqa: F401
 from app.orm.event import EventORM, EventParticipantORM  # noqa: F401
 from app.orm.user_settings import UserSettingsORM  # noqa: F401
+from app.orm.slave import SlaveORM  # noqa: F401
 from app.utils.auth import hash_password
-from app.core.config import engine, get_env_or_default, init_db
+from app.core.config import engine, get_env_or_default, init_db, DISTRIBOX_MODE, QEMUConfig
 from app.utils.crypto import encrypt_secret, is_encrypted_secret
 from app.services.vm_service import VmService
 
@@ -23,7 +25,6 @@
 
 frontend_url = get_env_or_default("FRONTEND_URL", "http://localhost:3000")
 
-# CORS configuration
 app.add_middleware(
     CORSMiddleware,
     allow_origins=[frontend_url],
@@ -33,21 +34,61 @@
 )
 
 
+@app.on_event("shutdown")
+async def shutdown_event():
+    if DISTRIBOX_MODE != "slave":
+        return
+
+    _stop_all_local_vms()
+
+    import httpx
+    from app.core.config import MASTER_URL, SLAVE_API_KEY
+
+    if not MASTER_URL or not SLAVE_API_KEY:
+        return
+    try:
+        async with httpx.AsyncClient(timeout=5.0) as client:
+            await client.post(
+                f"{MASTER_URL}/slaves/shutdown",
+                headers={"X-Slave-Token": SLAVE_API_KEY},
+            )
+        logger.info("Shutdown notification sent to master")
+    except Exception:
+        logger.warning("Failed to notify master of shutdown")
+
+
+def _stop_all_local_vms():
+    try:
+        conn = QEMUConfig.get_connection()
+        domains = conn.listAllDomains(0)
+        for domain in domains:
+            if domain.isActive():
+                try:
+                    domain.destroy()
+                    logger.info("Destroyed VM %s", domain.name())
+                except Exception:
+                    logger.warning("Failed to destroy VM %s", domain.name())
+    except Exception:
+        logger.warning("Failed to stop local VMs during shutdown")
+
+
 @app.on_event("startup")
 async def startup_event():
-    """Initialize database and create default admin user if it doesn't exist."""
-    init_db()  # This might be temporary
+    init_db()
+
+    if DISTRIBOX_MODE == "slave":
+        print(f"✓ Starting in SLAVE mode")
+        asyncio.create_task(_slave_heartbeat_loop())
+        return
 
     admin_username = get_env_or_default("ADMIN_USERNAME", "admin")
     admin_password = get_env_or_default("ADMIN_PASSWORD", "admin")
 
     with Session(engine) as session:
-        # Check if admin exists
         statement = select(UserORM).where(UserORM.username == admin_username)
         admin = session.exec(statement).first()
 
         if not admin:
-            # Create admin user
             admin = UserORM(
                 username=admin_username,
                 hashed_password=hash_password(admin_password),
@@ -93,6 +134,8 @@ async def startup_event():
             )
 
     asyncio.create_task(_enforce_event_deadlines())
+    asyncio.create_task(_check_stale_slaves())
+    print(f"✓ Starting in MASTER mode")
 
 
 async def _enforce_event_deadlines():
@@ -133,6 +176,54 @@ async def _enforce_event_deadlines():
         await asyncio.sleep(30)
 
 
+async def _check_stale_slaves():
+    while True:
+        try:
+            from app.services.slave_service import SlaveService
+            SlaveService.check_stale_slaves()
+        except Exception:
+            logger.exception("Error in stale slave check")
+        await asyncio.sleep(30)
+
+
+async def _slave_heartbeat_loop():
+    import httpx
+    from app.core.config import MASTER_URL, SLAVE_API_KEY
+    from app.services.host_service import HostService
+    import psutil
+
+    if not MASTER_URL or not SLAVE_API_KEY:
+        logger.warning(
+            "MASTER_URL or SLAVE_API_KEY not set, skipping heartbeat"
+        )
+        return
+
+    while True:
+        try:
+            host_info = HostService.get_host_info()
+            mem = psutil.virtual_memory()
+            heartbeat = {
+                "total_cpu": psutil.cpu_count() or 0,
+                "total_mem": round(mem.total / 2**30),
+                "total_disk": round(host_info.disk.total),
+                "available_cpu": round(
+                    100.0 - host_info.cpu.percent_used_total, 2
+                ),
+                "available_mem": round(host_info.mem.available, 2),
+                "available_disk": round(host_info.disk.available, 2),
+            }
+            async with httpx.AsyncClient(timeout=10.0) as client:
+                await client.post(
+                    f"{MASTER_URL}/slaves/heartbeat",
+                    json=heartbeat,
+                    headers={"X-Slave-Token": SLAVE_API_KEY},
+                )
+            logger.debug("Heartbeat sent to master")
+        except Exception:
+            logger.exception("Failed to send heartbeat to master")
+        await asyncio.sleep(30)
+
+
 @app.exception_handler(HTTPException)
 async def http_exception_handler(_, exc: HTTPException):
     return JSONResponse(
@@ -148,10 +239,14 @@ async def general_exception_handler(_, exc: Exception):
         content={"detail": str(exc)}
     )
 
-app.include_router(auth.router, prefix="/auth", tags=["auth"])
-app.include_router(user_management.router, tags=["users"])
-app.include_router(vm.router, prefix="/vms", tags=["vms"])
-app.include_router(image.router, prefix="/images", tags=["images"])
-app.include_router(host.router, prefix="/host", tags=["host"])
-app.include_router(tunnel.router, tags=["tunnel"])
-app.include_router(event.router, prefix="/events", tags=["events"])
+if DISTRIBOX_MODE == "slave":
+    app.include_router(slave_agent.router, tags=["slave-agent"])
+else:
+    app.include_router(auth.router, prefix="/auth", tags=["auth"])
+    app.include_router(user_management.router, tags=["users"])
+    app.include_router(vm.router, prefix="/vms", tags=["vms"])
+    app.include_router(image.router, prefix="/images", tags=["images"])
+    app.include_router(host.router, prefix="/host", tags=["host"])
+    app.include_router(tunnel.router, tags=["tunnel"])
+    app.include_router(event.router, prefix="/events", tags=["events"])
+    app.include_router(slave.router, prefix="/slaves", tags=["slaves"])

backend/app/models/host.py

@@ -1,3 +1,5 @@
+from typing import Optional
+from uuid import UUID
 from pydantic import BaseModel
 from app.models.resources import MemoryInfoBase, DiskInfoBase, CPUInfoBase
 
@@ -18,3 +20,22 @@ class HostInfoBase(BaseModel):
     disk: DiskInfoHost
     mem: MemoryInfoHost
     cpu: CPUInfoHost
+
+
+class NodeHostInfo(BaseModel):
+    node_id: Optional[UUID] = None
+    node_name: str
+    host_info: HostInfoBase
+
+
+class ClusterHostInfo(BaseModel):
+    nodes: list[NodeHostInfo]
+    totals: "ClusterTotals"
+
+
+class ClusterTotals(BaseModel):
+    cpu_count: int
+    mem_total: float
+    mem_available: float
+    disk_total: float
+    disk_available: float

backend/app/models/slave.py

@@ -0,0 +1,43 @@
+from typing import Optional
+from pydantic import BaseModel, Field
+from uuid import UUID
+from datetime import datetime
+
+
+class SlaveCreate(BaseModel):
+    name: str = Field(min_length=1)
+    hostname: str = Field(min_length=1)
+    port: int = Field(default=8080, ge=1, le=65535)
+
+
+class SlaveRead(BaseModel):
+    id: UUID
+    name: str
+    hostname: str
+    port: int
+    api_key: str
+    status: str
+    created_at: datetime
+    last_heartbeat: Optional[datetime] = None
+    total_cpu: int
+    total_mem: int
+    total_disk: int
+    available_cpu: float
+    available_mem: float
+    available_disk: float
+
+
+class SlaveHeartbeat(BaseModel):
+    total_cpu: int
+    total_mem: int
+    total_disk: int
+    available_cpu: float
+    available_mem: float
+    available_disk: float
+
+
+class SlaveHostInfo(BaseModel):
+    """Response from a slave's /host/info endpoint."""
+    disk: dict
+    mem: dict
+    cpu: dict

backend/app/models/vm.py

@@ -19,10 +19,14 @@ class VmRead(VmBase):
     state: str
     ipv4: Optional[str]
     credentials_count: int = 0
+    slave_id: Optional[UUID] = None
+    slave_name: Optional[str] = None
 
 
 class VmCreate(VmBase):
     activate_at_start: bool
+    slave_id: Optional[UUID] = None
+    auto_place: bool = False
 
 
 class VmCreateXML(VmBase):