Add LogoutAll: sign a user out of the backend on every device

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

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<string, Set<WebSocket>>;
   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<string, Set<WebSocket>>();
     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<Response> {
     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<Response> {
+    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<Response> {
+    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<Response> {
     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<boolean> {
+    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 {

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<void> => {
+    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<void> => {
     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<Result> => {
+      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<Result> => {
       ThrowErrorIfFailed(this.CheckParams(Data, {}));
       const ResponseData = {

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<void> => {
+  ThrowErrorIfFailed(await XMOJDatabase.Delete("session_token", {
+    user_id: Username
+  }));
+};
+
 export const RevokeSessionToken = async (XMOJDatabase: Database, Token: string): Promise<void> => {
   ThrowErrorIfFailed(await XMOJDatabase.Delete("session_token", {
     token_hash: HashSessionToken(Token)

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));
     }
 

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);
+});