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
+79
-9
No files matched your search
+15
-1
@@ -175,8 +175,22 @@ export function createWebServer(options: WebServerOptions): WebServer {
|
||||
socket.destroy();
|
||||
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) => {
|
||||
(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);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -7,6 +7,7 @@ import { createDatabase, type BotDatabase } from "../data/database.js";
|
||||
import { createUserStore } from "../data/users.js";
|
||||
import { createSessionStore } from "../data/sessions.js";
|
||||
import { validateSessionFromHeaders, SESSION_COOKIE_NAME } from "./auth/validateSession.js";
|
||||
import { setupWebSocket } from "./websocket.js";
|
||||
|
||||
function buildServer(sessions: ReturnType<typeof createSessionStore>) {
|
||||
const app = express();
|
||||
@@ -72,3 +73,44 @@ describe("WebSocket auth at upgrade", () => {
|
||||
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();
|
||||
});
|
||||
});
|
||||
+22
-8
@@ -10,6 +10,17 @@ export function setupWebSocket(
|
||||
): () => void {
|
||||
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) */
|
||||
const attachedBots = new Map<string, {
|
||||
bot: BotInstance;
|
||||
@@ -22,7 +33,10 @@ export function setupWebSocket(
|
||||
clients.add(ws);
|
||||
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.on("close", () => {
|
||||
@@ -36,17 +50,17 @@ export function setupWebSocket(
|
||||
});
|
||||
});
|
||||
|
||||
const broadcast = (data: object) => {
|
||||
const broadcast = (data: object, botId?: string) => {
|
||||
const message = JSON.stringify(data);
|
||||
for (const client of clients) {
|
||||
if (client.readyState === WebSocket.OPEN) {
|
||||
if (client.readyState !== WebSocket.OPEN) continue;
|
||||
if (botId !== undefined && !visibleToClient(client, botId)) continue;
|
||||
try {
|
||||
client.send(message);
|
||||
} catch {
|
||||
clients.delete(client);
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
function detachBotListener(id: string): void {
|
||||
@@ -72,7 +86,7 @@ export function setupWebSocket(
|
||||
botId: bot.id,
|
||||
status: bot.getStatus(),
|
||||
queue: bot.getQueue(),
|
||||
});
|
||||
}, bot.id);
|
||||
};
|
||||
|
||||
const onConnected = () => {
|
||||
@@ -80,7 +94,7 @@ export function setupWebSocket(
|
||||
type: "botConnected",
|
||||
botId: bot.id,
|
||||
status: bot.getStatus(),
|
||||
});
|
||||
}, bot.id);
|
||||
};
|
||||
|
||||
const onDisconnected = () => {
|
||||
@@ -88,7 +102,7 @@ export function setupWebSocket(
|
||||
type: "botDisconnected",
|
||||
botId: bot.id,
|
||||
status: bot.getStatus(),
|
||||
});
|
||||
}, bot.id);
|
||||
};
|
||||
|
||||
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
|
||||
const onBotInstanceRemoved = (id: string) => {
|
||||
detachBotListener(id);
|
||||
broadcast({ type: "botRemoved", botId: id });
|
||||
broadcast({ type: "botRemoved", botId: id }, id);
|
||||
};
|
||||
botManager.on("botInstanceRemoved", onBotInstanceRemoved);
|
||||
|
||||
|
||||
Reference in new issue
Block a user