diff --git a/CLAUDE.md b/CLAUDE.md index af16d8d..f0d658c 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -92,7 +92,7 @@ No routing configuration needed - pathname automatically maps to function name. #### Authentication - `CheckToken` accepts either a backend-issued `{Token}` (current clients) or a legacy `{SessionID, Username}` (old clients) -- `Login` (in `Source/Process.ts`) swaps a verified PHPSESSID for a token issued by `Source/SessionToken.ts`; tokens are stored as SHA3 hashes in `session_token`, expire after 30 days unused, and are revoked by `Logout`. The username always comes from the token row +- `Login` (in `Source/Process.ts`) swaps a verified PHPSESSID for a token issued by `Source/SessionToken.ts`; tokens are stored as SHA3 hashes in `session_token`, expire after 30 days unused, and are revoked by `Logout` (this device) or `LogoutAll` (every token for the user, plus their cached `phpsessid` rows; open notification sockets are closed best-effort, and the Durable Object re-checks the token when admitting a socket so a handshake racing `LogoutAll` can't stay subscribed). The username always comes from the token row - Legacy PHPSESSIDs are validated against the XMOJ website and cached as SHA3 hashes in the `phpsessid` table - Token requests carry no PHPSESSID. Code that must scrape xmoj as the user calls `this.RequireSession()` first, which answers `SessionRequired`; the client then retries once with `{Token, SessionID}`. Keep such scrapes rare, since we want to see users' PHPSESSIDs as little as possible - `problem_score` caches users' best scores so `GetStd` doesn't need to scrape (and so doesn't need a session) once a score has cleared the threshold diff --git a/Source/NotificationManager.ts b/Source/NotificationManager.ts index a911b55..260d517 100644 --- a/Source/NotificationManager.ts +++ b/Source/NotificationManager.ts @@ -34,18 +34,21 @@ interface HibernationWebSocket extends WebSocket { interface NotificationEnvironment { NOTIFICATION_PUSH_TOKEN?: string; + DB?: D1Database; } export class NotificationManager { private readonly state: DurableObjectState; private readonly sessions: Map>; private readonly pushToken: string; + private readonly db: D1Database | null; private static readonly MAX_SESSIONS_PER_USER = 20; constructor(state: DurableObjectState, env: NotificationEnvironment) { this.state = state; this.sessions = new Map>(); this.pushToken = env.NOTIFICATION_PUSH_TOKEN || ""; + this.db = env.DB || null; // `state.getWebSockets()` is synchronous in the current Cloudflare runtime. this.rebuildSessionIndex(); } @@ -123,26 +126,70 @@ export class NotificationManager { async fetch(request: Request): Promise { const url = new URL(request.url); - - // Internal push channel from Process.ts. if (url.pathname === "/notify") { - if (this.pushToken === "" || request.headers.get("X-Notification-Token") !== this.pushToken) { - return new Response("Unauthorized", {status: 401}); + return this.handleNotify(request); + } + if (url.pathname === "/disconnect") { + return this.handleDisconnect(request); + } + return this.handleConnect(request, url); + } + + // The internal channels from Process.ts are only for the worker. + private isInternal(request: Request): boolean { + return this.pushToken !== "" && request.headers.get("X-Notification-Token") === this.pushToken; + } + + // Internal push channel from Process.ts. + private async handleNotify(request: Request): Promise { + if (!this.isInternal(request)) { + return new Response("Unauthorized", {status: 401}); + } + + const body = await request.json() as { userId: string; notification: object }; + const userSessions = this.sessions.get(body.userId); + if (userSessions) { + const payload = JSON.stringify(body.notification); + for (const websocket of userSessions) { + if (websocket.readyState === 1) { + websocket.send(payload); + } } + } + return new Response("OK"); + } - const body = await request.json() as { userId: string; notification: object }; - const userSessions = this.sessions.get(body.userId); - if (userSessions) { - const payload = JSON.stringify(body.notification); - for (const websocket of userSessions) { - if (websocket.readyState === 1) { - websocket.send(payload); - } + // Internal: the user logged out everywhere. 4001 tells the client its + // token is gone, so it reconnects by proving its xmoj session again. + private async handleDisconnect(request: Request): Promise { + if (!this.isInternal(request)) { + return new Response("Unauthorized", {status: 401}); + } + + let body: { userId?: unknown }; + try { + body = await request.json() as { userId?: unknown }; + } catch (_) { + return new Response("Bad Request", {status: 400}); + } + if (!body || typeof body.userId !== "string") { + return new Response("Bad Request", {status: 400}); + } + const userSessions = this.sessions.get(body.userId); + if (userSessions) { + for (const websocket of Array.from(userSessions)) { + this.removeSession(body.userId, websocket); + try { + websocket.close(4001, "Logged out"); + } catch (_) { + // Already closing. } } - return new Response("OK"); } + return new Response("OK"); + } + private async handleConnect(request: Request, url: URL): Promise { const upgradeHeader = request.headers.get("Upgrade"); if (upgradeHeader !== "websocket") { return new Response("Expected WebSocket", {status: 426}); @@ -164,6 +211,17 @@ export class NotificationManager { }); this.addSession(userId, server); + // The worker checked the token, but LogoutAll may have revoked it since. + // The socket is registered before this check because D1 calls don't hold + // other requests back: /disconnect either finds the socket and closes it, + // or ran first, and then the token is already gone here. + const tokenHash = url.searchParams.get("tokenHash"); + if (tokenHash && !(await this.tokenStillValid(tokenHash))) { + this.removeSession(userId, server); + server.close(4001, "Logged out"); + return new Response(null, {status: 101, webSocket: client}); + } + // Only set when the worker minted a token for a client that connected // with its PHPSESSID. It is handed over once and never stored here. const issuedToken = url.searchParams.get("issuedToken"); @@ -187,11 +245,36 @@ export class NotificationManager { } } - webSocketClose(websocket: WebSocket): void { + private async tokenStillValid(tokenHash: string): Promise { + if (this.db === null) { + return true; + } + try { + const row = await this.db.prepare("SELECT 1 FROM session_token WHERE token_hash = ?").bind(tokenHash).first(); + return row !== null; + } catch (_) { + // The worker has just validated this token; a D1 hiccup shouldn't cost + // the user their notifications. + return true; + } + } + + webSocketClose(websocket: WebSocket, code?: number, reason?: string): void { const userId = this.getSocketUserId(websocket); if (userId !== "") { this.removeSession(userId, websocket); } + // Our compatibility date predates automatic close replies, so finish the + // close handshake ourselves; otherwise the client sits in CLOSING and + // never sees why it was closed (such as 4001 from /disconnect). + try { + // close() only takes 1000 and 3000-4999; peers also send 1001 (page + // closed) and others, which must still get an answer. + const ReplyCode = code !== undefined && (code === 1000 || (code >= 3000 && code <= 4999)) ? code : 1000; + websocket.close(ReplyCode, reason || ""); + } catch (_) { + // Already closed. + } } webSocketError(websocket: WebSocket): void { diff --git a/Source/Process.ts b/Source/Process.ts index e8c7ea1..73d589f 100644 --- a/Source/Process.ts +++ b/Source/Process.ts @@ -25,7 +25,7 @@ import * as sqlstring from 'sqlstring'; // @ts-ignore import CryptoJS from "crypto-js"; import {AnalyticsEngineDataset, D1Database, D1DatabaseSession, DurableObjectNamespace, KVNamespace} from "@cloudflare/workers-types"; -import {IssueSessionToken, ResolveSessionToken, RevokeSessionToken} from "./SessionToken"; +import {IssueSessionToken, ResolveSessionToken, RevokeAllSessionTokens, RevokeSessionToken} from "./SessionToken"; interface Environment { API_TOKEN: string; @@ -663,6 +663,25 @@ export class Process { } }; + // Closes the user's open notification sockets. They reconnect on their own, + // and only get back in if they can still prove an xmoj session. + private disconnectNotifications = async (userId: string): Promise => { + try { + const id = this.notifications.idFromName(userId); + const stub = this.notifications.get(id); + await stub.fetch(new Request("https://dummy/disconnect", { + method: "POST", + headers: { + "X-Notification-Token": this.notificationPushToken + }, + body: JSON.stringify({userId}) + })); + } catch (_) { + // The tokens are already gone; a socket left open only keeps receiving + // notifications until it next reconnects. + } + }; + private AddBBSMention = async (ToUserID: string, PostID: number, ReplyID: number): Promise => { if (ToUserID === this.Username) { return; @@ -1592,6 +1611,17 @@ export class Process { } return new Result(true, "登出成功"); }, + // Signs the user out of the backend on every device: all tokens, the + // cached PHPSESSID checks old clients rely on, and open sockets. + LogoutAll: async (Data: object): Promise => { + ThrowErrorIfFailed(this.CheckParams(Data, {})); + await RevokeAllSessionTokens(this.XMOJDatabase, this.Username); + ThrowErrorIfFailed(await this.XMOJDatabase.Delete("phpsessid", { + user_id: this.Username + })); + await this.disconnectNotifications(this.Username); + return new Result(true, "已在所有设备上登出"); + }, GetStdList: async (Data: object): Promise => { ThrowErrorIfFailed(this.CheckParams(Data, {})); const ResponseData = { diff --git a/Source/SessionToken.ts b/Source/SessionToken.ts index 802d4cc..6675c40 100644 --- a/Source/SessionToken.ts +++ b/Source/SessionToken.ts @@ -79,6 +79,14 @@ export const ResolveSessionToken = async (XMOJDatabase: Database, Token: string) return Rows[0]["user_id"]; }; +// Every token the user holds, on every device. Devices that still have a live +// xmoj session simply exchange it again; a token on its own is dead. +export const RevokeAllSessionTokens = async (XMOJDatabase: Database, Username: string): Promise => { + ThrowErrorIfFailed(await XMOJDatabase.Delete("session_token", { + user_id: Username + })); +}; + export const RevokeSessionToken = async (XMOJDatabase: Database, Token: string): Promise => { ThrowErrorIfFailed(await XMOJDatabase.Delete("session_token", { token_hash: HashSessionToken(Token) diff --git a/Source/index.ts b/Source/index.ts index 726b2ad..e30d3d7 100644 --- a/Source/index.ts +++ b/Source/index.ts @@ -18,7 +18,7 @@ import {Process, RebuildStdList} from "./Process"; import {Database} from "./Database"; import {NotificationManager} from "./NotificationManager"; -import {IssueSessionToken, ResolveSessionToken, SessionTokenLifetime} from "./SessionToken"; +import {HashSessionToken, IssueSessionToken, ResolveSessionToken, SessionTokenLifetime} from "./SessionToken"; import type {D1Database, KVNamespace, AnalyticsEngineDataset, DurableObjectNamespace, Ai} from "@cloudflare/workers-types"; interface Environment { @@ -170,6 +170,13 @@ const HandleRequest = async (RequestData: Request, Environment: Environment): Pr if (issuedToken !== "") { forwardURL.searchParams.set("issuedToken", issuedToken); } + // Lets the Durable Object re-check the token as it admits the socket, + // so a handshake racing LogoutAll can't slip in. Only the hash travels. + forwardURL.searchParams.delete("tokenHash"); + const admittedToken = token !== "" ? token : issuedToken; + if (admittedToken !== "") { + forwardURL.searchParams.set("tokenHash", HashSessionToken(admittedToken)); + } return await notificationStub.fetch(new Request(forwardURL.toString(), RequestData)); } diff --git a/test/index.test.js b/test/index.test.js index 44d656a..6c8c325 100644 --- a/test/index.test.js +++ b/test/index.test.js @@ -125,3 +125,25 @@ test('WebSocket drops an issuedToken the client supplied', async (t) => { await Worker.fetch(new Request('https://api.xmoj-script.uk/ws/notifications?SessionID=abc123&issuedToken=' + 'd'.repeat(64)), env); assert.strictEqual(forwarded[0].searchParams.get('issuedToken'), null); }); + +test('WebSocket forwards only the token hash, so the Durable Object can re-check it', async () => { + const token = 'b'.repeat(64); + const { env, forwarded } = wsEnvironment({ [token]: { user_id: 'alice', last_used: Date.now() } }); + await Worker.fetch(new Request('https://api.xmoj-script.uk/ws/notifications?Token=' + token + '&tokenHash=spoofed'), env); + assert.strictEqual(forwarded[0].searchParams.get('tokenHash'), HashSessionToken(token)); + assert.ok(!forwarded[0].toString().includes(token), 'the token itself is not forwarded'); +}); + +test('WebSocket forwards the hash of a freshly minted token too', async (t) => { + t.mock.method(global, 'fetch', async () => ProfilePage('alice')); + const { env, forwarded } = wsEnvironment(); + await Worker.fetch(new Request('https://api.xmoj-script.uk/ws/notifications?SessionID=abc123&IssueToken=1', { headers: { Upgrade: 'websocket' } }), env); + assert.strictEqual(forwarded[0].searchParams.get('tokenHash'), HashSessionToken(forwarded[0].searchParams.get('issuedToken'))); +}); + +test('WebSocket on the plain PHPSESSID path forwards no token hash', async (t) => { + t.mock.method(global, 'fetch', async () => ProfilePage('alice')); + const { env, forwarded } = wsEnvironment(); + await Worker.fetch(new Request('https://api.xmoj-script.uk/ws/notifications?SessionID=abc123&tokenHash=spoofed'), env); + assert.strictEqual(forwarded[0].searchParams.get('tokenHash'), null); +}); diff --git a/test/notification_manager.test.js b/test/notification_manager.test.js index ed8fef6..a4ba7c8 100644 --- a/test/notification_manager.test.js +++ b/test/notification_manager.test.js @@ -86,3 +86,119 @@ test('rejects notify without internal token', async () => { assert.strictEqual(response.status, 401); assert.deepStrictEqual(socket.getSent(), []); }); + +test('disconnect closes every socket of the user with 4001 and nobody else\'s', async () => { + const manager = createManager(); + const closed = []; + const socket = (user) => Object.assign(createFakeWebSocket(user), { close: (code, reason) => closed.push([user, code, reason]) }); + const aliceA = socket('alice'), aliceB = socket('alice'), bob = socket('bob'); + manager.addSession('alice', aliceA); + manager.addSession('alice', aliceB); + manager.addSession('bob', bob); + + const response = await manager.fetch(new Request('https://dummy/disconnect', { + method: 'POST', + headers: { 'X-Notification-Token': 'test-push-token' }, + body: JSON.stringify({ userId: 'alice' }), + })); + + assert.strictEqual(response.status, 200); + assert.deepStrictEqual(closed, [['alice', 4001, 'Logged out'], ['alice', 4001, 'Logged out']]); + await manager.fetch(new Request('https://dummy/notify', { + method: 'POST', + headers: { 'X-Notification-Token': 'test-push-token' }, + body: JSON.stringify({ userId: 'bob', notification: { type: 'pong' } }), + })); + assert.deepStrictEqual(bob.getSent(), [JSON.stringify({ type: 'pong' })], 'other users keep their socket'); +}); + +test('disconnect requires the internal token', async () => { + const manager = createManager(); + const response = await manager.fetch(new Request('https://dummy/disconnect', { + method: 'POST', + body: JSON.stringify({ userId: 'alice' }), + })); + assert.strictEqual(response.status, 401); +}); + +test('webSocketClose answers the close frame so the client sees the code', () => { + const manager = createManager(); + const closed = []; + const socket = Object.assign(createFakeWebSocket('alice'), { close: (code, reason) => closed.push([code, reason]) }); + manager.addSession('alice', socket); + manager.webSocketClose(socket, 4001, 'Logged out'); + manager.webSocketClose(Object.assign(createFakeWebSocket('alice'), { close: (code, reason) => closed.push([code, reason]) }), 1005, ''); + assert.deepStrictEqual(closed, [[4001, 'Logged out'], [1000, '']], 'reserved codes cannot be sent back'); +}); + +// Admits a socket through fetch() with stand-ins for the Workers-only pieces. +async function admit(manager, query) { + const client = createFakeWebSocket(), server = Object.assign(createFakeWebSocket(), { closed: null }); + server.close = (code, reason) => { server.closed = [code, reason]; }; + const realPair = global.WebSocketPair, RealResponse = global.Response; + global.WebSocketPair = function () { return { 0: client, 1: server }; }; + global.Response = function (body, init) { return { status: init.status, webSocket: init.webSocket }; }; + try { + const response = await manager.fetch(new Request('https://dummy/ws?' + query, { headers: { Upgrade: 'websocket' } })); + return { response, server }; + } finally { + global.WebSocketPair = realPair; + global.Response = RealResponse; + } +} + +function managerWithTokens(validHashes, failing = false) { + const state = { getWebSockets: () => [], acceptWebSocket: () => {} }; + const DB = { + prepare: () => ({ + bind: (hash) => ({ + first: async () => { + if (failing) throw new Error('D1 down'); + return validHashes.includes(hash) ? { 1: 1 } : null; + } + }) + }) + }; + return new NotificationManager(state, { NOTIFICATION_PUSH_TOKEN: 'test-push-token', DB }); +} + +test('a socket whose token was revoked mid-handshake is closed with 4001', async () => { + const manager = managerWithTokens([]); + const { server } = await admit(manager, 'userId=alice&tokenHash=gone'); + assert.deepStrictEqual(server.closed, [4001, 'Logged out']); + assert.strictEqual(manager.sessions.get('alice'), undefined, 'not left subscribed'); + assert.deepStrictEqual(server.getSent(), [], 'never told it is connected'); +}); + +test('a socket whose token is still valid is admitted', async () => { + const manager = managerWithTokens(['live']); + const { server } = await admit(manager, 'userId=alice&tokenHash=live'); + assert.strictEqual(server.closed, null); + assert.strictEqual(manager.sessions.get('alice').size, 1); + assert.match(server.getSent()[0], /"type":"connected"/); +}); + +test('a D1 failure during the re-check does not drop the socket', async () => { + const manager = managerWithTokens([], true); + const { server } = await admit(manager, 'userId=alice&tokenHash=whatever'); + assert.strictEqual(server.closed, null); +}); + +test('disconnect rejects a body it cannot read instead of failing silently', async () => { + const manager = createManager(); + for (const body of ['not json', '', JSON.stringify({ userId: 5 })]) { + const response = await manager.fetch(new Request('https://dummy/disconnect', { + method: 'POST', headers: { 'X-Notification-Token': 'test-push-token' }, body, + })); + assert.strictEqual(response.status, 400, JSON.stringify(body)); + } +}); + +test('webSocketClose answers codes close() would reject, like 1001', () => { + const manager = createManager(); + const closed = []; + for (const code of [1001, 1002, 1011, 2999, 5000]) { + manager.webSocketClose(Object.assign(createFakeWebSocket('alice'), { close: (c) => closed.push(c) }), code, ''); + } + assert.deepStrictEqual(closed, [1000, 1000, 1000, 1000, 1000]); +}); diff --git a/test/process.test.js b/test/process.test.js index 3e7d3f4..69d45e1 100644 --- a/test/process.test.js +++ b/test/process.test.js @@ -1541,3 +1541,25 @@ test('A score cached by a racing request wins over our older, lower scrape', asy proc.GetProblemScore = async () => 40; assert.strictEqual(await proc.GetProblemScoreChecker(1000, 50), 100); }); + +test('LogoutAll revokes every token, the cached PHPSESSID checks and open sockets', async () => { + const del = test.mock.fn(async () => new Result(true, '')); + const proc = createProcess({ db: { Delete: del } }); + proc.Username = 'alice'; + const result = await proc.ProcessFunctions['LogoutAll']({}); + assert.ok(result.Success); + assert.deepStrictEqual(del.mock.calls.map((call) => call.arguments), [ + ['session_token', { user_id: 'alice' }], + ['phpsessid', { user_id: 'alice' }] + ]); + const request = proc._notifyFetch.mock.calls[0].arguments[0]; + assert.strictEqual(new URL(request.url).pathname, '/disconnect'); + assert.strictEqual(request.headers.get('X-Notification-Token'), 'test-notification-token'); + assert.deepStrictEqual(await request.json(), { userId: 'alice' }); +}); + +test('LogoutAll still succeeds when the notification socket cannot be reached', async () => { + const proc = createProcess({ notifications: { fetch: async () => { throw new Error('DO down'); } } }); + const result = await proc.ProcessFunctions['LogoutAll']({}); + assert.ok(result.Success); +});