feat(server): token encryption

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

@@ -108,17 +108,17 @@ export class AuthController {
     @Query('redirect_uri') redirectUri?: string
   ) {
     const session = await this.auth.signOut(
-      req.cookies[this.auth.sessionCookieName],
-      parseAuthUserSeqNum(req.headers[this.auth.authUserSeqCookieName])
+      req.cookies[AuthService.sessionCookieName],
+      parseAuthUserSeqNum(req.headers[AuthService.authUserSeqHeaderName])
     );
 
     if (session) {
-      res.cookie(this.auth.sessionCookieName, session.id, {
+      res.cookie(AuthService.sessionCookieName, session.id, {
         expires: session.expiresAt ?? void 0, // expiredAt is `string | null`
         ...this.auth.cookieOptions,
       });
     } else {
-      res.clearCookie(this.auth.sessionCookieName);
+      res.clearCookie(AuthService.sessionCookieName);
     }
 
     if (redirectUri) {
@@ -188,7 +188,7 @@ export class AuthController {
   @Public()
   @Get('/sessions')
   async currentSessionUsers(@Req() req: Request) {
-    const token = req.cookies[this.auth.sessionCookieName];
+    const token = req.cookies[AuthService.sessionCookieName];
     if (!token) {
       return {
         users: [],

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

@@ -16,7 +16,7 @@ import { AuthService, parseAuthUserSeqNum } from './service';
 
 function extractTokenFromHeader(authorization: string) {
   if (!/^Bearer\s/i.test(authorization)) {
-    return null;
+    return;
   }
 
   return authorization.substring(7);
@@ -38,11 +38,10 @@ export class AuthGuard implements CanActivate, OnModuleInit {
 
   async canActivate(context: ExecutionContext) {
     const { req } = getRequestResponseFromContext(context);
-    const token = req.headers.authorization;
 
     // check cookie
     let sessionToken: string | undefined =
-      req.cookies[this.auth.sessionCookieName];
+      req.cookies[AuthService.sessionCookieName];
 
     // backward compatibility for client older then 0.12
     // TODO: remove
@@ -55,28 +54,22 @@ export class AuthGuard implements CanActivate, OnModuleInit {
         ];
     }
 
+    if (!sessionToken && req.headers.authorization) {
+      sessionToken = extractTokenFromHeader(req.headers.authorization);
+    }
+
     if (sessionToken) {
-      const user = await this.auth.getUser(
-        sessionToken,
-        parseAuthUserSeqNum(req.headers[this.auth.authUserSeqCookieName])
+      const userSeq = parseAuthUserSeqNum(
+        req.headers[AuthService.authUserSeqHeaderName]
       );
 
+      const user = await this.auth.getUser(sessionToken, userSeq);
+
       if (user) {
         req.user = user;
       }
     }
 
-    // check authorization token
-    else if (token) {
-      const accessToken = extractTokenFromHeader(token);
-      if (accessToken) {
-        const user = await this.auth.getUser(accessToken);
-        if (user) {
-          req.user = user;
-        }
-      }
-    }
-
     // api is public
     const isPublic = this.reflector.get<boolean>(
       'isPublic',

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

@@ -5,12 +5,16 @@ import {
   NotFoundException,
   OnApplicationBootstrap,
 } from '@nestjs/common';
-import { hash, verify } from '@node-rs/argon2';
 import { PrismaClient, type User } from '@prisma/client';
 import type { CookieOptions, Request, Response } from 'express';
 import { assign, omit } from 'lodash-es';
 
-import { Config, MailService } from '../../fundamentals';
+import {
+  Config,
+  CryptoHelper,
+  MailService,
+  SessionCache,
+} from '../../fundamentals';
 import { FeatureManagementService } from '../features/management';
 import { UserService } from '../user/service';
 import type { CurrentUser } from './current-user';
@@ -47,15 +51,17 @@ export class AuthService implements OnApplicationBootstrap {
     domain: this.config.host,
     secure: this.config.https,
   };
-  readonly sessionCookieName = 'sid';
-  readonly authUserSeqCookieName = 'authuser';
+  static readonly sessionCookieName = 'sid';
+  static readonly authUserSeqHeaderName = 'x-auth-user';
 
   constructor(
     private readonly config: Config,
     private readonly db: PrismaClient,
     private readonly mailer: MailService,
     private readonly feature: FeatureManagementService,
-    private readonly user: UserService
+    private readonly user: UserService,
+    private readonly crypto: CryptoHelper,
+    private readonly cache: SessionCache
   ) {}
 
   async onApplicationBootstrap() {
@@ -81,7 +87,7 @@ export class AuthService implements OnApplicationBootstrap {
       throw new BadRequestException('Email was taken');
     }
 
-    const hashedPassword = await hash(password);
+    const hashedPassword = await this.crypto.encryptPassword(password);
 
     return this.user
       .createUser({
@@ -105,7 +111,10 @@ export class AuthService implements OnApplicationBootstrap {
       );
     }
 
-    const passwordMatches = await verify(user.password, password);
+    const passwordMatches = await this.crypto.verifyPassword(
+      password,
+      user.password
+    );
 
     if (!passwordMatches) {
       throw new NotAcceptableException('Incorrect Password');
@@ -114,6 +123,22 @@ export class AuthService implements OnApplicationBootstrap {
     return sessionUser(user);
   }
 
+  async getUserWithCache(token: string, seq = 0) {
+    const cacheKey = `session:${token}:${seq}`;
+    let user = await this.cache.get<CurrentUser | null>(cacheKey);
+    if (user) {
+      return user;
+    }
+
+    user = await this.getUser(token, seq);
+
+    if (user) {
+      await this.cache.set(cacheKey, user);
+    }
+
+    return user;
+  }
+
   async getUser(token: string, seq = 0): Promise<CurrentUser | null> {
     const session = await this.getSession(token);
 
@@ -272,10 +297,10 @@ export class AuthService implements OnApplicationBootstrap {
   async setCookie(req: Request, res: Response, user: { id: string }) {
     const session = await this.createUserSession(
       user,
-      req.cookies[this.sessionCookieName]
+      req.cookies[AuthService.sessionCookieName]
     );
 
-    res.cookie(this.sessionCookieName, session.sessionId, {
+    res.cookie(AuthService.sessionCookieName, session.sessionId, {
       expires: session.expiresAt ?? void 0,
       ...this.cookieOptions,
     });
@@ -292,7 +317,7 @@ export class AuthService implements OnApplicationBootstrap {
       throw new BadRequestException('Invalid email');
     }
 
-    const hashedPassword = await hash(newPassword);
+    const hashedPassword = await this.crypto.encryptPassword(newPassword);
 
     return this.db.user.update({
       where: {

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

@@ -3,6 +3,8 @@ import { randomUUID } from 'node:crypto';
 import { Injectable } from '@nestjs/common';
 import { PrismaClient } from '@prisma/client';
 
+import { CryptoHelper } from '../../fundamentals/helpers';
+
 export enum TokenType {
   SignIn,
   VerifyEmail,
@@ -13,7 +15,10 @@ export enum TokenType {
 
 @Injectable()
 export class TokenService {
-  constructor(private readonly db: PrismaClient) {}
+  constructor(
+    private readonly db: PrismaClient,
+    private readonly crypto: CryptoHelper
+  ) {}
 
   async createToken(
     type: TokenType,
@@ -31,8 +36,7 @@ export class TokenService {
       },
     });
 
-    // TODO: encrypt the token
-    return token;
+    return this.crypto.encrypt(token);
   }
 
   async verifyToken(
@@ -46,6 +50,7 @@ export class TokenService {
       keep?: boolean;
     } = {}
   ) {
+    token = this.crypto.decrypt(token);
     const record = await this.db.verificationToken.findUnique({
       where: {
         type_token: {

packages/backend/server/src/data/migrations/1605053000403-self-host-admin.ts

@@ -1,13 +1,15 @@
 import { ModuleRef } from '@nestjs/core';
-import { hash } from '@node-rs/argon2';
 import { PrismaClient } from '@prisma/client';
 
-import { Config } from '../../fundamentals';
+import { UserService } from '../../core/user';
+import { Config, CryptoHelper } from '../../fundamentals';
 
 export class SelfHostAdmin1605053000403 {
   // do the migration
-  static async up(db: PrismaClient, ref: ModuleRef) {
+  static async up(_db: PrismaClient, ref: ModuleRef) {
     const config = ref.get(Config, { strict: false });
+    const crypto = ref.get(CryptoHelper, { strict: false });
+    const user = ref.get(UserService, { strict: false });
     if (config.isSelfhosted) {
       if (
         !process.env.AFFINE_ADMIN_EMAIL ||
@@ -17,13 +19,12 @@ export class SelfHostAdmin1605053000403 {
           'You have to set AFFINE_ADMIN_EMAIL and AFFINE_ADMIN_PASSWORD environment variables to generate the initial user for self-hosted AFFiNE Server.'
         );
       }
-      await db.user.create({
-        data: {
-          name: 'AFFINE First User',
-          email: process.env.AFFINE_ADMIN_EMAIL,
-          emailVerifiedAt: new Date(),
-          password: await hash(process.env.AFFINE_ADMIN_PASSWORD),
-        },
+      await user.findOrCreateUser(process.env.AFFINE_ADMIN_EMAIL, {
+        name: 'AFFINE First User',
+        emailVerifiedAt: new Date(),
+        password: await crypto.encryptPassword(
+          process.env.AFFINE_ADMIN_PASSWORD
+        ),
       });
     }
   }

packages/backend/server/src/fundamentals/helpers/__tests__/crypto.spec.ts

@@ -0,0 +1,105 @@
+import { createPrivateKey, createPublicKey } from 'node:crypto';
+
+import { Test } from '@nestjs/testing';
+import ava, { TestFn } from 'ava';
+import Sinon from 'sinon';
+
+import { ConfigModule } from '../../config';
+import { CryptoHelper } from '../crypto';
+
+const test = ava as TestFn<{
+  crypto: CryptoHelper;
+}>;
+
+const key = `-----BEGIN EC PRIVATE KEY-----
+MHcCAQEEIEtyAJLIULkphVhqXqxk4Nr8Ggty3XLwUJWBxzAWCWTMoAoGCCqGSM49
+AwEHoUQDQgAEF3U/0wIeJ3jRKXeFKqQyBKlr9F7xaAUScRrAuSP33rajm3cdfihI
+3JvMxVNsS2lE8PSGQrvDrJZaDo0L+Lq9Gg==
+-----END EC PRIVATE KEY-----`;
+const privateKey = createPrivateKey({
+  key,
+  format: 'pem',
+  type: 'sec1',
+})
+  .export({
+    type: 'pkcs8',
+    format: 'pem',
+  })
+  .toString('utf8');
+
+const publicKey = createPublicKey({
+  key,
+  format: 'pem',
+  type: 'spki',
+})
+  .export({
+    format: 'pem',
+    type: 'spki',
+  })
+  .toString('utf8');
+
+test.beforeEach(async t => {
+  const module = await Test.createTestingModule({
+    imports: [
+      ConfigModule.forRoot({
+        secrets: {
+          publicKey,
+          privateKey,
+        },
+      }),
+    ],
+    providers: [CryptoHelper],
+  }).compile();
+
+  t.context.crypto = module.get(CryptoHelper);
+});
+
+test('should be able to sign and verify', t => {
+  const data = 'hello world';
+  const signature = t.context.crypto.sign(data);
+  t.true(t.context.crypto.verify(data, signature));
+  t.false(t.context.crypto.verify(data, 'fake-signature'));
+});
+
+test('should be able to encrypt and decrypt', t => {
+  const data = 'top secret';
+  const stub = Sinon.stub(t.context.crypto, 'randomBytes').returns(
+    Buffer.alloc(12, 0)
+  );
+
+  const encrypted = t.context.crypto.encrypt(data);
+  const decrypted = t.context.crypto.decrypt(encrypted);
+
+  // we are using a stub to make sure the iv is always 0,
+  // the encrypted result will always be the same
+  t.is(encrypted, 'AAAAAAAAAAAAAAAAWUDlJRhzP+SZ3avvmLcgnou+q4E11w==');
+  t.is(decrypted, data);
+
+  stub.restore();
+});
+
+test('should be able to get random bytes', t => {
+  const bytes = t.context.crypto.randomBytes();
+  t.is(bytes.length, 12);
+  const bytes2 = t.context.crypto.randomBytes();
+
+  t.notDeepEqual(bytes, bytes2);
+});
+
+test('should be able to digest', t => {
+  const data = 'hello world';
+  const hash = t.context.crypto.sha256(data).toString('base64');
+  t.is(hash, 'uU0nuZNNPgilLlLX2n2r+sSE7+N6U4DukIj3rOLvzek=');
+});
+
+test('should be able to safe compare', t => {
+  t.true(t.context.crypto.compare('abc', 'abc'));
+  t.false(t.context.crypto.compare('abc', 'def'));
+});
+
+test('should be able to hash and verify password', async t => {
+  const password = 'mySecurePassword';
+  const hash = await t.context.crypto.encryptPassword(password);
+  t.true(await t.context.crypto.verifyPassword(password, hash));
+  t.false(await t.context.crypto.verifyPassword('wrong-password', hash));
+});