ollama-plus/coderunner/index.ts

import http from "node:http";
import { spawn } from "node:child_process";

const PORT = Number(process.env.PORT || 8787);
const HOST = "0.0.0.0";

/**
 * @example const imgObj = { image: "gcc:14", installcommands: [aptUpdateAndInstall("gpp")], filename: "main.cpp", run: ... },
 */
const aptUpdateAndInstall = (pkgs: string[] | string) => {
	const istr = typeof pkgs === 'string' ? pkgs : pkgs.join(" ");

	return `apt-get update \
 && apt-get install -y --no-install-recommends ${istr} \
 && rm -rf /var/lib/apt/lists/*;`
}

// allow-listed images and run commands per language
const LANGS = {
	python: { image: "python:3.12-alpine", filename: "main.py", run: ["python", "main.py"] },
	node: { image: "node:22-alpine", filename: "main.mjs", run: ["node", "main.mjs"] },
	bun: { image: "oven/bun:1.2.2-alpine", filename: "main.ts", run: ["bun", "run", "main.ts"] },
	bash: { image: "alpine:3.20", filename: "main.sh", run: ["sh", "main.sh"] },
	ruby: { image: "ruby:3.3-alpine", filename: "main.rb", run: ["ruby", "main.rb"] },
	go: { image: "golang:1.22-alpine", filename: "main.go", run: ["go run main.go"] },
	rust: { image: "rust:1-alpine", filename: "main.rs", run: ["rustc -O main.rs -o main && ./main"] },
	java: { image: "eclipse-temurin:21-jdk", filename: "Main.java", run: ["javac Main.java && java Main"] },
	c: { image: "gcc:14", filename: "main.c", run: ["gcc -O2 main.c -o main.out && ./main.out"] },
	cpp: { image: "gcc:14", filename: "main.cpp", run: ["g++ -O2 main.cpp -o main.out && ./main.out"] },
};

type langObj = {
	image: string,
	filename: string,
	run: string[],
	installcommands?: string[]
}

type fileType = {
	path: string,
	content: string
}

// docker binary (or set DOCKER_BIN=podman)
const DOCKER_BIN = process.env.DOCKER_BIN || "docker";

// basic openapi for open webui
const OPENAPI = JSON.parse((await import('fs')).readFileSync('openapi.json'))

function sendJson(res, status, obj) {
	const body = JSON.stringify(obj);
	res.writeHead(status, { "content-type": "application/json; charset=utf-8" });
	res.end(body);
}

async function ensureImage(spec: langObj) {
	// note: 'docker run' has a --pull policy (missing|always|never),
	const { spawn } = await import("node:child_process"),
		DOCKER_BIN = process.env.DOCKER_BIN || "docker";

	// check if image exists locally
	const inspect = spawn(DOCKER_BIN, ["image", "inspect", spec.image], { stdio: "ignore" }),
		ok = await new Promise((r) => inspect.on("close", (c) => r(c === 0)));
	if (ok) return;

	// pull with a bigger timeout the first time (4 minutes)
	await new Promise((resolve) => {
		const child = spawn(DOCKER_BIN, ["image", "pull", "--quiet", spec.image]);
		let timer = setTimeout(() => {
			child.kill("SIGKILL");
			resolve(new Error("pull timeout"));
		}, 240_000);

		child.on("close", (code) => {
			clearTimeout(timer);
			resolve(code === 0 ? null : new Error(`pull failed: ${code}`));
		});
	}).then((err) => {
		if (err) throw err;
	});

	// 	// dependancies
	// 	if (spec.installcommands) {
	// 		const commandFull = spec.installcommands.join(" && ");
	// 		await new Promise((resolve) => {
	// 			const child = spawn(DOCKER_BIN, ["image", "pull", "--quiet", spec.image]);
	// 			let timer = setTimeout(() => {
	// 				child.kill("SIGKILL");
	// 				resolve(new Error("pull timeout"));
	// 			}, 240_000);

	// 			child.on("close", (code) => {
	// 				clearTimeout(timer);
	// 				resolve(code === 0 ? null : new Error(`pull failed: ${code}`));
	// 			});
	// 		}).then((err) => {
	// 			if (err) throw err;
	// 		})
	// 	}
}

async function runInContainer({ language, code, args = [], files = [] }: {
	language: string,
	code: string,
	args: string[],
	files: fileType[]
}) {
	if (!LANGS[language]) throw new Error(`language not allowed: ${language}`);
	const spec = LANGS[language];

	await ensureImage(spec);

	// build the Docker args
	const dockerArgs = [
		"run", "--rm",
		"--network=none", "--read-only",
		"--pids-limit=256",
		"--cpus=1", "--memory=512m",
		"--cap-drop=ALL", "--security-opt", "no-new-privileges",
		"--tmpfs", "/work:rw,exec,size=64m",
		"-w", "/work",
		"--pull=never",
		spec.image
	];

	// inside the container, write the files and run code
	const script = [
		// write the main file using base64
		`echo ${JSON.stringify(Buffer.from(code, "utf8").toString("base64"))} | base64 -d > ${spec.filename}`,
		// write any extra files using base64
		...files.flatMap((f) => [
			`mkdir -p "$(dirname "${f.path}")"`,
			`echo ${JSON.stringify(Buffer.from(f.content, "utf8").toString("base64"))} | base64 -d > "${f.path}"`,
		]),
		// run it
		`${spec.run.join(' ')} ${args.map((a) => JSON.stringify(a)).join(' ')}`
	].join('\n');

	dockerArgs.push("sh", "-lc", script);

	const result = await new Promise((resolve) => {
		const child = spawn(DOCKER_BIN, dockerArgs, { stdio: ["ignore", "pipe", "pipe"] });
		let stdout = "", stderr = "";

		const timer = setTimeout(() => {
			child.kill("SIGKILL");
			resolve({ stdout, stderr: stderr + "\n[killed: timeout]", exitCode: 137, timedOut: true });
		}, 25_000);

		child.stdout.on("data", (d) => { stdout += d.toString(); });
		child.stderr.on("data", (d) => { stderr += d.toString(); });
		child.on("close", (code) => {
			clearTimeout(timer);
			resolve({ stdout, stderr, exitCode: code ?? 1, timedOut: false });
		});
	});

	return result;
}

const server = http.createServer(async (req, res) => {
	try {
		console.debug(`recieved ${req.method} request on ${req.url}`);

		if (req.method === "GET" && req.url === "/openapi.json") {
			sendJson(res, 200, OPENAPI);
			return;
		}

		if (req.method === "GET" && req.url === "/") {
			res.writeHead(200);
			res.end("Ok");
			return;
		}

		if (req.method === "POST" && req.url === "/execute") {
			let body = "";
			req.on("data", (c) => { body += c.toString(); });
			req.on("end", async () => {
				try {
					const payload = JSON.parse(body || "{}");
					const out = await runInContainer(payload);
					sendJson(res, 200, out);
				} catch (e) {
					sendJson(res, 400, { error: String(e?.message || e) });
				}
			});
			return;
		}

		res.writeHead(404);
		res.end("not found");
	} catch (e) {
		sendJson(res, 500, { error: String(e?.message || e) });
	}
});

server.listen(PORT, HOST, () => {
	console.log(`[runner] listening on http://${HOST}:${PORT}`);
});