Merge pull request #222 from import-ai/fix/mem

refactor: replace KVStore with CacheService

Author: LucienShui

src/api-key/api-key.service.spec.ts

@@ -2,7 +2,6 @@
 import { Test, TestingModule } from '@nestjs/testing';
 import { getRepositoryToken } from '@nestjs/typeorm';
 import { Repository } from 'typeorm';
-import { ForbiddenException } from '@nestjs/common';
 import { APIKeyService } from './api-key.service';
 import {
   APIKey,
@@ -74,8 +73,10 @@ describe('APIKeyService', () => {
       t: jest.fn((key: string, options?: any) => {
         // Return mock translations for test purposes
         const translations: Record<string, string> = {
-          'apikey.errors.noPermissionForNamespace': 'User {{userId}} does not have permission to namespace {{namespaceId}}',
-          'apikey.errors.noWritePermission': 'User {{userId}} does not have write permission to resource {{resourceId}} in namespace {{namespaceId}}',
+          'apikey.errors.noPermissionForNamespace':
+            'User {{userId}} does not have permission to namespace {{namespaceId}}',
+          'apikey.errors.noWritePermission':
+            'User {{userId}} does not have write permission to resource {{resourceId}} in namespace {{namespaceId}}',
         };
         let translation = translations[key] || key;
         if (options?.args) {

src/auth/api-key/api-key-auth.guard.spec.ts

@@ -1,9 +1,5 @@
 import { Test, TestingModule } from '@nestjs/testing';
-import {
-  ExecutionContext,
-  UnauthorizedException,
-  ForbiddenException,
-} from '@nestjs/common';
+import { ExecutionContext } from '@nestjs/common';
 import { Reflector } from '@nestjs/core';
 import { APIKeyAuthGuard } from 'omniboxd/auth';
 import { APIKeyService } from 'omniboxd/api-key/api-key.service';
@@ -20,7 +16,6 @@ describe('APIKeyAuthGuard', () => {
   let guard: APIKeyAuthGuard;
   let apiKeyService: jest.Mocked<APIKeyService>;
   let reflector: jest.Mocked<Reflector>;
-  let i18nService: jest.Mocked<I18nService>;
 
   beforeEach(async () => {
     const module: TestingModule = await Test.createTestingModule({
@@ -44,11 +39,14 @@ describe('APIKeyAuthGuard', () => {
             t: jest.fn((key: string) => {
               // Return mock translations for test purposes
               const translations: Record<string, string> = {
-                'apikey.errors.authorizationHeaderRequired': 'Authorization header is required',
+                'apikey.errors.authorizationHeaderRequired':
+                  'Authorization header is required',
                 'apikey.errors.invalidApiKeyFormat': 'Invalid API key format',
                 'apikey.errors.invalidApiKey': 'Invalid API key',
-                'apikey.errors.noPermissionForTarget': 'No permission for target {{target}}',
-                'apikey.errors.noSpecificPermission': 'No {{permission}} permission for target {{target}}',
+                'apikey.errors.noPermissionForTarget':
+                  'No permission for target {{target}}',
+                'apikey.errors.noSpecificPermission':
+                  'No {{permission}} permission for target {{target}}',
               };
               return translations[key] || key;
             }),
@@ -60,7 +58,6 @@ describe('APIKeyAuthGuard', () => {
     guard = module.get<APIKeyAuthGuard>(APIKeyAuthGuard);
     apiKeyService = module.get(APIKeyService);
     reflector = module.get(Reflector);
-    i18nService = module.get(I18nService);
   });
 
   const createMockExecutionContext = (
@@ -111,9 +108,7 @@ describe('APIKeyAuthGuard', () => {
       .mockReturnValueOnce({ enabled: true }); // apiKeyAuthOptions = { enabled: true }
     const context = createMockExecutionContext();
 
-    await expect(guard.canActivate(context)).rejects.toThrow(
-      AppException,
-    );
+    await expect(guard.canActivate(context)).rejects.toThrow(AppException);
   });
 
   it('should throw AppException when API key does not start with sk-', async () => {
@@ -122,9 +117,7 @@ describe('APIKeyAuthGuard', () => {
       .mockReturnValueOnce({ enabled: true }); // apiKeyAuthOptions = { enabled: true }
     const context = createMockExecutionContext('Bearer invalid-key');
 
-    await expect(guard.canActivate(context)).rejects.toThrow(
-      AppException,
-    );
+    await expect(guard.canActivate(context)).rejects.toThrow(AppException);
   });
 
   it('should throw AppException when API key is not found', async () => {
@@ -134,9 +127,7 @@ describe('APIKeyAuthGuard', () => {
     const context = createMockExecutionContext('Bearer sk-validformat');
     apiKeyService.findByValue.mockResolvedValue(null);
 
-    await expect(guard.canActivate(context)).rejects.toThrow(
-      AppException,
-    );
+    await expect(guard.canActivate(context)).rejects.toThrow(AppException);
   });
 
   it('should allow valid API key and set apiKey and user on request', async () => {
@@ -245,9 +236,7 @@ describe('APIKeyAuthGuard', () => {
     const context = createMockExecutionContext('Bearer sk-validkey');
     apiKeyService.findByValue.mockResolvedValue(mockApiKey);
 
-    await expect(guard.canActivate(context)).rejects.toThrow(
-      AppException,
-    );
+    await expect(guard.canActivate(context)).rejects.toThrow(AppException);
   });
 
   it('should throw AppException when API key lacks specific permission type', async () => {
@@ -283,9 +272,7 @@ describe('APIKeyAuthGuard', () => {
     const context = createMockExecutionContext('Bearer sk-validkey');
     apiKeyService.findByValue.mockResolvedValue(mockApiKey);
 
-    await expect(guard.canActivate(context)).rejects.toThrow(
-      AppException,
-    );
+    await expect(guard.canActivate(context)).rejects.toThrow(AppException);
   });
 
   it('should allow API key with multiple matching permissions', async () => {

src/auth/api-key/api-key-auth.guard.ts

@@ -77,7 +77,7 @@ export class APIKeyAuthGuard implements CanActivate {
       apiKeyAuthOptions.permissions &&
       apiKeyAuthOptions.permissions.length > 0
     ) {
-      await this.validatePermissions(
+      this.validatePermissions(
         apiKey.attrs.permissions || [],
         apiKeyAuthOptions.permissions,
       );
@@ -90,10 +90,10 @@ export class APIKeyAuthGuard implements CanActivate {
     return true;
   }
 
-  private async validatePermissions(
+  private validatePermissions(
     apiKeyPermissions: APIKeyPermission[],
     requiredPermissions: APIKeyPermission[],
-  ): Promise<void> {
+  ): void {
     for (const required of requiredPermissions) {
       const apiKeyPermission = apiKeyPermissions.find(
         (p) => p.target === required.target,

src/auth/auth.module.ts

@@ -23,6 +23,7 @@ import { SocialService } from 'omniboxd/auth/social.service';
 import { APIKeyModule } from 'omniboxd/api-key/api-key.module';
 import { APIKeyAuthGuard } from 'omniboxd/auth/api-key/api-key-auth.guard';
 import { CookieAuthGuard } from 'omniboxd/auth/cookie/cookie-auth.guard';
+import { CacheService } from 'omniboxd/common/cache.service';
 
 @Module({
   exports: [AuthService, WechatService, GoogleService, SocialService],
@@ -39,6 +40,7 @@ import { CookieAuthGuard } from 'omniboxd/auth/cookie/cookie-auth.guard';
     GoogleService,
     JwtStrategy,
     LocalStrategy,
+    CacheService,
     {
       provide: APP_GUARD,
       useClass: JwtAuthGuard,

src/auth/auth.service.ts

@@ -288,7 +288,7 @@ export class AuthService {
     }
   }
 
-  async jwtVerify(token: string) {
+  jwtVerify(token: string) {
     if (!token) {
       const message = this.i18n.t('auth.errors.noToken');
       throw new AppException(message, 'NO_TOKEN', HttpStatus.UNAUTHORIZED);

src/auth/cookie/cookie-auth.guard.spec.ts

@@ -1,5 +1,5 @@
 import { Test, TestingModule } from '@nestjs/testing';
-import { ExecutionContext, UnauthorizedException } from '@nestjs/common';
+import { ExecutionContext } from '@nestjs/common';
 import { Reflector } from '@nestjs/core';
 import { CookieAuthGuard } from 'omniboxd/auth/cookie/cookie-auth.guard';
 import { AuthService } from 'omniboxd/auth/auth.service';
@@ -10,7 +10,6 @@ describe('CookieAuthGuard', () => {
   let guard: CookieAuthGuard;
   let reflector: jest.Mocked<Reflector>;
   let authService: jest.Mocked<AuthService>;
-  let i18nService: jest.Mocked<I18nService>;
 
   beforeEach(async () => {
     const module: TestingModule = await Test.createTestingModule({
@@ -34,7 +33,8 @@ describe('CookieAuthGuard', () => {
             t: jest.fn((key: string) => {
               // Return mock translations for test purposes
               const translations: Record<string, string> = {
-                'auth.errors.tokenCookieRequired': 'Authentication token cookie is required',
+                'auth.errors.tokenCookieRequired':
+                  'Authentication token cookie is required',
                 'auth.errors.tokenInvalid': 'Invalid or expired token',
                 'auth.errors.invalidTokenPayload': 'Invalid token payload',
               };
@@ -48,7 +48,6 @@ describe('CookieAuthGuard', () => {
     guard = module.get<CookieAuthGuard>(CookieAuthGuard);
     reflector = module.get(Reflector);
     authService = module.get(AuthService);
-    i18nService = module.get(I18nService);
   });
 
   const createMockExecutionContext = (cookies: any = {}): ExecutionContext => {
@@ -94,9 +93,7 @@ describe('CookieAuthGuard', () => {
       .mockReturnValueOnce({ enabled: true }); // cookieAuthOptions with default onAuthFail = 'reject'
     const context = createMockExecutionContext();
 
-    await expect(guard.canActivate(context)).rejects.toThrow(
-      AppException,
-    );
+    await expect(guard.canActivate(context)).rejects.toThrow(AppException);
   });
 
   it('should throw AppException when token is invalid', async () => {
@@ -106,9 +103,7 @@ describe('CookieAuthGuard', () => {
     const context = createMockExecutionContext({ token: 'invalid-token' });
     authService.jwtVerify.mockRejectedValue(new Error('Invalid token'));
 
-    await expect(guard.canActivate(context)).rejects.toThrow(
-      AppException,
-    );
+    await expect(guard.canActivate(context)).rejects.toThrow(AppException);
   });
 
   it('should throw AppException when token payload is invalid', async () => {
@@ -118,9 +113,7 @@ describe('CookieAuthGuard', () => {
     const context = createMockExecutionContext({ token: 'valid-token' });
     authService.jwtVerify.mockResolvedValue({}); // No sub field
 
-    await expect(guard.canActivate(context)).rejects.toThrow(
-      AppException,
-    );
+    await expect(guard.canActivate(context)).rejects.toThrow(AppException);
   });
 
   it('should successfully authenticate with valid token and set cookie auth data', async () => {

src/auth/social.service.ts

@@ -2,12 +2,10 @@ import { EntityManager } from 'typeorm';
 import generateId from 'omniboxd/utils/generate-id';
 import { UserService } from 'omniboxd/user/user.service';
 import { WechatCheckResponseDto } from 'omniboxd/auth/dto/wechat-login.dto';
-import { HttpStatus, Inject, Injectable } from '@nestjs/common';
+import { HttpStatus, Injectable } from '@nestjs/common';
 import { AppException } from 'omniboxd/common/exceptions/app.exception';
 import { I18nService } from 'nestjs-i18n';
-import { Cache, CACHE_MANAGER } from '@nestjs/cache-manager';
-import { KVStore } from 'omniboxd/common/kv-store';
-import { ConfigService } from '@nestjs/config';
+import { CacheService } from 'omniboxd/common/cache.service';
 
 export interface UserSocialState {
   type: string;
@@ -18,22 +16,15 @@ export interface UserSocialState {
 
 @Injectable()
 export class SocialService {
+  private readonly namespace = '/social/states';
   private readonly minUsernameLength = 2;
   private readonly maxUsernameLength = 32;
-  private readonly kvStore: KVStore<UserSocialState>;
 
   constructor(
     private readonly userService: UserService,
     private readonly i18n: I18nService,
-    @Inject(CACHE_MANAGER) private readonly cacheManager: Cache,
-    private readonly configService: ConfigService,
-  ) {
-    this.kvStore = new KVStore(
-      this.cacheManager,
-      '/social/states',
-      this.configService,
-    );
-  }
+    private readonly cacheService: CacheService,
+  ) {}
 
   /**
    * Generate a kv state, return a kv key
@@ -44,7 +35,8 @@ export class SocialService {
   async generateState(type: string, prefix: string = ''): Promise<string> {
     const key = `${prefix ? prefix + '_' : ''}${generateId()}`;
     const expiresIn = 5 * 60 * 1000; // Expires in 5 minutes
-    await this.kvStore.set(
+    await this.cacheService.set<UserSocialState>(
+      this.namespace,
       key,
       {
         type,
@@ -57,13 +49,18 @@ export class SocialService {
   }
 
   async getState(state: string) {
-    return await this.kvStore.get(state);
+    return await this.cacheService.get<UserSocialState>(this.namespace, state);
   }
 
   async updateState(state: string, data: UserSocialState) {
     const ttl = data.expiresIn - (Date.now() - data.createdAt);
     if (ttl > 0) {
-      await this.kvStore.set(state, data, ttl);
+      await this.cacheService.set<UserSocialState>(
+        this.namespace,
+        state,
+        data,
+        ttl,
+      );
     }
   }
 

src/common/cache.service.ts

@@ -0,0 +1,36 @@
+import { Inject, Injectable } from '@nestjs/common';
+import { Cache, CACHE_MANAGER } from '@nestjs/cache-manager';
+import { ConfigService } from '@nestjs/config';
+
+@Injectable()
+export class CacheService {
+  private readonly env: string;
+
+  constructor(
+    @Inject(CACHE_MANAGER) private readonly cacheManager: Cache,
+    private readonly configService: ConfigService,
+  ) {
+    this.env = this.configService.get<string>('ENV', 'unknown');
+  }
+
+  private getKey(namespace: string, key: string): string {
+    return `/${this.env}${namespace}/${key}`;
+  }
+
+  async get<T>(namespace: string, key: string): Promise<T | null> {
+    return await this.cacheManager.get<T>(this.getKey(namespace, key));
+  }
+
+  async set<T>(
+    namespace: string,
+    key: string,
+    value: T,
+    ttl?: number,
+  ): Promise<void> {
+    await this.cacheManager.set(this.getKey(namespace, key), value, ttl);
+  }
+
+  async delete(namespace: string, key: string): Promise<void> {
+    await this.cacheManager.del(this.getKey(namespace, key));
+  }
+}

src/common/kv-store.ts

@@ -18,7 +18,11 @@ export const UserId = createParamDecorator(
 
     if (!userId && !optional) {
       // Note: Decorators cannot easily inject I18nService, using static message
-      throw new AppException('Not authorized', 'NOT_AUTHORIZED', HttpStatus.UNAUTHORIZED);
+      throw new AppException(
+        'Not authorized',
+        'NOT_AUTHORIZED',
+        HttpStatus.UNAUTHORIZED,
+      );
     }
 
     return userId;