From 0c1a65570be7f34f12a35da45669676f4479abd4 Mon Sep 17 00:00:00 2001 From: Stas Medvedev Date: Wed, 12 Jun 2024 16:56:07 +0300 Subject: add utils.restricted_exec add app.tasks --- app/main.py | 28 ++++++++++++++++++++++++++++ app/tasks.py | 17 +++++++++++++++++ 2 files changed, 45 insertions(+) create mode 100644 app/tasks.py (limited to 'app') 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 + ] diff --git a/app/tasks.py b/app/tasks.py new file mode 100644 index 0000000..bba3fae --- /dev/null +++ b/app/tasks.py @@ -0,0 +1,17 @@ +from celery import Celery + +from utils import restricted_exec + +app = Celery( + 'tasks', + broker='redis://localhost:6379/0', + backend='redis://localhost:6379/0', + task_send_sent_event=True, + worker_send_task_events=True, + worker_enable_remote_control=True, +) + + +@app.task +def restricted_exec_task(code: str): + return restricted_exec.getoutput(code) -- cgit v1.2.3