From 1c41d55ceb9fb9b4ccbb3d715c94fac0964a0db5 Mon Sep 17 00:00:00 2001
From: Arthur Lu <alu@tronnet.net>
Date: Mon, 8 Apr 2024 21:19:57 +0000
Subject: [PATCH] add pool selection to instance create, allow multiple pools
 per user

---
 src/routes/cluster.js |  3 ++-
 src/routes/sync.js    | 37 ++++++++++++++++++++++---------------
 template.localdb.json |  5 ++++-
 3 files changed, 28 insertions(+), 17 deletions(-)

diff --git a/src/routes/cluster.js b/src/routes/cluster.js
index 0570b7b..69cd384 100644
--- a/src/routes/cluster.js
+++ b/src/routes/cluster.js
@@ -145,6 +145,7 @@ router.post(`${basePath}/create`, async (req, res) => {
 		vmid: req.params.vmid,
 		hostname: req.body.hostname,
 		name: req.body.name,
+		pool: req.body.pool,
 		cores: req.body.cores,
 		memory: req.body.memory,
 		swap: req.body.swap,
@@ -206,7 +207,7 @@ router.post(`${basePath}/create`, async (req, res) => {
 		vmid: params.vmid,
 		cores: Number(params.cores),
 		memory: Number(params.memory),
-		pool: user.cluster.pool
+		pool: params.pool // TODO allow user to select pool to assign VM
 	};
 	for (const key of Object.keys(user.templates.instances[params.type])) {
 		action[key] = user.templates.instances[params.type][key].value;
diff --git a/src/routes/sync.js b/src/routes/sync.js
index d1f95fb..e1c61ac 100644
--- a/src/routes/sync.js
+++ b/src/routes/sync.js
@@ -63,7 +63,7 @@ if (schemes.hash.enabled) {
 // setup interupt scheme
 if (schemes.interrupt.enabled) {
 	const wsServer = new WebSocketServer({ noServer: true, path: "/api/sync/interrupt" });
-	wsServer.on("connection", (socket, username, pool) => {
+	wsServer.on("connection", (socket, username, pools) => {
 		// add new socket to userSocketmap
 		if (userSocketMap[username]) {
 			const index = Object.keys(userSocketMap[username]).length;
@@ -74,13 +74,15 @@ if (schemes.interrupt.enabled) {
 			userSocketMap[username] = { 0: socket };
 			socket.userIndex = 0;
 		}
-		// add user to associated pool in poolUserMap
-		if (poolUserMap[pool]) {
-			poolUserMap[pool][username] = true;
-		}
-		else {
-			poolUserMap[pool] = {};
-			poolUserMap[pool][username] = true;
+		// for each pool, add user to associated pool in poolUserMap
+		for (const pool of pools) {
+			if (poolUserMap[pool]) {
+				poolUserMap[pool][username] = true;
+			}
+			else {
+				poolUserMap[pool] = {};
+				poolUserMap[pool][username] = true;
+			}
 		}
 		// add socket entry into requestedRates
 		const index = Object.keys(requestedRates).length;
@@ -96,11 +98,13 @@ if (schemes.interrupt.enabled) {
 			if (Object.keys(userSocketMap[username]).length === 0) {
 				// delete the user entry
 				delete userSocketMap[username];
-				// remove user from poolUserMap pool entry
-				delete poolUserMap[pool][username];
-				// if the poolUserMap pool entry is empty, delete the entry
-				if (Object.keys(poolUserMap[pool]).length === 0) {
-					delete poolUserMap[pool];
+				// remove user from poolUserMap entry for each pool
+				for (const pool of pools) {
+					delete poolUserMap[pool][username];
+					// if the poolUserMap pool entry is empty, delete the entry
+					if (Object.keys(poolUserMap[pool]).length === 0) {
+						delete poolUserMap[pool];
+					}
 				}
 			}
 			// remove socket entry from requestedRates
@@ -154,6 +158,7 @@ if (schemes.interrupt.enabled) {
 	});
 	// handle the wss upgrade request
 	global.server.on("upgrade", async (req, socket, head) => {
+		// basic auth validation
 		const cookies = cookie.parse(req.headers.cookie || "");
 		const auth = (await global.pve.requestPVE("/version", "GET", { cookies })).status === 200;
 		if (!auth) {
@@ -161,8 +166,10 @@ if (schemes.interrupt.enabled) {
 		}
 		else {
 			wsServer.handleUpgrade(req, socket, head, (socket) => {
-				const pool = global.db.getUser(cookies.username).cluster.pool;
-				wsServer.emit("connection", socket, cookies.username, pool);
+				// get the user pools
+				const pools = global.db.getUser(cookies.username).cluster.pools;
+				// emit the connection to initialize socket
+				wsServer.emit("connection", socket, cookies.username, pools);
 			});
 		}
 	});
diff --git a/template.localdb.json b/template.localdb.json
index 5bf5422..d20c6e3 100644
--- a/template.localdb.json
+++ b/template.localdb.json
@@ -81,7 +81,10 @@
 					"min": 100,
 					"max": 199
 				},
-				"pool": "examplepool"
+				"pools": [
+					"examplepool1",
+					"examplepool2"
+				]
 			},
 			"templates": {
 				"instances": {