From 38d1ea6a3e7bcbf1a639fcb8252a5a02b941f36e Mon Sep 17 00:00:00 2001
From: tgt <tgt@incognitotgt.me>
Date: Mon, 2 Dec 2024 14:08:18 -0500
Subject: [PATCH] feat: usage limits

---
 .config/config-schema.json | 13 +++++++++++++
 src/lib/session/index.ts   | 19 +++++++++++++++++++
 src/types/config.d.ts      |  7 +++++++
 3 files changed, 39 insertions(+)

diff --git a/.config/config-schema.json b/.config/config-schema.json
index 1cbd2e5..4706f2a 100644
--- a/.config/config-schema.json
+++ b/.config/config-schema.json
@@ -130,6 +130,19 @@
 					"default": 1440,
 					"description": "The amount of time to keep an inactive session alive for, in minutes.",
 					"type": "number"
+				},
+				"usageLimits": {
+					"additionalProperties": false,
+					"description": "Session per user usage limit configuration",
+					"properties": {
+						"instance": {
+							"type": "number"
+						},
+						"user": {
+							"type": "number"
+						}
+					},
+					"type": "object"
 				}
 			},
 			"type": "object"
diff --git a/src/lib/session/index.ts b/src/lib/session/index.ts
index d7c5ebb..f9e77ae 100644
--- a/src/lib/session/index.ts
+++ b/src/lib/session/index.ts
@@ -20,6 +20,25 @@ async function createSession(image: string) {
 		const config = getConfig();
 		const userSession = await auth();
 		if (!userSession?.user) throw new Error("User not found");
+		if (config.session?.usageLimits) {
+			const { userIsAdmin, allSessions } = await db.transaction(async (tx) => ({
+				userIsAdmin: (
+					await tx.query.user.findFirst({
+						where: (user, { eq }) => eq(user.id, userSession?.user?.id as string),
+						columns: {
+							isAdmin: true,
+						},
+					})
+				)?.isAdmin,
+				allSessions: await tx.select().from(session),
+			}));
+			if (config.session?.usageLimits.instance && config.session?.usageLimits.instance <= allSessions.length)
+				throw new Error("Instance limit exceeded");
+			if (config.session.usageLimits.user && !userIsAdmin) {
+				const userSessions = allSessions.filter((v) => v.userId === userSession.user.id);
+				if (config.session?.usageLimits.user <= userSessions.length) throw new Error("User session limit exceeded");
+			}
+		}
 		const container = await docker.createContainer({
 			name: `stardust-${createId()}-${image.split("/")[2] || image.split("/")[1]}`,
 			Image: image,
diff --git a/src/types/config.d.ts b/src/types/config.d.ts
index 4fa2749..c767936 100644
--- a/src/types/config.d.ts
+++ b/src/types/config.d.ts
@@ -96,6 +96,13 @@ export interface SessionConfig {
 	 * @default system default
 	 */
 	dnsServers?: string[];
+	/**
+	 * Session per user usage limit configuration
+	 */
+	usageLimits?: {
+		instance?: number;
+		user?: number;
+	};
 }
 
 export interface TurnstileConfig {