Add self restarts for workstations

backend/ng/containers/models/IndvidualContainer.py

@@ -1,10 +1,15 @@
 from CTFd.models import db
 from CTFd.utils import get_app_config
 import docker
+import redis_lock
 from typing import TypedDict
 from ..constants import DOCKER_RUNNING, DOCKER_BRIDGE, DOCKER_MEM_REGEX
 from ..utils.get_client import get_client
+from ...core import BusinessLogicError
 from .ContainerInstance import ContainerInstance
+from ..utils.redis import get_redis_client
+
+LOCK_EXPIRE_SECONDS = 5*60
 
 class SerializedIndvidualContainerInfo(TypedDict):
     id: int
@@ -32,6 +37,8 @@ def get_indvidual_container_by_dockerid(cls, docker_id: str):
 
     @classmethod
     def create_indvidual_container(cls, user_id: int, commit: bool = True):
+        redis_client = get_redis_client(3)
+
         db_exists = cls.query.filter_by(user=user_id).first()
         DOCKER_HOST = get_app_config("DOCKER_HOST")
         client = get_client(DOCKER_HOST)
@@ -41,9 +48,14 @@ def create_indvidual_container(cls, user_id: int, commit: bool = True):
             try:
                 ctr = client.get_running(db_exists.dockerid)
             except docker.errors.NotFound:
-                ctr = cls.run_container(client, container_name)
-                db_exists.dockerid = ctr.id
-                db.session.commit()
+                lock = redis_lock.Lock(redis_client, cls.render_lock_key(user_id), expire=LOCK_EXPIRE_SECONDS)
+                if lock.acquire(blocking=False):
+                    ctr = cls.run_container(client, container_name)
+                    db_exists.dockerid = ctr.id
+                    db.session.commit()
+                    lock.release()
+                else:
+                     raise BusinessLogicError("Workspace is already being started/reset") from None
 
             return db_exists
 
@@ -61,25 +73,35 @@ def create_indvidual_container(cls, user_id: int, commit: bool = True):
             return indv
 
         except docker.errors.NotFound:
-            ctr = cls.run_container(client, container_name)
+            lock = redis_lock.Lock(redis_client, cls.render_lock_key(user_id), expire=LOCK_EXPIRE_SECONDS)
+            if lock.acquire(blocking=False):
+                ctr = cls.run_container(client, container_name)
 
-            indvidual_container = cls(
-                user=user_id,
-                hostip=DOCKER_HOST,
-                dockerid=ctr.id,
-            )
+                indvidual_container = cls(
+                    user=user_id,
+                    hostip=DOCKER_HOST,
+                    dockerid=ctr.id,
+                )
+
+                db.session.add(indvidual_container)
+                if commit:
+                    db.session.commit()
+                lock.release()
+            else:
+                 raise BusinessLogicError("Workspace is already being started/reset") from None
 
-            db.session.add(indvidual_container)
-            if commit:
-                db.session.commit()
             return indvidual_container
 
     @staticmethod
     def render_container_name(user_id) -> str:
         return f"{user_id}-indv"
 
     @staticmethod
-    def run_container(client, container_name):
+    def render_lock_key(user_id) -> str:
+        return f"{user_id}-vnc-lock"
+
+    @staticmethod
+    def run_container(client, container_name, lock=None):
         NOVNC_CONTAINER = get_app_config("NOVNC_CONTAINER")
 
         NOVNC_RAM = get_app_config("NOVNC_RAM", "4g")
@@ -181,55 +203,80 @@ def get_current_challenge(self) -> int | None:
             return None
 
     def restart(self):
-        client = get_client(self.hostip)
-        try:
-            ctr = client.containers.get(self.dockerid)
-            ctr.restart()
-        except docker.errors.NotFound as exc:
-            raise ValueError("Container not found please recycle") from exc
+        redis_client = get_redis_client(3)
+        lock = redis_lock.Lock(redis_client, self.render_lock_key(self.user), expire=LOCK_EXPIRE_SECONDS)
+        if lock.acquire(blocking=False):
+            client = get_client(self.hostip)
+            try:
+                ctr = client.containers.get(self.dockerid)
+                ctr.restart()
+            except docker.errors.NotFound as exc:
+                raise ValueError("Container not found please recycle") from exc
+            finally:
+                lock.release()
+        else:
+            raise BusinessLogicError("Workspace is already being started/reset")
 
     def recycle(self):
-        client = get_client(self.hostip)
-        try:
-            ctr = client.containers.get(self.dockerid)
-            if ctr.status == DOCKER_RUNNING:
-                ctr.kill()
+        redis_client = get_redis_client(3)
+        lock = redis_lock.Lock(redis_client, self.render_lock_key(self.user), expire=LOCK_EXPIRE_SECONDS)
+        if lock.acquire(blocking=False):
+            client = get_client(self.hostip)
+            try:
+                ctr = client.containers.get(self.dockerid)
+                if ctr.status == DOCKER_RUNNING:
+                    ctr.kill()
 
-            ctr.remove()
+                ctr.remove()
 
-        except docker.errors.NotFound:
-            pass
+            except docker.errors.NotFound:
+                pass
 
-        finally:
-            container_name = self.render_container_name(self.user)
-            new_ctr = self.run_container(client, container_name)
+            finally:
+                container_name = self.render_container_name(self.user)
+                new_ctr = self.run_container(client, container_name)
 
-            self.dockerid = new_ctr.id
-            db.session.commit()
+                self.dockerid = new_ctr.id
+                db.session.commit()
+                lock.release()
+        else:
+            raise BusinessLogicError("Workspace is already being started/reset")
 
     def stop(self):
