From 4fe45d443092ff39056eb986198de8ea96b819d0 Mon Sep 17 00:00:00 2001
From: Timothy Jaeryang Baek <tim@openwebui.com>
Date: Sun, 8 Jun 2025 20:26:07 +0400
Subject: [PATCH] refac/security: python code format endpoint

---
 backend/open_webui/routers/utils.py         |  2 +-
 scripts/prepare-pyodide.js                  |  3 +-
 src/lib/components/common/CodeEditor.svelte | 76 ++++++++++++++++++++-
 3 files changed, 77 insertions(+), 4 deletions(-)

diff --git a/backend/open_webui/routers/utils.py b/backend/open_webui/routers/utils.py
index b64adafb4..0e6768a67 100644
--- a/backend/open_webui/routers/utils.py
+++ b/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}
diff --git a/scripts/prepare-pyodide.js b/scripts/prepare-pyodide.js
index 70f3cf5c6..664683a30 100644
--- a/scripts/prepare-pyodide.js
+++ b/scripts/prepare-pyodide.js
@@ -12,7 +12,8 @@ const packages = [
 	'sympy',
 	'tiktoken',
 	'seaborn',
-	'pytz'
+	'pytz',
+	'black'
 ];
 
 import { loadPyodide } from 'pyodide';
diff --git a/src/lib/components/common/CodeEditor.svelte b/src/lib/components/common/CodeEditor.svelte
index 5afbf536d..32b725ee2 100644
--- a/src/lib/components/common/CodeEditor.svelte
+++ b/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({