fix: cookie issues in Electron (#4115)

apps/core/src/components/affine/auth/sign-in-with-password.tsx

@@ -9,10 +9,11 @@ import { useAFFiNEI18N } from '@affine/i18n/hooks';
 import { Button } from '@toeverything/components/button';
 import { useSetAtom } from 'jotai';
 // eslint-disable-next-line @typescript-eslint/no-restricted-imports
-import { signIn, useSession } from 'next-auth/react';
+import { useSession } from 'next-auth/react';
 import type { FC } from 'react';
 import { useCallback, useState } from 'react';
 
+import { signInCloud } from '../../../utils/cloud-utils';
 import type { AuthPanelProps } from './index';
 import { forgetPasswordButton } from './style.css';
 
@@ -30,7 +31,7 @@ export const SignInWithPassword: FC<AuthPanelProps> = ({
   const [passwordError, setPasswordError] = useState(false);
 
   const onSignIn = useCallback(async () => {
-    const res = await signIn('credentials', {
+    const res = await signInCloud('credentials', {
       redirect: false,
       email,
       password,

apps/core/src/utils/cloud-utils.tsx

@@ -17,26 +17,30 @@ export const signInCloud: typeof signIn = async (provider, ...rest) => {
         '_target'
       );
       return;
-    } else if (provider === 'email') {
+    } else {
       const [options, ...tail] = rest;
+      const callbackUrl =
+        runtimeConfig.serverUrlPrefix +
+        (provider === 'email' ? '/open-app/oauth-jwt' : location.pathname);
       return signIn(
         provider,
         {
           ...options,
-          callbackUrl: buildCallbackUrl('/open-app/oauth-jwt'),
+          callbackUrl: buildCallbackUrl(callbackUrl),
         },
         ...tail
       );
-    } else {
-      throw new Error('Unsupported provider');
     }
   } else {
     return signIn(provider, ...rest);
   }
 };
 
-export const signOutCloud: typeof signOut = async (...args) => {
-  return signOut(...args).then(result => {
+export const signOutCloud: typeof signOut = async options => {
+  return signOut({
+    ...options,
+    callbackUrl: '/',
+  }).then(result => {
     if (result) {
       startTransition(() => {
         getCurrentStore().set(refreshRootMetadataAtom);

apps/electron/src/main/deep-link.ts

@@ -2,10 +2,11 @@ import path from 'node:path';
 
 import type { App } from 'electron';
 
-import { buildType, isDev } from './config';
+import { buildType, CLOUD_BASE_URL, isDev } from './config';
 import { logger } from './logger';
 import {
   handleOpenUrlInHiddenWindow,
+  mainWindowOrigin,
   restoreOrCreateWindow,
   setCookie,
 } from './main-window';
@@ -70,24 +71,36 @@ async function handleOauthJwt(url: string) {
       mainWindow.show();
       const urlObj = new URL(url);
       const token = urlObj.searchParams.get('token');
-      const mainOrigin = new URL(mainWindow.webContents.getURL()).origin;
 
       if (!token) {
         logger.error('no token in url', url);
         return;
       }
 
+      const isSecure = CLOUD_BASE_URL.startsWith('https://');
+
       // set token to cookie
       await setCookie({
-        url: mainOrigin,
+        url: CLOUD_BASE_URL,
         httpOnly: true,
         value: token,
-        name: 'next-auth.session-token',
+        secure: true,
+        name: isSecure
+          ? '__Secure-next-auth.session-token'
+          : 'next-auth.session-token',
+        expirationDate: Math.floor(Date.now() / 1000 + 3600 * 24 * 7),
+      });
+
+      // force reset next-auth.callback-url
+      await setCookie({
+        url: CLOUD_BASE_URL,
+        httpOnly: true,
+        name: 'next-auth.callback-url',
       });
 
       // hacks to refresh auth state in the main window
       const window = await handleOpenUrlInHiddenWindow(
-        mainOrigin + '/auth/signIn'
+        mainWindowOrigin + '/auth/signIn'
       );
       uiSubjects.onFinishLogin.next({
         success: true,

apps/electron/src/main/main-window.ts

@@ -15,6 +15,8 @@ const IS_DEV: boolean =
 
 const DEV_TOOL = process.env.DEV_TOOL === 'true';
 
+export const mainWindowOrigin = process.env.DEV_SERVER_URL || 'file://.';
+
 async function createWindow() {
   logger.info('create window');
   const mainWindowState = electronWindowState({
@@ -114,7 +116,7 @@ async function createWindow() {
   /**
    * URL for main window.
    */
-  const pageUrl = process.env.DEV_SERVER_URL || 'file://.'; // see protocol.ts
+  const pageUrl = mainWindowOrigin; // see protocol.ts
 
   logger.info('loading page at', pageUrl);
 
@@ -126,35 +128,30 @@ async function createWindow() {
 }
 
 // singleton
-let browserWindow: BrowserWindow | undefined;
+let browserWindow$: Promise<BrowserWindow> | undefined;
 
 /**
  * Restore existing BrowserWindow or Create new BrowserWindow
  */
 export async function restoreOrCreateWindow() {
-  if (!browserWindow || browserWindow.isDestroyed()) {
-    browserWindow = await createWindow();
+  if (!browserWindow$ || (await browserWindow$.then(w => w.isDestroyed()))) {
+    browserWindow$ = createWindow();
   }
+  const mainWindow = await browserWindow$;
 
-  if (browserWindow.isMinimized()) {
-    browserWindow.restore();
+  if (mainWindow.isMinimized()) {
+    mainWindow.restore();
     logger.info('restore main window');
   }
-
-  return browserWindow;
+  return mainWindow;
 }
 
 export async function handleOpenUrlInHiddenWindow(url: string) {
-  const mainExposedMeta = getExposedMeta();
   const win = new BrowserWindow({
     width: 1200,
     height: 600,
     webPreferences: {
       preload: join(__dirname, './preload.js'),
-      additionalArguments: [
-        `--main-exposed-meta=` + JSON.stringify(mainExposedMeta),
-        // popup window does not need helper process, right?
-      ],
     },
     show: false,
   });
@@ -169,10 +166,6 @@ export async function handleOpenUrlInHiddenWindow(url: string) {
   return win;
 }
 
-export function reloadApp() {
-  browserWindow?.reload();
-}
-
 export async function setCookie(cookie: CookiesSetDetails): Promise<void>;
 export async function setCookie(origin: string, cookie: string): Promise<void>;
 
@@ -186,9 +179,20 @@ export async function setCookie(
       ? parseCookie(arg0, arg1)
       : arg0;
 
+  logger.info('setting cookie to main window', details);
+
   if (typeof details !== 'object') {
     throw new Error('invalid cookie details');
   }
 
   await window.webContents.session.cookies.set(details);
 }
+
+export async function getCookie(url?: string, name?: string) {
+  const window = await restoreOrCreateWindow();
+  const cookies = await window.webContents.session.cookies.get({
+    url,
+    name,
+  });
+  return cookies;
+}

apps/electron/src/main/protocol.ts

@@ -2,6 +2,8 @@ import { net, protocol, session } from 'electron';
 import { join } from 'path';
 
 import { CLOUD_BASE_URL } from './config';
+import { logger } from './logger';
+import { getCookie } from './main-window';
 
 protocol.registerSchemesAsPrivileged([
   {
@@ -70,9 +72,49 @@ export function registerProtocol() {
           'DELETE',
           'OPTIONS',
         ];
+        // replace SameSite=Lax with SameSite=None
+        const originalCookie =
+          responseHeaders['set-cookie'] || responseHeaders['Set-Cookie'];
+
+        if (originalCookie) {
+          delete responseHeaders['set-cookie'];
+          delete responseHeaders['Set-Cookie'];
+          responseHeaders['Set-Cookie'] = originalCookie.map(cookie => {
+            let newCookie = cookie.replace(/SameSite=Lax/gi, 'SameSite=None');
+
+            // if the cookie is not secure, set it to secure
+            if (!newCookie.includes('Secure')) {
+              newCookie = newCookie + '; Secure';
+            }
+            return newCookie;
+          });
+        }
       }
 
       callback({ responseHeaders });
     }
   );
+
+  session.defaultSession.webRequest.onBeforeSendHeaders((details, callback) => {
+    (async () => {
+      const url = new URL(details.url);
+      const pathname = url.pathname;
+      // if sending request to the cloud, attach the session cookie
+      if (isNetworkResource(pathname)) {
+        const cookie = await getCookie(CLOUD_BASE_URL);
+        const cookieString = cookie.map(c => `${c.name}=${c.value}`).join('; ');
+        details.requestHeaders['cookie'] = cookieString;
+      }
+      callback({
+        cancel: false,
+        requestHeaders: details.requestHeaders,
+      });
+    })().catch(e => {
+      logger.error('failed to attach cookie', e);
+      callback({
+        cancel: false,
+        requestHeaders: details.requestHeaders,
+      });
+    });
+  });
 }

apps/server/src/modules/auth/next-auth-options.ts

@@ -121,7 +121,7 @@ export const NextAuthOptionsProvider: FactoryProvider<NextAuthOptions> = {
       adapter: prismaAdapter,
       debug: !config.node.prod,
       session: {
-        strategy: config.node.prod ? 'database' : 'jwt',
+        strategy: 'jwt',
       },
       // @ts-expect-error Third part library type mismatch
       logger: console,