refac/security: python code format endpoint

backend/open_webui/routers/utils.py

@@ -33,7 +33,7 @@ class CodeForm(BaseModel):
 
 
 @router.post("/code/format")
-async def format_code(form_data: CodeForm, user=Depends(get_verified_user)):
+async def format_code(form_data: CodeForm, user=Depends(get_admin_user)):
     try:
         formatted_code = black.format_str(form_data.code, mode=black.Mode())
         return {"code": formatted_code}

scripts/prepare-pyodide.js

@@ -12,7 +12,8 @@ const packages = [
 	'sympy',
 	'tiktoken',
 	'seaborn',
-	'pytz'
+	'pytz',
+	'black'
 ];
 
 import { loadPyodide } from 'pyodide';

src/lib/components/common/CodeEditor.svelte

@@ -13,8 +13,11 @@
 
 	import { onMount, createEventDispatcher, getContext, tick } from 'svelte';
 
+	import PyodideWorker from '$lib/workers/pyodide.worker?worker';
+
 	import { formatPythonCode } from '$lib/apis/utils';
 	import { toast } from 'svelte-sonner';
+	import { user } from '$lib/stores';
 
 	const dispatch = createEventDispatcher();
 	const i18n = getContext('i18n');
@@ -113,13 +116,82 @@
 		return await language?.load();
 	};
 
+	let pyodideWorkerInstance = null;
+
+	const getPyodideWorker = () => {
+		if (!pyodideWorkerInstance) {
+			pyodideWorkerInstance = new PyodideWorker(); // Your worker constructor
+		}
+		return pyodideWorkerInstance;
+	};
+
+	// Generate unique IDs for requests
+	let _formatReqId = 0;
+
+	const formatPythonCodePyodide = (code) => {
+		return new Promise((resolve, reject) => {
+			const id = `format-${++_formatReqId}`;
+			let timeout;
+			const worker = getPyodideWorker();
+
+			const script = `
+import black
+print(black.format_str("""${code.replace(/\\/g, '\\\\').replace(/`/g, '\\`').replace(/"/g, '\\"')}""", mode=black.Mode()))
+`;
+
+			const packages = ['black'];
+
+			function handleMessage(event) {
+				const { id: eventId, stdout, stderr } = event.data;
+				if (eventId !== id) return; // Only handle our message
+				clearTimeout(timeout);
+				worker.removeEventListener('message', handleMessage);
+				worker.removeEventListener('error', handleError);
+
+				if (stderr) {
+					reject(stderr);
+				} else {
+					const formatted = stdout && typeof stdout === 'string' ? stdout.trim() : '';
+					resolve({ code: formatted });
+				}
+			}
+
+			function handleError(event) {
+				clearTimeout(timeout);
+				worker.removeEventListener('message', handleMessage);
+				worker.removeEventListener('error', handleError);
+				reject(event.message || 'Pyodide worker error');
+			}
+
+			worker.addEventListener('message', handleMessage);
+			worker.addEventListener('error', handleError);
+
+			// Send to worker
+			worker.postMessage({ id, code: script, packages });
+
+			// Timeout
+			timeout = setTimeout(() => {
+				worker.removeEventListener('message', handleMessage);
+				worker.removeEventListener('error', handleError);
+				try {
+					worker.terminate();
+				} catch {}
+				pyodideWorkerInstance = null;
+				reject('Execution Time Limit Exceeded');
+			}, 60000);
+		});
+	};
+
 	export const formatPythonCodeHandler = async () => {
 		if (codeEditor) {
-			const res = await formatPythonCode(localStorage.token, _value).catch((error) => {
+			const res = await (
+				$user?.role === 'admin'
+					? formatPythonCode(localStorage.token, _value)
+					: formatPythonCodePyodide(_value)
+			).catch((error) => {
 				toast.error(`${error}`);
 				return null;
 			});
-
 			if (res && res.code) {
 				const formattedCode = res.code;
 				codeEditor.dispatch({