Revamp connection strategy

server.ts

@@ -2,11 +2,10 @@ import "dotenv/config";
 
 import { createServer } from "node:http";
 import { Socket } from "node:net";
-import { db, session } from "@/lib/drizzle/db";
 import { consola } from "consola";
-import { eq } from "drizzle-orm";
 import next from "next";
 import { WebSocketServer } from "ws";
+import docker from "@/lib/docker";
 const port = Number.parseInt(process.env.PORT as string) || 3000;
 const dev = process.env.NODE_ENV !== "production";
 if (process.argv.includes("--turbo")) {
@@ -20,40 +19,44 @@ const nextRequest = app.getRequestHandler();
 const nextUpgrade = app.getUpgradeHandler();
 const websockify = new WebSocketServer({ noServer: true });
 websockify.on("connection", async (ws, req) => {
-	const id = req.url?.split("/")[2];
-	if (!id) {
-		ws.close(1008, "Missing ID");
-		return;
-	}
-	const [{ vncPort }] = await db.select({ vncPort: session.vncPort }).from(session).where(eq(session.id, id)).limit(1);
-	if (!vncPort) {
-		ws.close(1008, "Session not found");
-		return;
-	}
-	const socket = new Socket();
-	socket.connect(vncPort, process.env.CONTAINER_HOST as string);
-	ws.on("message", (message: Uint8Array) => {
-		socket.write(message);
-	});
-	ws.on("close", (code, reason) => {
-		consola.info(
-			`✨ Stardust: Connection closed with code ${code} and ${reason ? `reason ${reason.toString()}` : "no reason"}`,
-		);
-		socket.end();
-	});
+	try {
+		const id = req.url?.split("/")[2];
+		if (!id) {
+			ws.close(1008, "Missing ID");
+			return;
+		}
+		const container = docker.getContainer(id);
+		const ip = await container
+			.inspect()
+			.then((data) => data.NetworkSettings.Networks[data.HostConfig.NetworkMode as string].IPAddress);
+		const socket = new Socket();
+		socket.connect(5901, ip);
+		ws.on("message", (message: Uint8Array) => {
+			socket.write(message);
+		});
+		ws.on("close", (code, reason) => {
+			consola.info(
+				`✨ Stardust: Connection closed with code ${code} and ${reason ? `reason ${reason.toString()}` : "no reason"}`,
+			);
+			socket.end();
+		});
 
-	socket.on("data", (data) => {
-		ws.send(data);
-	});
+		socket.on("data", (data) => {
+			ws.send(data);
+		});
 
-	socket.on("error", (err) => {
-		consola.warn(`✨ Stardust: ${err.message}`);
-		ws.close();
-	});
+		socket.on("error", (err) => {
+			consola.warn(`✨ Stardust: ${err.message}`);
+			ws.close();
+		});
 
-	socket.on("close", () => {
-		ws.close();
-	});
+		socket.on("close", () => {
+			ws.close();
+		});
+	} catch (error) {
+		ws.close(1008, "Server error");
+		consola.error(`✨ Stardust: ${(error as Error).message}`);
+	}
 });
 server.on("request", nextRequest);
 server.on("upgrade", async (req, socket, head) => {

src/app/api/session/[slug]/files/route.ts

@@ -6,14 +6,14 @@ import type { NextRequest } from "next/server";
 export async function GET(req: NextRequest, { params }: { params: { slug: string } }) {
 	const userSession = await auth();
 	const fileName = req.nextUrl.searchParams.get("name");
-	const { id, agentPort } = (await getSession(params.slug, userSession)) || {};
-	if (!id || !agentPort) {
+	const { id, ip } = (await getSession(params.slug, userSession)) || {};
+	if (!id || !ip) {
 		return Response.json(["not found"], { status: 404 });
 	}
-	await sessionRunning(agentPort);
+	await sessionRunning(ip);
 	if (fileName) {
 		try {
-			const download = await fetch(`http://${process.env.CONTAINER_HOST}:${agentPort}/files/download/${fileName}`);
+			const download = await fetch(`http://${ip}:6080/files/download/${fileName}`);
 			return new Response(download.body, {
 				headers: {
 					"Content-Disposition": `attachment; filename=${fileName}`,
@@ -25,17 +25,17 @@ export async function GET(req: NextRequest, { params }: { params: { slug: string
 			return Response.json({ error: "Download failed" }, { status: 500 });
 		}
 	}
-	return fetch(`http://${process.env.CONTAINER_HOST}:${agentPort}/files/list`);
+	return fetch(`http://${ip}:6080/files/list`);
 }
 export async function PUT(req: NextRequest, { params }: { params: { slug: string } }) {
 	const userSession = await auth();
 	const buffer = await req.arrayBuffer();
 	const fileName = req.nextUrl.searchParams.get("name");
-	const { id, agentPort } = (await getSession(params.slug, userSession)) || {};
+	const { id, ip } = (await getSession(params.slug, userSession)) || {};
 	if (!id) {
 		return Response.json({ error: "Not Found" }, { status: 404 });
 	}
-	await fetch(`http://${process.env.CONTAINER_HOST}:${agentPort}/files/upload/${fileName}`, {
+	await fetch(`http://${ip}:6080/files/upload/${fileName}`, {
 		method: "PUT",
 		body: Buffer.from(buffer),
 	});

src/app/api/session/[slug]/preview/route.ts

@@ -4,10 +4,10 @@ import { sessionRunning } from "@/lib/session/session-running";
 import type { NextRequest } from "next/server";
 export async function GET(_req: NextRequest, { params }: { params: { slug: string } }) {
 	const userSession = await auth();
-	const { agentPort } = (await getSession(params.slug, userSession)) || {};
-	if (!agentPort) {
+	const { ip } = (await getSession(params.slug, userSession)) || {};
+	if (!ip) {
 		return Response.json({ error: "Not Found" }, { status: 404 });
 	}
-	await sessionRunning(agentPort);
-	return fetch(`http://${process.env.CONTAINER_HOST}:${agentPort}/screenshot`);
+	await sessionRunning(ip);
+	return fetch(`http://${ip}:6080/screenshot`);
 }

src/app/api/session/[slug]/route.ts

@@ -18,10 +18,8 @@ export async function GET(_req: NextRequest, { params }: { params: { slug: strin
 	if (State.Paused) {
 		await container.unpause();
 	}
-	await sessionRunning(containerSession.agentPort);
-	const password = await fetch(`http://${process.env.CONTAINER_HOST}:${containerSession.agentPort}/password`).then(
-		(res) => res.text(),
-	);
+	await sessionRunning(containerSession.ip);
+	const password = await fetch(`http://${containerSession.ip}:6080/password`).then((res) => res.text());
 	return Response.json({
 		exists: true,
 		password,

src/lib/drizzle/schema.ts

@@ -1,5 +1,5 @@
 import { relations } from "drizzle-orm";
-import { bigint, boolean, integer, pgTable, text } from "drizzle-orm/pg-core";
+import { bigint, boolean, pgTable, text } from "drizzle-orm/pg-core";
 
 export const user = pgTable("User", {
 	email: text("email").notNull().unique(),
@@ -30,8 +30,6 @@ export const session = pgTable("Session", {
 	userId: text("userId")
 		.notNull()
 		.references(() => user.id),
-	vncPort: integer("vncPort").notNull(),
-	agentPort: integer("agentPort").notNull(),
 });
 export type SelectSession = typeof session.$inferSelect;
 export const sessionRelations = relations(session, ({ one }) => ({

src/lib/session/get-session.ts

@@ -1,3 +1,4 @@
+import docker from "@/lib/docker";
 import { db, session, user } from "@/lib/drizzle/db";
 import { and, eq } from "drizzle-orm";
 import type { Session } from "next-auth";
@@ -9,20 +10,26 @@ import type { Session } from "next-auth";
  */
 async function getSession(containerId: string, userSession: Session | null) {
 	const [containerSession] = await db.transaction(async (tx) => {
-		if (!userSession?.user) throw new Error("User not found");
 		const [{ userId }] = await tx
 			.select({
 				userId: user.id,
 			})
 			.from(user)
-			.where(eq(user.email, userSession.user.email as string));
+			.where(eq(user.email, userSession?.user?.email as string));
 		return tx
 			.select()
 			.from(session)
 			.where(and(eq(session.id, containerId), eq(session.userId, userId)));
 	});
 	if (!containerSession) return null;
-	return containerSession;
+	const ip = await docker
+		.getContainer(containerId)
+		.inspect()
+		.then((data) => data.NetworkSettings.Networks[data.HostConfig.NetworkMode as string].IPAddress);
+	return {
+		ip,
+		...containerSession,
+	};
 }
 
 export { getSession };

src/lib/session/index.ts

@@ -8,7 +8,6 @@ import { eq } from "drizzle-orm";
 import { revalidatePath } from "next/cache";
 import { auth } from "../auth";
 import { getSession } from "./get-session";
-const getRandomNumber = (min: number, max: number) => Math.floor(Math.random() * (max - min + 1)) + min;
 /**
  * Creates a new Stardust session
  * @param Image Docker image to use for making the session
@@ -18,19 +17,6 @@ async function createSession(Image: string) {
 	try {
 		const userSession = await auth();
 		if (!userSession?.user) throw new Error("User not found");
-		if (!process.env.DOCKER_PORT_RANGE) throw new Error("Docker port range not set");
-		const portsRange = process.env.DOCKER_PORT_RANGE.split("-").map(Number);
-		let vncPort: number = getRandomNumber(portsRange[0], portsRange[1]);
-		let agentPort: number = getRandomNumber(portsRange[0], portsRange[1]);
-		const portInUse = await docker
-			.listContainers({ all: true })
-			.then((containers) => containers.flatMap((container) => container.Ports?.map((port) => port.PublicPort)));
-		while (portInUse.includes(vncPort)) {
-			vncPort = getRandomNumber(portsRange[0], portsRange[1]);
-		}
-		while (portInUse.includes(agentPort)) {
-			agentPort = getRandomNumber(portsRange[0], portsRange[1]);
-		}
 		const id = `stardust-${crypto.randomUUID()}-${Image.split("/")[2]}`;
 		await docker.createNetwork({
 			Name: id,
@@ -41,10 +27,6 @@ async function createSession(Image: string) {
 			HostConfig: {
 				NetworkMode: id,
 				ShmSize: 1024,
-				PortBindings: {
-					"5901/tcp": [{ hostIp: "127.0.0.1" }, { HostPort: vncPort.toString() }],
-					"6080/tcp": [{ hostIp: "127.0.0.1" }, { HostPort: agentPort.toString() }],
-				},
 			},
 		});
 		await container.start().catch((e) => {
@@ -58,8 +40,6 @@ async function createSession(Image: string) {
 				return await tx
 					.insert(session)
 					.values({
-						vncPort,
-						agentPort,
 						id: container.id,
 						dockerImage: Image,
 						createdAt: Date.now(),

src/lib/session/session-running.ts

@@ -6,17 +6,17 @@ import { consola } from "consola";
  * While loop for checking if the container is running.
  * @param port The port of the container to check
  */
-async function sessionRunning(port: number) {
-	let containerRunning = await fetch(`http://${process.env.CONTAINER_HOST}:${port}/healthcheck`)
+async function sessionRunning(ip: string) {
+	let containerRunning = await fetch(`http://${ip}:6080/healthcheck`)
 		.then((res) => res.json())
 		.then((data) => data.message)
-		.catch(() => consola.warn(`✨ Stardust: Container on port ${port} not running, retrying...`));
+		.catch(() => consola.warn(`✨ Stardust: Container on ${ip} not running, retrying...`));
 	while (!containerRunning) {
 		await new Promise((resolve) => setTimeout(resolve, 2000));
-		containerRunning = await fetch(`http://${process.env.CONTAINER_HOST}:${port}/healthcheck`)
+		containerRunning = await fetch(`http://${ip}:6080/healthcheck`)
 			.then((res) => res.json())
 			.then((data) => data.message)
-			.catch(() => consola.warn(`✨ Stardust: Container on port ${port} not running, retrying...`));
+			.catch(() => consola.warn(`✨ Stardust: Container on ${ip} not running, retrying...`));
 	}
 }
 export { sessionRunning };