refactor(server): use userSession model in auth service

Author: fengmk2

packages/backend/server/src/__tests__/auth/guard.spec.ts

@@ -48,7 +48,7 @@ test.before(async t => {
   u1 = await auth.signUp('u1@affine.pro', '1');
 
   const models = app.get(Models);
-  const session = await models.session.create();
+  const session = await models.session.createSession();
   sessionId = session.id;
   await auth.createUserSession(u1.id, sessionId);
 

packages/backend/server/src/__tests__/auth/job.spec.ts

@@ -0,0 +1,47 @@
+import { ScheduleModule } from '@nestjs/schedule';
+import { TestingModule } from '@nestjs/testing';
+import { PrismaClient } from '@prisma/client';
+import test from 'ava';
+
+import { AuthModule, AuthService } from '../../core/auth';
+import { AuthCronJob } from '../../core/auth/job';
+import { createTestingModule } from '../utils';
+
+let m: TestingModule;
+let db: PrismaClient;
+
+test.before(async () => {
+  m = await createTestingModule({
+    imports: [ScheduleModule.forRoot(), AuthModule],
+  });
+
+  db = m.get(PrismaClient);
+});
+
+test.after.always(async () => {
+  await m.close();
+});
+
+test('should clean expired user sessions', async t => {
+  const auth = m.get(AuthService);
+  const job = m.get(AuthCronJob);
+  const user1 = await auth.signUp('u1@affine.pro', '1');
+  const user2 = await auth.signUp('u2@affine.pro', '1');
+  await auth.createUserSession(user1.id);
+  await auth.createUserSession(user2.id);
+  let userSessions = await db.userSession.findMany();
+  t.is(userSessions.length, 2);
+
+  // no expired sessions
+  await job.cleanExpiredUserSessions();
+  userSessions = await db.userSession.findMany();
+  t.is(userSessions.length, 2);
+
+  // clean all expired sessions
+  await db.userSession.updateMany({
+    data: { expiresAt: new Date(Date.now() - 1000) },
+  });
+  await job.cleanExpiredUserSessions();
+  userSessions = await db.userSession.findMany();
+  t.is(userSessions.length, 0);
+});

packages/backend/server/src/__tests__/auth/service.spec.ts

@@ -192,8 +192,10 @@ test('should be able to signout multi accounts session', async t => {
 
   const session = await auth.createSession();
 
-  await auth.createUserSession(u1.id, session.id);
-  await auth.createUserSession(u2.id, session.id);
+  const userSession1 = await auth.createUserSession(u1.id, session.id);
+  const userSession2 = await auth.createUserSession(u2.id, session.id);
+  t.not(userSession1.id, userSession2.id);
+  t.is(userSession1.sessionId, userSession2.sessionId);
 
   await auth.signOut(session.id, u1.id);
 

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

@@ -7,6 +7,7 @@ import { QuotaModule } from '../quota';
 import { UserModule } from '../user';
 import { AuthController } from './controller';
 import { AuthGuard, AuthWebsocketOptionsProvider } from './guard';
+import { AuthCronJob } from './job';
 import { AuthResolver } from './resolver';
 import { AuthService } from './service';
 
@@ -16,6 +17,7 @@ import { AuthService } from './service';
     AuthService,
     AuthResolver,
     AuthGuard,
+    AuthCronJob,
     AuthWebsocketOptionsProvider,
   ],
   exports: [AuthService, AuthGuard, AuthWebsocketOptionsProvider],

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

@@ -0,0 +1,14 @@
+import { Injectable } from '@nestjs/common';
+import { Cron, CronExpression } from '@nestjs/schedule';
+
+import { Models } from '../../models';
+
+@Injectable()
+export class AuthCronJob {
+  constructor(private readonly models: Models) {}
+
+  @Cron(CronExpression.EVERY_DAY_AT_MIDNIGHT)
+  async cleanExpiredUserSessions() {
+    await this.models.session.cleanExpiredUserSessions();
+  }
+}

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

@@ -1,7 +1,5 @@
 import { Injectable, OnApplicationBootstrap } from '@nestjs/common';
-import { Cron, CronExpression } from '@nestjs/schedule';
 import type { User, UserSession } from '@prisma/client';
-import { PrismaClient } from '@prisma/client';
 import type { CookieOptions, Request, Response } from 'express';
 import { assign, pick } from 'lodash-es';
 
@@ -47,7 +45,6 @@ export class AuthService implements OnApplicationBootstrap {
 
   constructor(
     private readonly config: Config,
-    private readonly db: PrismaClient,
     private readonly models: Models,
     private readonly mailer: MailService,
     private readonly feature: FeatureManagementService,
@@ -105,14 +102,9 @@ export class AuthService implements OnApplicationBootstrap {
   async signOut(sessionId: string, userId?: string) {
     // sign out all users in the session
     if (!userId) {
-      await this.models.session.delete(sessionId);
+      await this.models.session.deleteSession(sessionId);
     } else {
-      await this.db.userSession.deleteMany({
-        where: {
-          sessionId,
-          userId,
-        },
-      });
+      await this.models.session.deleteUserSession(userId, sessionId);
     }
   }
 
@@ -149,117 +141,50 @@ export class AuthService implements OnApplicationBootstrap {
   }
 
   async getUserSessions(sessionId: string) {
-    return this.db.userSession.findMany({
-      where: {
-        sessionId,
-        OR: [{ expiresAt: { gt: new Date() } }, { expiresAt: null }],
-      },
-      orderBy: {
-        createdAt: 'asc',
-      },
-    });
+    return await this.models.session.findUserSessionsBySessionId(sessionId);
   }
 
-  async createUserSession(
-    userId: string,
-    sessionId?: string,
-    ttl = this.config.auth.session.ttl
-  ) {
-    // check whether given session is valid
-    if (sessionId) {
-      const session = await this.getSession(sessionId);
-
-      if (!session) {
-        sessionId = undefined;
-      }
-    }
-
-    if (!sessionId) {
-      const session = await this.createSession();
-      sessionId = session.id;
-    }
-
-    const expiresAt = new Date(Date.now() + ttl * 1000);
-
-    return this.db.userSession.upsert({
-      where: {
-        sessionId_userId: {
-          sessionId,
-          userId,
-        },
-      },
-      update: {
-        expiresAt,
-      },
-      create: {
-        sessionId,
-        userId,
-        expiresAt,
-      },
-    });
+  async createUserSession(userId: string, sessionId?: string, ttl?: number) {
+    return await this.models.session.createOrRefreshUserSession(
+      userId,
+      sessionId,
+      ttl
+    );
   }
 
   async getUserList(sessionId: string) {
-    const sessions = await this.db.userSession.findMany({
-      where: {
-        sessionId,
-        OR: [
-          {
-            expiresAt: null,
-          },
-          {
-            expiresAt: {
-              gt: new Date(),
-            },
-          },
-        ],
-      },
-      include: {
+    const sessions = await this.models.session.findUserSessionsBySessionId(
+      sessionId,
+      {
         user: true,
-      },
-      orderBy: {
-        createdAt: 'asc',
-      },
-    });
-
+      }
+    );
     return sessions.map(({ user }) => sessionUser(user));
   }
 
   async createSession() {
-    return await this.models.session.create();
+    return await this.models.session.createSession();
   }
 
   async getSession(sessionId: string) {
-    return await this.models.session.get(sessionId);
+    return await this.models.session.getSession(sessionId);
   }
 
   async refreshUserSessionIfNeeded(
     res: Response,
-    session: UserSession,
-    ttr = this.config.auth.session.ttr
+    userSession: UserSession,
+    ttr?: number
   ): Promise<boolean> {
-    if (
-      session.expiresAt &&
-      session.expiresAt.getTime() - Date.now() > ttr * 1000
-    ) {
+    const newExpiresAt = await this.models.session.refreshUserSessionIfNeeded(
+      userSession,
+      ttr
+    );
+    if (!newExpiresAt) {
       // no need to refresh
       return false;
     }
 
-    const newExpiresAt = new Date(
-      Date.now() + this.config.auth.session.ttl * 1000
-    );
-
-    await this.db.userSession.update({
-      where: {
-        id: session.id,
-      },
-      data: {
-        expiresAt: newExpiresAt,
-      },
-    });
-
-    res.cookie(AuthService.sessionCookieName, session.sessionId, {
+    res.cookie(AuthService.sessionCookieName, userSession.sessionId, {
       expires: newExpiresAt,
       ...this.cookieOptions,
     });
@@ -268,11 +193,7 @@ export class AuthService implements OnApplicationBootstrap {
   }
 
   async revokeUserSessions(userId: string) {
-    return this.db.userSession.deleteMany({
-      where: {
-        userId,
-      },
-    });
+    return await this.models.session.deleteUserSession(userId);
   }
 
   getSessionOptionsFromRequest(req: Request) {
@@ -412,15 +333,4 @@ export class AuthService implements OnApplicationBootstrap {
           to: email,
         });
   }
-
-  @Cron(CronExpression.EVERY_DAY_AT_MIDNIGHT)
-  async cleanExpiredSessions() {
-    await this.db.userSession.deleteMany({
-      where: {
-        expiresAt: {
-          lte: new Date(),
-        },
-      },
-    });
-  }
 }

packages/backend/server/src/models/session.ts

@@ -4,14 +4,15 @@ import {
   PrismaClient,
   type Session,
   type User,
-  type UserSession as _UserSession,
+  type UserSession,
 } from '@prisma/client';
 
 export type { Session };
 
 import { Config } from '../base';
 
-export type UserSession = _UserSession & { user?: User };
+export type { UserSession };
+export type UserSessionWithUser = UserSession & { user: User };
 
 @Injectable()
 export class SessionModel {
@@ -115,10 +116,12 @@ export class SessionModel {
     return newExpiresAt;
   }
 
-  async findUserSessionsBySessionId(
+  async findUserSessionsBySessionId<
+    T extends Prisma.UserSessionInclude | undefined,
+  >(
     sessionId: string,
-    include?: Prisma.UserSessionInclude
-  ): Promise<UserSession[]> {
+    include?: T
+  ): Promise<(T extends { user: true } ? UserSessionWithUser : UserSession)[]> {
     return await this.db.userSession.findMany({
       where: {
         sessionId,
@@ -127,7 +130,7 @@ export class SessionModel {
       orderBy: {
         createdAt: 'asc',
       },
-      include,
+      include: include as Prisma.UserSessionInclude,
     });
   }