mirror of
https://github.com/ZHANGTIANYAO1/teamspeak-music-bot.git
synced 2026-10-02 04:52:50 +08:00
feat(ws): scope guest WebSocket feed to allowed bots
This commit is contained in:
1 parent
0073d7d612
commit
d2ab888114
3 files changed
+83
-13
No files matched your search
+15
-1
@@ -175,8 +175,22 @@ export function createWebServer(options: WebServerOptions): WebServer {
|
|||||||
socket.destroy();
|
socket.destroy();
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
// Guest sessions are only valid while guest mode is enabled.
|
||||||
|
if (result.role === "guest" && !options.config.guestMode.enabled) {
|
||||||
|
socket.write("HTTP/1.1 401 Unauthorized\r\nConnection: close\r\n\r\n");
|
||||||
|
socket.destroy();
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
const guestBots = options.config.guestMode.bots;
|
||||||
|
const botScope: "all" | Set<string> =
|
||||||
|
result.role === "guest"
|
||||||
|
? guestBots === "all" ? "all" : new Set(guestBots)
|
||||||
|
: "all";
|
||||||
wss.handleUpgrade(req, socket, head, (ws) => {
|
wss.handleUpgrade(req, socket, head, (ws) => {
|
||||||
(ws as unknown as { userId: string }).userId = result.userId;
|
const w = ws as unknown as { userId: string; isGuest: boolean; botScope: "all" | Set<string> };
|
||||||
|
w.userId = result.userId;
|
||||||
|
w.isGuest = result.role === "guest";
|
||||||
|
w.botScope = botScope;
|
||||||
wss.emit("connection", ws, req);
|
wss.emit("connection", ws, req);
|
||||||
});
|
});
|
||||||
});
|
});
|
||||||
|
|||||||
@@ -7,6 +7,7 @@ import { createDatabase, type BotDatabase } from "../data/database.js";
|
|||||||
import { createUserStore } from "../data/users.js";
|
import { createUserStore } from "../data/users.js";
|
||||||
import { createSessionStore } from "../data/sessions.js";
|
import { createSessionStore } from "../data/sessions.js";
|
||||||
import { validateSessionFromHeaders, SESSION_COOKIE_NAME } from "./auth/validateSession.js";
|
import { validateSessionFromHeaders, SESSION_COOKIE_NAME } from "./auth/validateSession.js";
|
||||||
|
import { setupWebSocket } from "./websocket.js";
|
||||||
|
|
||||||
function buildServer(sessions: ReturnType<typeof createSessionStore>) {
|
function buildServer(sessions: ReturnType<typeof createSessionStore>) {
|
||||||
const app = express();
|
const app = express();
|
||||||
@@ -72,3 +73,44 @@ describe("WebSocket auth at upgrade", () => {
|
|||||||
ws.close();
|
ws.close();
|
||||||
});
|
});
|
||||||
});
|
});
|
||||||
|
|
||||||
|
describe("WebSocket guest bot scope", () => {
|
||||||
|
it("guest init is filtered to the guest bot scope", () => {
|
||||||
|
const sent: any[] = [];
|
||||||
|
const fakeWs: any = {
|
||||||
|
readyState: 1,
|
||||||
|
isGuest: true,
|
||||||
|
botScope: new Set(["bot1"]),
|
||||||
|
send: (m: string) => sent.push(JSON.parse(m)),
|
||||||
|
on: () => {},
|
||||||
|
};
|
||||||
|
const fakeWss: any = {
|
||||||
|
on: (ev: string, cb: any) => {
|
||||||
|
if (ev === "connection") fakeWss._conn = cb;
|
||||||
|
},
|
||||||
|
};
|
||||||
|
const makeBot = (id: string) => ({
|
||||||
|
id,
|
||||||
|
getStatus: () => ({ id }),
|
||||||
|
getQueue: () => [],
|
||||||
|
on: () => {},
|
||||||
|
removeListener: () => {},
|
||||||
|
});
|
||||||
|
const botManager: any = {
|
||||||
|
getAllBots: () => [makeBot("bot1"), makeBot("bot2")],
|
||||||
|
on: () => {},
|
||||||
|
off: () => {},
|
||||||
|
removeListener: () => {},
|
||||||
|
};
|
||||||
|
const cleanup = setupWebSocket(fakeWss, botManager, {
|
||||||
|
debug() {},
|
||||||
|
error() {},
|
||||||
|
info() {},
|
||||||
|
warn() {},
|
||||||
|
} as any);
|
||||||
|
fakeWss._conn(fakeWs);
|
||||||
|
const init = sent.find((m) => m.type === "init");
|
||||||
|
expect(init.bots.map((b: any) => b.id)).toEqual(["bot1"]);
|
||||||
|
cleanup();
|
||||||
|
});
|
||||||
|
});
|
||||||
+26
-12
@@ -10,6 +10,17 @@ export function setupWebSocket(
|
|||||||
): () => void {
|
): () => void {
|
||||||
const clients = new Set<WebSocket>();
|
const clients = new Set<WebSocket>();
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Whether a given bot is visible to a WebSocket client. Member/admin clients
|
||||||
|
* (non-guest) and guests with full scope see everything; scoped guests only
|
||||||
|
* see bots in their allowed set.
|
||||||
|
*/
|
||||||
|
function visibleToClient(ws: WebSocket, botId: string): boolean {
|
||||||
|
const w = ws as unknown as { isGuest?: boolean; botScope?: "all" | Set<string> };
|
||||||
|
if (!w.isGuest || w.botScope === "all" || !w.botScope) return true;
|
||||||
|
return w.botScope.has(botId);
|
||||||
|
}
|
||||||
|
|
||||||
/** Track which bot instances have listeners attached (keyed by id, storing ref) */
|
/** Track which bot instances have listeners attached (keyed by id, storing ref) */
|
||||||
const attachedBots = new Map<string, {
|
const attachedBots = new Map<string, {
|
||||||
bot: BotInstance;
|
bot: BotInstance;
|
||||||
@@ -22,7 +33,10 @@ export function setupWebSocket(
|
|||||||
clients.add(ws);
|
clients.add(ws);
|
||||||
logger.debug("WebSocket client connected");
|
logger.debug("WebSocket client connected");
|
||||||
|
|
||||||
const bots = botManager.getAllBots().map((b) => b.getStatus());
|
const bots = botManager
|
||||||
|
.getAllBots()
|
||||||
|
.filter((b) => visibleToClient(ws, b.id))
|
||||||
|
.map((b) => b.getStatus());
|
||||||
ws.send(JSON.stringify({ type: "init", bots }));
|
ws.send(JSON.stringify({ type: "init", bots }));
|
||||||
|
|
||||||
ws.on("close", () => {
|
ws.on("close", () => {
|
||||||
@@ -36,15 +50,15 @@ export function setupWebSocket(
|
|||||||
});
|
});
|
||||||
});
|
});
|
||||||
|
|
||||||
const broadcast = (data: object) => {
|
const broadcast = (data: object, botId?: string) => {
|
||||||
const message = JSON.stringify(data);
|
const message = JSON.stringify(data);
|
||||||
for (const client of clients) {
|
for (const client of clients) {
|
||||||
if (client.readyState === WebSocket.OPEN) {
|
if (client.readyState !== WebSocket.OPEN) continue;
|
||||||
try {
|
if (botId !== undefined && !visibleToClient(client, botId)) continue;
|
||||||
client.send(message);
|
try {
|
||||||
} catch {
|
client.send(message);
|
||||||
clients.delete(client);
|
} catch {
|
||||||
}
|
clients.delete(client);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@@ -72,7 +86,7 @@ export function setupWebSocket(
|
|||||||
botId: bot.id,
|
botId: bot.id,
|
||||||
status: bot.getStatus(),
|
status: bot.getStatus(),
|
||||||
queue: bot.getQueue(),
|
queue: bot.getQueue(),
|
||||||
});
|
}, bot.id);
|
||||||
};
|
};
|
||||||
|
|
||||||
const onConnected = () => {
|
const onConnected = () => {
|
||||||
@@ -80,7 +94,7 @@ export function setupWebSocket(
|
|||||||
type: "botConnected",
|
type: "botConnected",
|
||||||
botId: bot.id,
|
botId: bot.id,
|
||||||
status: bot.getStatus(),
|
status: bot.getStatus(),
|
||||||
});
|
}, bot.id);
|
||||||
};
|
};
|
||||||
|
|
||||||
const onDisconnected = () => {
|
const onDisconnected = () => {
|
||||||
@@ -88,7 +102,7 @@ export function setupWebSocket(
|
|||||||
type: "botDisconnected",
|
type: "botDisconnected",
|
||||||
botId: bot.id,
|
botId: bot.id,
|
||||||
status: bot.getStatus(),
|
status: bot.getStatus(),
|
||||||
});
|
}, bot.id);
|
||||||
};
|
};
|
||||||
|
|
||||||
bot.on("stateChange", onStateChange);
|
bot.on("stateChange", onStateChange);
|
||||||
@@ -117,7 +131,7 @@ export function setupWebSocket(
|
|||||||
// React when a bot is removed: detach its listener and tell clients to drop it
|
// React when a bot is removed: detach its listener and tell clients to drop it
|
||||||
const onBotInstanceRemoved = (id: string) => {
|
const onBotInstanceRemoved = (id: string) => {
|
||||||
detachBotListener(id);
|
detachBotListener(id);
|
||||||
broadcast({ type: "botRemoved", botId: id });
|
broadcast({ type: "botRemoved", botId: id }, id);
|
||||||
};
|
};
|
||||||
botManager.on("botInstanceRemoved", onBotInstanceRemoved);
|
botManager.on("botInstanceRemoved", onBotInstanceRemoved);
|
||||||
|
|
||||||
|
|||||||
Reference in new issue
Block a user