feat(server): manage auth cookies (#8317)

packages/backend/server/src/core/auth/controller.ts

@@ -172,16 +172,19 @@ export class AuthController {
     });
   }
 
+  @Public()
   @Get('/sign-out')
   async signOut(
     @Res() res: Response,
-    @Session() session: Session,
-    @Body() { all }: { all: boolean }
+    @Session() session: Session | undefined,
+    @Query('user_id') userId: string | undefined
   ) {
-    await this.auth.signOut(
-      session.sessionId,
-      all ? undefined : session.userId
-    );
+    if (!session) {
+      return;
+    }
+
+    await this.auth.signOut(session.sessionId, userId);
+    await this.auth.refreshCookies(res, session.sessionId);
 
     res.status(HttpStatus.OK).send({});
   }

packages/backend/server/src/core/auth/guard.ts

@@ -6,7 +6,7 @@ import type {
 } from '@nestjs/common';
 import { Injectable, SetMetadata } from '@nestjs/common';
 import { ModuleRef, Reflector } from '@nestjs/core';
-import type { Request } from 'express';
+import type { Request, Response } from 'express';
 
 import {
   AuthenticationRequired,
@@ -37,7 +37,7 @@ export class AuthGuard implements CanActivate, OnModuleInit {
   async canActivate(context: ExecutionContext) {
     const { req, res } = getRequestResponseFromContext(context);
 
-    const userSession = await this.signIn(req);
+    const userSession = await this.signIn(req, res);
     if (res && userSession && userSession.expiresAt) {
       await this.auth.refreshUserSessionIfNeeded(res, userSession);
     }
@@ -59,7 +59,7 @@ export class AuthGuard implements CanActivate, OnModuleInit {
     return true;
   }
 
-  async signIn(req: Request): Promise<Session | null> {
+  async signIn(req: Request, res?: Response): Promise<Session | null> {
     if (req.session) {
       return req.session;
     }
@@ -68,7 +68,7 @@ export class AuthGuard implements CanActivate, OnModuleInit {
     parseCookies(req);
 
     // TODO(@forehalo): a cache for user session
-    const userSession = await this.auth.getUserSessionFromRequest(req);
+    const userSession = await this.auth.getUserSessionFromRequest(req, res);
 
     if (userSession) {
       req.session = {

packages/backend/server/src/core/auth/service.ts

@@ -122,35 +122,45 @@ export class AuthService implements OnApplicationBootstrap {
     sessionId: string,
     userId?: string
   ): Promise<{ user: CurrentUser; session: UserSession } | null> {
-    const userSession = await this.db.userSession.findFirst({
+    const sessions = await this.getUserSessions(sessionId);
+
+    if (!sessions.length) {
+      return null;
+    }
+
+    let userSession: UserSession | undefined;
+
+    // try read from user provided cookies.userId
+    if (userId) {
+      userSession = sessions.find(s => s.userId === userId);
+    }
+
+    // fallback to the first valid session if user provided userId is invalid
+    if (!userSession) {
+      // checked
+      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+      userSession = sessions.at(-1)!;
+    }
+
+    const user = await this.user.findUserById(userSession.userId);
+
+    if (!user) {
+      return null;
+    }
+
+    return { user: sessionUser(user), session: userSession };
+  }
+
+  async getUserSessions(sessionId: string) {
+    return this.db.userSession.findMany({
       where: {
         sessionId,
-        userId,
-      },
-      select: {
-        id: true,
-        sessionId: true,
-        userId: true,
-        createdAt: true,
-        expiresAt: true,
-        user: true,
+        OR: [{ expiresAt: { gt: new Date() } }, { expiresAt: null }],
       },
       orderBy: {
         createdAt: 'asc',
       },
     });
-
-    // no such session
-    if (!userSession) {
-      return null;
-    }
-
-    // user session expired
-    if (userSession.expiresAt && userSession.expiresAt <= new Date()) {
-      return null;
-    }
-
-    return { user: sessionUser(userSession.user), session: userSession };
   }
 
   async createUserSession(
@@ -309,6 +319,25 @@ export class AuthService implements OnApplicationBootstrap {
     this.setUserCookie(res, userId);
   }
 
+  async refreshCookies(res: Response, sessionId?: string) {
+    if (sessionId) {
+      const users = await this.getUserList(sessionId);
+      const candidateUser = users.at(-1);
+
+      if (candidateUser) {
+        this.setUserCookie(res, candidateUser.id);
+        return;
+      }
+    }
+
+    this.clearCookies(res);
+  }
+
+  private clearCookies(res: Response<any, Record<string, any>>) {
+    res.clearCookie(AuthService.sessionCookieName);
+    res.clearCookie(AuthService.userCookieName);
+  }
+
   setUserCookie(res: Response, userId: string) {
     res.cookie(AuthService.userCookieName, userId, {
       ...this.cookieOptions,
@@ -319,14 +348,28 @@ export class AuthService implements OnApplicationBootstrap {
     });
   }
 
-  async getUserSessionFromRequest(req: Request) {
+  async getUserSessionFromRequest(req: Request, res?: Response) {
     const { sessionId, userId } = this.getSessionOptionsFromRequest(req);
 
     if (!sessionId) {
       return null;
     }
 
-    return this.getUserSession(sessionId, userId);
+    const session = await this.getUserSession(sessionId, userId);
+
+    if (res) {
+      if (session) {
+        // set user id cookie for fast authentication
+        if (!userId || userId !== session.user.id) {
+          this.setUserCookie(res, session.user.id);
+        }
+      } else if (sessionId) {
+        // clear invalid cookies.session and cookies.userId
+        this.clearCookies(res);
+      }
+    }
+
+    return session;
   }
 
   async changePassword(

packages/backend/server/tests/auth/controller.spec.ts

@@ -161,12 +161,155 @@ test('should be able to sign out', async t => {
   t.falsy(session.user);
 });
 
-test('should not be able to sign out if not signed in', async t => {
-  const { app } = t.context;
+test('should be able to correct user id cookie', async t => {
+  const { app, u1 } = t.context;
 
-  await request(app.getHttpServer())
-    .get('/api/auth/sign-out')
-    .expect(HttpStatus.UNAUTHORIZED);
+  const signInRes = await request(app.getHttpServer())
+    .post('/api/auth/sign-in')
+    .send({ email: u1.email, password: '1' })
+    .expect(200);
 
-  t.assert(true);
+  const cookie = sessionCookie(signInRes.headers);
+
+  let session = await request(app.getHttpServer())
+    .get('/api/auth/session')
+    .set('cookie', cookie)
+    .expect(200);
+
+  let userIdCookie = session.get('Set-Cookie')?.find(c => {
+    return c.startsWith(`${AuthService.userCookieName}=`);
+  });
+
+  t.true(userIdCookie?.startsWith(`${AuthService.userCookieName}=${u1.id}`));
+
+  session = await request(app.getHttpServer())
+    .get('/api/auth/session')
+    .set('cookie', `${cookie};${AuthService.userCookieName}=invalid_user_id`)
+    .expect(200);
+
+  userIdCookie = session.get('Set-Cookie')?.find(c => {
+    return c.startsWith(`${AuthService.userCookieName}=`);
+  });
+
+  t.true(userIdCookie?.startsWith(`${AuthService.userCookieName}=${u1.id}`));
+  t.is(session.body.user.id, u1.id);
+});
+
+// multiple accounts session tests
+test('should be able to sign in another account in one session', async t => {
+  const { app, u1, auth } = t.context;
+
+  const u2 = await auth.signUp('u3@affine.pro', '3');
+
+  // sign in u1
+  const signInRes = await request(app.getHttpServer())
+    .post('/api/auth/sign-in')
+    .send({ email: u1.email, password: '1' })
+    .expect(200);
+
+  const cookie = sessionCookie(signInRes.headers);
+
+  // avoid create session at the exact same time, leads to same random session users order
+  await new Promise(resolve => setTimeout(resolve, 1));
+
+  // sign in u2 in the same session
+  await request(app.getHttpServer())
+    .post('/api/auth/sign-in')
+    .set('cookie', cookie)
+    .send({ email: u2.email, password: '3' })
+    .expect(200);
+
+  // list [u1, u2]
+  const sessions = await request(app.getHttpServer())
+    .get('/api/auth/sessions')
+    .set('cookie', cookie)
+    .expect(200);
+
+  t.is(sessions.body.users.length, 2);
+  t.is(sessions.body.users[0].id, u1.id);
+  t.is(sessions.body.users[1].id, u2.id);
+
+  // default to latest signed in user: u2
+  let session = await request(app.getHttpServer())
+    .get('/api/auth/session')
+    .set('cookie', cookie)
+    .expect(200);
+
+  t.is(session.body.user.id, u2.id);
+
+  // switch to u1
+  session = await request(app.getHttpServer())
+    .get('/api/auth/session')
+    .set('cookie', `${cookie};${AuthService.userCookieName}=${u1.id}`)
+    .expect(200);
+
+  t.is(session.body.user.id, u1.id);
+});
+
+test('should be able to sign out multiple accounts in one session', async t => {
+  const { app, u1, auth } = t.context;
+
+  const u2 = await auth.signUp('u4@affine.pro', '4');
+
+  // sign in u1
+  const signInRes = await request(app.getHttpServer())
+    .post('/api/auth/sign-in')
+    .send({ email: u1.email, password: '1' })
+    .expect(200);
+
+  const cookie = sessionCookie(signInRes.headers);
+
+  await new Promise(resolve => setTimeout(resolve, 1));
+
+  // sign in u2 in the same session
+  await request(app.getHttpServer())
+    .post('/api/auth/sign-in')
+    .set('cookie', cookie)
+    .send({ email: u2.email, password: '4' })
+    .expect(200);
+
+  // sign out u2
+  let signOut = await request(app.getHttpServer())
+    .get(`/api/auth/sign-out?user_id=${u2.id}`)
+    .set('cookie', `${cookie};${AuthService.userCookieName}=${u2.id}`)
+    .expect(200);
+
+  // auto switch to u1 after sign out u2
+  const userIdCookie = signOut.get('Set-Cookie')?.find(c => {
+    return c.startsWith(`${AuthService.userCookieName}=`);
+  });
+
+  t.true(userIdCookie?.startsWith(`${AuthService.userCookieName}=${u1.id}`));
+
+  // list [u1]
+  const session = await request(app.getHttpServer())
+    .get('/api/auth/session')
+    .set('cookie', cookie)
+    .expect(200);
+
+  t.is(session.body.user.id, u1.id);
+
+  // sign in u2 in the same session
+  await request(app.getHttpServer())
+    .post('/api/auth/sign-in')
+    .set('cookie', cookie)
+    .send({ email: u2.email, password: '4' })
+    .expect(200);
+
+  // sign out all account in session
+  signOut = await request(app.getHttpServer())
+    .get('/api/auth/sign-out')
+    .set('cookie', cookie)
+    .expect(200);
+
+  t.true(
+    signOut
+      .get('Set-Cookie')
+      ?.some(c => c.startsWith(`${AuthService.sessionCookieName}=;`))
+  );
+  t.true(
+    signOut
+      .get('Set-Cookie')
+      ?.some(c => c.startsWith(`${AuthService.userCookieName}=;`))
+  );
 });

packages/backend/server/tests/auth/service.spec.ts

@@ -202,16 +202,17 @@ test('should be able to signout multi accounts session', async t => {
   t.is(list.length, 1);
   t.is(list[0]!.id, u2.id);
 
-  const u1Session = await auth.getUserSession(session.id, u1.id);
+  const u2Session = await auth.getUserSession(session.id, u1.id);
 
-  t.is(u1Session, null);
+  t.is(u2Session?.session.sessionId, session.id);
+  t.is(u2Session?.user.id, u2.id);
 
   await auth.signOut(session.id, u2.id);
   list = await auth.getUserList(session.id);
 
   t.is(list.length, 0);
 
-  const u2Session = await auth.getUserSession(session.id, u2.id);
+  const nullSession = await auth.getUserSession(session.id, u2.id);
 
-  t.is(u2Session, null);
+  t.is(nullSession, null);
 });