From 0c1a65570be7f34f12a35da45669676f4479abd4 Mon Sep 17 00:00:00 2001
From: Stas Medvedev <medvedevsa97@gmail.com>
Date: Wed, 12 Jun 2024 16:56:07 +0300
Subject: add utils.restricted_exec add app.tasks

---
 app/main.py | 28 ++++++++++++++++++++++++++++
 1 file changed, 28 insertions(+)

(limited to 'app/main.py')

diff --git a/app/main.py b/app/main.py
index 59e7a8b..b90fef9 100644
--- a/app/main.py
+++ b/app/main.py
@@ -1,11 +1,15 @@
 from typing import Annotated
+import asyncio
+from datetime import datetime, timedelta
 
 from fastapi import FastAPI, Request, Depends
 from fastapi.responses import HTMLResponse
 from fastapi.staticfiles import StaticFiles
 from starlette.templating import Jinja2Templates
+from pydantic import BaseModel
 
 from utils import get_avatar_urls, get_client_geo
+from app.tasks import restricted_exec_task
 
 templates = Jinja2Templates(directory="templates")
 
@@ -33,3 +37,27 @@ async def client_addr(
         "partials/client_geo.html",
         {"request": request, "client_geo": client_geo},
     )
+
+
+class RestrictedExecBase(BaseModel):
+    code: str
+
+
+@app.post('/restricted_exec')
+async def post_restricted_exec(
+    body: RestrictedExecBase
+):
+    result = restricted_exec_task.delay(body.code)
+
+    start_time = datetime.now()
+    while datetime.now() - start_time < timedelta(seconds=5):
+        if result.ready():
+            return result.get()
+
+        await asyncio.sleep(0.1)
+
+    result.revoke(terminate=True, signal='SIGTERM')
+    return [
+        'Execution timeout, task revoked',
+        result.status
+    ]
-- 
cgit v1.2.3