fix(core): selfhosted auth handling (#15295)

fix #15284
fix #15266
fix #15268
fix #15267

#### PR Dependency Tree


* **PR #15295** 👈

This tree was auto-generated by
[Charcoal](https://github.com/danerwilliams/charcoal)

Cargo.lock

@@ -76,9 +76,9 @@ dependencies = [
 
 [[package]]
 name = "affine_doc_loader"
-version = "0.1.2"
+version = "0.1.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6af9fa0b8359900e4c150b717e837eec813f29ee6981fa72d83e671974c0f78f"
+checksum = "1bd208d52725dd0c63583b171ffc3121a60d7686bacedfee5c135f9c0c58d00b"
 dependencies = [
  "nanoid 0.5.0",
  "pulldown-cmark 0.13.1",
@@ -90,13 +90,14 @@ dependencies = [
 
 [[package]]
 name = "affine_importer"
-version = "0.1.1"
+version = "0.1.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0a698069c4f0ae26bbc3c07deaf84bb0d55b0ed80dae44d6879ca1bf63eee8e9"
+checksum = "5826d670f6d43faa7809ef6a06d4a8c5ff493b970399c5c3c5659d25e572746b"
 dependencies = [
  "affine_doc_loader",
  "chrono",
  "nanoid 0.5.0",
+ "pulldown-cmark 0.13.1",
  "serde",
  "serde_json",
  "sha2 0.11.0",

Cargo.toml

@@ -16,8 +16,8 @@ resolver = "3"
   [workspace.dependencies]
   aes-gcm = "0.10"
   affine_common = { path = "./packages/common/native" }
-  affine_doc_loader = "0.1.2"
+  affine_doc_loader = "0.1.3"
-  affine_importer = "0.1.1"
+  affine_importer = "0.1.2"
   affine_nbstore = { path = "./packages/frontend/native/nbstore" }
   affine_preview = { version = "0.1.0", default-features = false }
   anyhow = "1"

packages/common/graphql/src/graphql/get-current-user.gql

@@ -5,8 +5,7 @@ query getCurrentUser {
     email
     emailVerified
     avatarUrl
-    token {
+    hasPassword
-      sessionToken
+    features
-    }
   }
 }

packages/common/graphql/src/graphql/index.ts

@@ -1967,12 +1967,10 @@ export const getCurrentUserQuery = {
     email
     emailVerified
     avatarUrl
-    token {
+    hasPassword
-      sessionToken
+    features
-    }
   }
 }`,
-  deprecations: ["'token' is deprecated: use auth session exchange instead"],
 };
 
 export const getDocCreatedByUpdatedByListQuery = {

packages/common/graphql/src/schema.ts

@@ -6519,7 +6519,8 @@ export type GetCurrentUserQuery = {
     email: string;
     emailVerified: boolean;
     avatarUrl: string | null;
-    token: { __typename?: 'tokenType'; sessionToken: string | null };
+    hasPassword: boolean | null;
+    features: Array<FeatureType>;
   } | null;
 };
 

packages/common/nbstore/src/connection/connection.ts

@@ -186,17 +186,27 @@ export abstract class AutoReconnectConnection<
         return;
       }
 
+      if (signal?.aborted) {
+        reject(signal.reason);
+        return;
+      }
+
       const off = this.onStatusChanged(status => {
         if (status === 'connected') {
           resolve();
-          off();
+          cleanup();
         }
       });
 
-      signal?.addEventListener('abort', reason => {
+      const onAbort = () => {
-        reject(reason);
+        reject(signal?.reason);
+        cleanup();
+      };
+      const cleanup = () => {
         off();
-      });
+        signal?.removeEventListener('abort', onAbort);
+      };
+      signal?.addEventListener('abort', onAbort, { once: true });
     });
   }
 

packages/common/nbstore/src/realtime/__tests__/manager.spec.ts

@@ -41,8 +41,9 @@ class FakeSocket {
   }
 }
 
-const { resetSharedConnection } = vi.hoisted(() => ({
+const { resetSharedConnection, waitForConnected } = vi.hoisted(() => ({
   resetSharedConnection: vi.fn(),
+  waitForConnected: vi.fn(async () => {}),
 }));
 const socket = new FakeSocket();
 
@@ -56,7 +57,7 @@ vi.mock('../../impls/cloud/socket', () => ({
 
     connect() {}
 
-    async waitForConnected() {}
+    waitForConnected = waitForConnected;
 
     disconnect() {
       socket.disconnected = true;
@@ -74,6 +75,8 @@ beforeEach(() => {
   socket.connected = true;
   socket.disconnected = false;
   resetSharedConnection.mockClear();
+  waitForConnected.mockReset();
+  waitForConnected.mockResolvedValue(undefined);
 });
 
 test('getRealtimeInputKey is deterministic for realtime subscription inputs', () => {
@@ -157,7 +160,7 @@ test('non-bootstrap request still requires authenticated context', async () => {
   );
 });
 
-test('request rejects when aborted', async () => {
+test('request rejects when connection times out or is aborted', async () => {
   const manager = new RealtimeManager();
   manager.setContext({
     endpoint: 'http://server',
@@ -175,6 +178,14 @@ test('request rejects when aborted', async () => {
   controller.abort();
 
   await expect(request).rejects.toThrow('Realtime request aborted');
+
+  waitForConnected.mockImplementationOnce(() => new Promise(() => {}));
+  await expect(
+    manager.request('notification.count.get', {}, { timeoutMs: 1 })
+  ).rejects.toMatchObject({
+    name: 'RealtimeRequestTimeout',
+    message: 'Realtime request timed out: notification.count.get',
+  });
 });
 
 test('subscribe routes events by topic and stable input key', async () => {

packages/common/nbstore/src/realtime/manager.ts

@@ -79,8 +79,8 @@ export class RealtimeManager {
     input: RealtimeRequestInputOf<Op>,
     options?: { timeoutMs?: number; signal?: AbortSignal }
   ): Promise<RealtimeRequestOutputOf<Op>> {
-    const socket = await this.connect(op === 'user.profile.get');
     const timeoutMs = options?.timeoutMs ?? DEFAULT_REQUEST_TIMEOUT;
+    const connectAbort = new AbortController();
     let timeoutId: ReturnType<typeof setTimeout> | undefined;
     let abortHandler: (() => void) | undefined;
     const abort = () => {
@@ -104,6 +104,12 @@ export class RealtimeManager {
       options?.signal?.addEventListener('abort', abortHandler, { once: true });
     });
 
+    try {
+      const socket = await Promise.race([
+        this.connect(op === 'user.profile.get', connectAbort.signal),
+        timeout,
+        aborted,
+      ]);
       const ack = await Promise.race([
         socket.emitWithAck('realtime:request', {
           op,
@@ -112,20 +118,22 @@ export class RealtimeManager {
         }),
         timeout,
         aborted,
-    ]).finally(() => {
+      ]);
-      if (timeoutId) {
-        clearTimeout(timeoutId);
-      }
-      if (abortHandler) {
-        options?.signal?.removeEventListener('abort', abortHandler);
-      }
-    });
 
       if ('error' in ack) {
         throw rejectAck(ack.error);
       }
 
       return ack.data as unknown as RealtimeRequestOutputOf<Op>;
+    } finally {
+      connectAbort.abort();
+      if (timeoutId) {
+        clearTimeout(timeoutId);
+      }
+      if (abortHandler) {
+        options?.signal?.removeEventListener('abort', abortHandler);
+      }
+    }
   }
 
   subscribe<Topic extends RealtimeTopicName>(
@@ -224,7 +232,7 @@ export class RealtimeManager {
     };
   }
 
-  private async connect(allowUnauthenticated = false) {
+  private async connect(allowUnauthenticated = false, signal?: AbortSignal) {
     if (
       !this.context?.endpoint ||
       (!this.context.authenticated && !allowUnauthenticated)
@@ -245,7 +253,7 @@ export class RealtimeManager {
       this.socketConnection.connect();
     }
 
-    await this.socketConnection.waitForConnected();
+    await this.socketConnection.waitForConnected(signal);
     this.socketConnection.inner.socket.off('realtime:event', this.handleEvent);
     this.socketConnection.inner.socket.on('realtime:event', this.handleEvent);
     this.socketConnection.inner.socket.off('connect', this.handleReconnect);

packages/frontend/core/src/components/providers/workspace-side-effects.tsx

@@ -1,4 +1,4 @@
-import { toast } from '@affine/component';
+import { notify, toast } from '@affine/component';
 import {
   pushGlobalLoadingEventAtom,
   resolveGlobalLoadingEventAtom,
@@ -15,6 +15,7 @@ import {
   AuthService,
   EventSourceService,
   GraphQLService,
+  RealtimeService,
 } from '@affine/core/modules/cloud';
 import {
   GlobalDialogService,
@@ -39,6 +40,7 @@ import {
   fromPromise,
   onStart,
   throwIfAborted,
+  useLiveData,
   useService,
   useServices,
 } from '@toeverything/infra';
@@ -142,6 +144,20 @@ export const WorkspaceSideEffects = () => {
   const eventSourceService = useService(EventSourceService);
   const authService = useService(AuthService);
   const nbstoreService = useService(NbstoreService);
+  const realtimeConnectionError = useLiveData(
+    useService(RealtimeService).connectionError$
+  );
+
+  useEffect(() => {
+    if (!realtimeConnectionError) return;
+    notify.warning(
+      {
+        title: t['com.affine.realtime.connection-error.title'](),
+        message: t['com.affine.realtime.connection-error.message'](),
+      },
+      { id: `realtime-connection-error:${realtimeConnectionError.endpoint}` }
+    );
+  }, [realtimeConnectionError, t]);
 
   useEffect(() => {
     const dispose = setupAIProvider(

packages/frontend/core/src/modules/cloud/entities/session.ts

@@ -93,6 +93,14 @@ export class AuthSession extends Entity {
     )
   );
 
+  async revalidateOnce() {
+    const sessionInfo = await this.getSession();
+    if (!isEqual(this.store.getCachedAuthSession(), sessionInfo)) {
+      this.store.setCachedAuthSession(sessionInfo);
+    }
+    return sessionInfo;
+  }
+
   private async getSession(): Promise<AuthSessionInfo | null> {
     try {
       const session = await this.store.fetchSession();

packages/frontend/core/src/modules/cloud/services/auth.ts

@@ -166,7 +166,7 @@ export class AuthService extends Service {
     try {
       await this.store.signInMagicLink(email, token);
 
-      this.session.revalidate();
+      await this.session.revalidateOnce();
       track.$.$.auth.signedIn({ method });
     } catch (e) {
       track.$.$.auth.signInFail({
@@ -218,7 +218,7 @@ export class AuthService extends Service {
         provider
       );
 
-      this.session.revalidate();
+      await this.session.revalidateOnce();
 
       track.$.$.auth.signedIn({ method: 'oauth', provider });
       return { redirectUri };
@@ -251,7 +251,7 @@ export class AuthService extends Service {
   async signInOpenAppSignInCode(code: string) {
     await this.store.signInOpenAppSignInCode(code);
 
-    this.session.revalidate();
+    await this.session.revalidateOnce();
   }
 
   async signInPassword(credential: {
@@ -266,8 +266,10 @@ export class AuthService extends Service {
       const user = await this.store.signInPassword(credential);
       if (user) {
         this.store.setCachedSignInUser(user);
-      }
         this.session.revalidate();
+      } else {
+        await this.session.revalidateOnce();
+      }
       track.$.$.auth.signedIn({ method: 'password' });
     } catch (e) {
       track.$.$.auth.signInFail({

packages/frontend/core/src/modules/cloud/services/realtime.ts

@@ -10,6 +10,12 @@ import type { ServersService } from './servers';
 
 @OnEvent(ApplicationStarted, service => service.onApplicationStarted)
 export class RealtimeService extends Service {
+  readonly connectionError$ = new LiveData<{
+    endpoint: string;
+    error: unknown;
+  } | null>(null);
+  private contextGeneration = 0;
+
   private readonly currentServer$ =
     this.globalContextService.globalContext.serverId.$.selector(id =>
       id
@@ -43,7 +49,8 @@ export class RealtimeService extends Service {
     super();
 
     const subscription = this.currentServer$.subscribe(context => {
-      this.nbstoreService.realtime.configure(context).catch(error => {
+      const generation = ++this.contextGeneration;
+      this.configure(context, generation).catch(error => {
         console.error('Failed to configure realtime context', error);
       });
     });
@@ -51,4 +58,35 @@ export class RealtimeService extends Service {
   }
 
   onApplicationStarted() {}
+
+  private async configure(
+    context: {
+      endpoint: string;
+      authenticated: boolean;
+      isSelfHosted: boolean;
+    },
+    generation: number
+  ) {
+    await this.nbstoreService.realtime.configure(context);
+    if (generation !== this.contextGeneration) return;
+    if (!context.endpoint || !context.authenticated) {
+      this.connectionError$.next(null);
+      return;
+    }
+
+    try {
+      await this.nbstoreService.realtime.request(
+        'user.profile.get',
+        {},
+        { timeoutMs: 10_000 }
+      );
+      if (generation === this.contextGeneration) {
+        this.connectionError$.next(null);
+      }
+    } catch (error) {
+      if (generation === this.contextGeneration) {
+        this.connectionError$.next({ endpoint: context.endpoint, error });
+      }
+    }
+  }
 }

packages/frontend/core/src/modules/cloud/stores/auth.spec.ts

@@ -9,14 +9,14 @@ import { describe, expect, test, vi } from 'vitest';
 
 function createStore({
   fetch,
-  request,
+  gql = vi.fn(),
 }: {
   fetch: (input: string, init?: RequestInit) => Promise<Response>;
-  request: (op: string, input: object) => Promise<unknown>;
+  gql?: () => Promise<unknown>;
 }) {
   const framework = new Framework();
   framework.service(FetchService, { fetch } as any);
-  framework.service(GraphQLService, {} as any);
+  framework.service(GraphQLService, { gql } as any);
   framework.impl(GlobalState, {} as any);
   framework.service(ServerService, {
     server: {
@@ -26,7 +26,7 @@ function createStore({
   } as any);
   framework.impl(AuthProvider, {} as any);
   framework.service(NbstoreService, {
-    realtime: { request },
+    realtime: {},
   } as any);
   framework.store(AuthStore, [
     FetchService,
@@ -57,7 +57,7 @@ describe('AuthStore', () => {
         ok: true,
         json: async () => ({ registered: true, methods: {} }),
       });
-    const store = createStore({ fetch, request: vi.fn() });
+    const store = createStore({ fetch });
 
     await expect(store.checkUserByEmail('user@affine.pro')).resolves.toEqual(
       validResponse
@@ -70,7 +70,7 @@ describe('AuthStore', () => {
     });
   });
 
-  test('loads account profile from realtime after auth session bootstrap', async () => {
+  test('loads account profile from scoped GraphQL after auth session bootstrap', async () => {
     const authMethods = {
       password: { bound: true },
       oauth: { bound: false, providers: [] },
@@ -90,8 +90,8 @@ describe('AuthStore', () => {
       }
       throw new Error(`Unexpected request: ${input}`);
     });
-    const request = vi.fn(async () => ({
+    const gql = vi.fn(async () => ({
-      user: {
+      currentUser: {
         id: 'u1',
         email: 'u1@affine.pro',
         name: 'User',
@@ -101,7 +101,7 @@ describe('AuthStore', () => {
         features: ['Admin'],
       },
     }));
-    const store = createStore({ fetch, request });
+    const store = createStore({ fetch, gql });
 
     await expect(store.fetchSession()).resolves.toEqual({
       user: {
@@ -115,17 +115,17 @@ describe('AuthStore', () => {
         authMethods,
       },
     });
-    expect(request).toHaveBeenCalledWith('user.profile.get', {});
+    expect(gql).toHaveBeenCalledOnce();
   });
 
-  test('rejects mismatched realtime profile and auth session', async () => {
+  test('rejects mismatched GraphQL profile and auth session', async () => {
     const fetch = vi.fn(async () => {
       return {
         json: async () => ({ user: { id: 'u1' } }),
       } as Response;
     });
-    const request = vi.fn(async () => ({
+    const gql = vi.fn(async () => ({
-      user: {
+      currentUser: {
         id: 'u2',
         email: 'u2@affine.pro',
         name: 'User',
@@ -135,10 +135,10 @@ describe('AuthStore', () => {
         features: [],
       },
     }));
-    const store = createStore({ fetch, request });
+    const store = createStore({ fetch, gql });
 
     await expect(store.fetchSession()).rejects.toThrow(
-      'Realtime user profile does not match auth session'
+      'User profile does not match auth session'
     );
   });
 });

packages/frontend/core/src/modules/cloud/stores/auth.ts

@@ -1,5 +1,6 @@
 import {
   deleteAccountMutation,
+  getCurrentUserQuery,
   removeAvatarMutation,
   ServerDeploymentType,
   updateUserProfileMutation,
@@ -101,12 +102,11 @@ export class AuthStore extends Store {
     const session = await this.fetchAuthSession();
     if (!session.user) return { user: null };
 
-    const { user } = await this.nbstoreService.realtime.request(
+    const { currentUser: user } = await this.gqlService.gql({
-      'user.profile.get',
+      query: getCurrentUserQuery,
-      {}
+    });
-    );
     if (!user || user.id !== session.user.id) {
-      throw new Error('Realtime user profile does not match auth session');
+      throw new Error('User profile does not match auth session');
     }
     const authMethods = await this.fetchAuthMethods();
     return { user: { ...user, authMethods } };

packages/frontend/i18n/src/i18n.gen.ts

@@ -9278,6 +9278,14 @@ export function useAFFiNEI18N(): {
       * `Failed to sign out devices`
       */
     ["com.affine.settings.devices.sign-out-all-failed"](): string;
+    /**
+      * `Real-time connection failed`
+      */
+    ["com.affine.realtime.connection-error.title"](): string;
+    /**
+      * `Check that your server proxy forwards /socket.io over WebSocket or HTTP polling.`
+      */
+    ["com.affine.realtime.connection-error.message"](): string;
     /**
       * `An internal error occurred.`
       */

packages/frontend/i18n/src/resources/en.json

@@ -2318,6 +2318,8 @@
   "com.affine.settings.devices.load-failed": "Failed to load devices",
   "com.affine.settings.devices.sign-out-failed": "Failed to sign out device",
   "com.affine.settings.devices.sign-out-all-failed": "Failed to sign out devices",
+  "com.affine.realtime.connection-error.title": "Real-time connection failed",
+  "com.affine.realtime.connection-error.message": "Check that your server proxy forwards /socket.io over WebSocket or HTTP polling.",
   "error.INTERNAL_SERVER_ERROR": "An internal error occurred.",
   "error.NETWORK_ERROR": "Network error.",
   "error.TOO_MANY_REQUEST": "Too many requests.",

packages/frontend/i18n/src/resources/zh-Hans.json

@@ -2308,6 +2308,8 @@
   "com.affine.docIconPicker.placeholder": "添加图标",
   "error.INTERNAL_SERVER_ERROR": "发生内部错误。",
   "error.NETWORK_ERROR": "网络错误。",
+  "com.affine.realtime.connection-error.title": "实时同步连接失败",
+  "com.affine.realtime.connection-error.message": "请检查服务器反向代理是否通过 WebSocket 或 HTTP polling 正确转发 /socket.io。",
   "error.TOO_MANY_REQUEST": "请求过多。",
   "error.SSRF_BLOCKED_ERROR": "无效的 URL",
   "error.RESPONSE_TOO_LARGE_ERROR": "响应过大（{{receivedBytes}} 字节），限制为 {{limitBytes}} 字节",