-        client = get_client(self.hostip)
-        try:
-            ctr = client.containers.get(self.dockerid)
-            if ctr.status == DOCKER_RUNNING:
-                ctr.stop(timeout=5)
-        except docker.errors.NotFound:
-            pass
+        redis_client = get_redis_client(3)
+        lock = redis_lock.Lock(redis_client, self.render_lock_key(self.user), expire=LOCK_EXPIRE_SECONDS)
+        if lock.acquire(blocking=False):
+            client = get_client(self.hostip)
+            try:
+                ctr = client.containers.get(self.dockerid)
+                if ctr.status == DOCKER_RUNNING:
+                    ctr.stop(timeout=5)
+            except docker.errors.NotFound:
+                pass
+            lock.release()
+        else:
+            raise BusinessLogicError("Workspace is already being started/reset")
 
     def delete(self, commit=True):
-        client = get_client(self.hostip)
-        try:
-            ctr = client.containers.get(self.dockerid)
-            ctr.remove(force=True)
-        except docker.errors.NotFound:
+        redis_client = get_redis_client(3)
+        lock = redis_lock.Lock(redis_client, self.render_lock_key(self.user), expire=LOCK_EXPIRE_SECONDS)
+        if lock.acquire(blocking=False):
+            client = get_client(self.hostip)
             try:
-                ctr = client.containers.get(self.render_container_name(self.user))
+                ctr = client.containers.get(self.dockerid)
                 ctr.remove(force=True)
             except docker.errors.NotFound:
-                pass
-        db.session.delete(self)
-        if commit:
-            db.session.commit()
+                try:
+                    ctr = client.containers.get(self.render_container_name(self.user))
+                    ctr.remove(force=True)
+                except docker.errors.NotFound:
+                    pass
+            db.session.delete(self)
+            if commit:
+                db.session.commit()
+            lock.release()
+        else:
+            raise BusinessLogicError("Workspace is already being started/reset")
 
     def get_status(self) -> str:
         client = get_client(self.hostip)

backend/ng/containers/routes/routes.py

@@ -1,6 +1,10 @@
+from datetime import datetime
 from flask_restx import Namespace, Resource
 from ..controllers.vnc import forward_vnc
 from ..controllers.get_current_connected_challenge import get_current_connected_challenge
+from ...core import BusinessLogicError
+from ..models.IndvidualContainer import IndvidualContainer
+from ...core.utils.rate_limit import limiter
 
 
 from ...core.middleware import (
@@ -41,3 +45,28 @@ def get(self, current_user):
         current_chall = get_current_connected_challenge(current_user.id)
         return success_response(current_chall)
 
+
+@container_namespace.route("/me/restart")
+class WorkspaceRestart(Resource):
+    @container_namespace.doc(
+        description="Restart users workspace",
+        responses={
+            200: "Sucess",
+            400: "Bad request"
+        }
+    )
+    @limiter.limit("1 per 5 minutes")
+    @user_endpoint()
+    def post(self, current_user):
+        teams = current_user.get_teams()
+        now = datetime.utcnow()
+        # Checks if user is in an event they can activley participate in
+        team_stats = [
+            tm.start_timestamp is not None and ((tm.end_time is not None and tm.end_time > now) or tm.end_time is None) for tm in teams
+        ]
+        if True in team_stats:
+            ctr = IndvidualContainer.get_user_indvidual_container(current_user.id)
+            ctr.restart()
+            return success_response(True)
+        else:
+            raise BusinessLogicError("Must be apart of an active event to restart your workspace")

frontend/src/hooks/container.ts

@@ -94,3 +94,9 @@ export function connectDeployment(challengeId: number, teamId: number) {
     method : 'POST',
   });
 }
+
+export function restartWorkspace() {
+  return apiMutation('/container/me/restart', undefined, {
+    method : 'POST',
+  });
+}

frontend/src/routes/profile/WorkspaceRestartModal.tsx

@@ -0,0 +1,29 @@
+import { COLOR_NEGATIVE } from '@/constants';
+import { restartWorkspace } from '@/hooks/container';
+import { Button, Text } from '@radix-ui/themes';
+import Modal from 'components/Modal';
+import { TbReload } from 'react-icons/tb';
+
+export default function WorkspaceRestartModal() {
+  return (
+    <Modal
+      title="Restart Workspace"
+      description=""
+      submitVerb="Restart"
+      submitColor={COLOR_NEGATIVE}
+      onSubmit={async () => restartWorkspace()}
+      trigger={(
+        <Button
+          color={COLOR_NEGATIVE}
+        >
+          <TbReload />
+          Restart
+        </Button>
+      )}
+    >
+      <Text color="gray">
+        This will restart your workspace. If you need your workspace completely reset please contact support.
+      </Text>
+    </Modal>
+  );
+}

frontend/src/routes/profile/index.tsx

@@ -18,6 +18,7 @@ import {
 } from 'lodash';
 import { useState } from 'react';
 import SponsorImageCard from './SponsorImageCard';
+import WorkspaceRestartModal from './WorkspaceRestartModal';
 
 export default function Profile() {
   const [ isEditing, setIsEditing ] = useState<boolean>(false);
@@ -84,6 +85,9 @@ export default function Profile() {
             </>
           )
         )}
+
+        <Heading size="4" as="h2" className="pt-4">Workspace</Heading>
+        <WorkspaceRestartModal />
       </Container>
     </>
   );