feat(server): add administrator feature (#6995)

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

@@ -22,6 +22,8 @@ function extractTokenFromHeader(authorization: string) {
   return authorization.substring(7);
 }
 
+const PUBLIC_ENTRYPOINT_SYMBOL = Symbol('public');
+
 @Injectable()
 export class AuthGuard implements CanActivate, OnModuleInit {
   private auth!: AuthService;
@@ -72,9 +74,9 @@ export class AuthGuard implements CanActivate, OnModuleInit {
     }
 
     // api is public
-    const isPublic = this.reflector.get<boolean>(
-      'isPublic',
-      context.getHandler()
+    const isPublic = this.reflector.getAllAndOverride<boolean>(
+      PUBLIC_ENTRYPOINT_SYMBOL,
+      [context.getClass(), context.getHandler()]
     );
 
     if (isPublic) {
@@ -110,4 +112,4 @@ export const Auth = () => {
 };
 
 // api is public accessible
-export const Public = () => SetMetadata('isPublic', true);
+export const Public = () => SetMetadata(PUBLIC_ENTRYPOINT_SYMBOL, true);

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

@@ -89,6 +89,7 @@ export class AuthService implements OnApplicationBootstrap {
           });
         }
         await this.quota.switchUserQuota(devUser.id, QuotaType.ProPlanV1);
+        await this.feature.addAdmin(devUser.id);
         await this.feature.addCopilot(devUser.id);
       } catch (e) {
         // ignore

packages/backend/server/src/core/common/admin-guard.ts

@@ -0,0 +1,52 @@
+import type {
+  CanActivate,
+  ExecutionContext,
+  OnModuleInit,
+} from '@nestjs/common';
+import { Injectable, UnauthorizedException, UseGuards } from '@nestjs/common';
+import { ModuleRef } from '@nestjs/core';
+
+import { getRequestResponseFromContext } from '../../fundamentals';
+import { FeatureManagementService } from '../features';
+
+@Injectable()
+export class AdminGuard implements CanActivate, OnModuleInit {
+  private feature!: FeatureManagementService;
+
+  constructor(private readonly ref: ModuleRef) {}
+
+  onModuleInit() {
+    this.feature = this.ref.get(FeatureManagementService, { strict: false });
+  }
+
+  async canActivate(context: ExecutionContext) {
+    const { req } = getRequestResponseFromContext(context);
+    let allow = false;
+    if (req.user) {
+      allow = await this.feature.isAdmin(req.user.id);
+    }
+
+    if (!allow) {
+      throw new UnauthorizedException('Your operation is not allowed.');
+    }
+
+    return true;
+  }
+}
+
+/**
+ * This guard is used to protect routes/queries/mutations that require a user to be administrator.
+ *
+ * @example
+ *
+ * ```typescript
+ * \@Admin()
+ * \@Mutation(() => UserType)
+ * createAccount(userInput: UserInput) {
+ *   // ...
+ * }
+ * ```
+ */
+export const Admin = () => {
+  return UseGuards(AdminGuard);
+};

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

@@ -0,0 +1 @@
+export * from './admin-guard';

packages/backend/server/src/core/features/feature.ts

@@ -1,12 +1,14 @@
 import { PrismaTransaction } from '../../fundamentals';
 import { Feature, FeatureSchema, FeatureType } from './types';
 
-class FeatureConfig {
-  readonly config: Feature;
+class FeatureConfig<T extends FeatureType> {
+  readonly config: Feature & { feature: T };
 
   constructor(data: any) {
     const config = FeatureSchema.safeParse(data);
+
     if (config.success) {
+      // @ts-expect-error allow
       this.config = config.data;
     } else {
       throw new Error(`Invalid quota config: ${config.error.message}`);
@@ -19,83 +21,15 @@ class FeatureConfig {
   }
 }
 
-export class CopilotFeatureConfig extends FeatureConfig {
-  override config!: Feature & { feature: FeatureType.Copilot };
-  constructor(data: any) {
-    super(data);
-
-    if (this.config.feature !== FeatureType.Copilot) {
-      throw new Error('Invalid feature config: type is not Copilot');
-    }
-  }
-}
-
-export class EarlyAccessFeatureConfig extends FeatureConfig {
-  override config!: Feature & { feature: FeatureType.EarlyAccess };
-
-  constructor(data: any) {
-    super(data);
-
-    if (this.config.feature !== FeatureType.EarlyAccess) {
-      throw new Error('Invalid feature config: type is not EarlyAccess');
-    }
-  }
-}
-
-export class UnlimitedWorkspaceFeatureConfig extends FeatureConfig {
-  override config!: Feature & { feature: FeatureType.UnlimitedWorkspace };
-
-  constructor(data: any) {
-    super(data);
-
-    if (this.config.feature !== FeatureType.UnlimitedWorkspace) {
-      throw new Error('Invalid feature config: type is not UnlimitedWorkspace');
-    }
-  }
-}
-
-export class UnlimitedCopilotFeatureConfig extends FeatureConfig {
-  override config!: Feature & { feature: FeatureType.UnlimitedCopilot };
-
-  constructor(data: any) {
-    super(data);
-
-    if (this.config.feature !== FeatureType.UnlimitedCopilot) {
-      throw new Error('Invalid feature config: type is not AIEarlyAccess');
-    }
-  }
-}
-export class AIEarlyAccessFeatureConfig extends FeatureConfig {
-  override config!: Feature & { feature: FeatureType.AIEarlyAccess };
-
-  constructor(data: any) {
-    super(data);
-
-    if (this.config.feature !== FeatureType.AIEarlyAccess) {
-      throw new Error('Invalid feature config: type is not AIEarlyAccess');
-    }
-  }
-}
-
-const FeatureConfigMap = {
-  [FeatureType.Copilot]: CopilotFeatureConfig,
-  [FeatureType.EarlyAccess]: EarlyAccessFeatureConfig,
-  [FeatureType.AIEarlyAccess]: AIEarlyAccessFeatureConfig,
-  [FeatureType.UnlimitedWorkspace]: UnlimitedWorkspaceFeatureConfig,
-  [FeatureType.UnlimitedCopilot]: UnlimitedCopilotFeatureConfig,
-};
-
-export type FeatureConfigType<F extends FeatureType> = InstanceType<
-  (typeof FeatureConfigMap)[F]
->;
+export type FeatureConfigType<F extends FeatureType> = FeatureConfig<F>;
 
 const FeatureCache = new Map<number, FeatureConfigType<FeatureType>>();
 
 export async function getFeature(prisma: PrismaTransaction, featureId: number) {
-  const cachedQuota = FeatureCache.get(featureId);
+  const cachedFeature = FeatureCache.get(featureId);
 
-  if (cachedQuota) {
-    return cachedQuota;
+  if (cachedFeature) {
+    return cachedFeature;
   }
 
   const feature = await prisma.features.findFirst({
@@ -107,13 +41,8 @@ export async function getFeature(prisma: PrismaTransaction, featureId: number) {
     // this should unreachable
     throw new Error(`Quota config ${featureId} not found`);
   }
-  const ConfigClass = FeatureConfigMap[feature.feature as FeatureType];
-
-  if (!ConfigClass) {
-    throw new Error(`Feature config ${featureId} not found`);
-  }
 
-  const config = new ConfigClass(feature);
+  const config = new FeatureConfig(feature);
   // we always edit quota config as a new quota config
   // so we can cache it by featureId
   FeatureCache.set(featureId, config);

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

@@ -1,6 +1,8 @@
 import { Module } from '@nestjs/common';
 
+import { UserModule } from '../user';
 import { EarlyAccessType, FeatureManagementService } from './management';
+import { FeatureManagementResolver } from './resolver';
 import { FeatureService } from './service';
 
 /**
@@ -10,7 +12,12 @@ import { FeatureService } from './service';
  * - feature statistics
  */
 @Module({
-  providers: [FeatureService, FeatureManagementService],
+  imports: [UserModule],
+  providers: [
+    FeatureService,
+    FeatureManagementService,
+    FeatureManagementResolver,
+  ],
   exports: [FeatureService, FeatureManagementService],
 })
 export class FeatureModule {}

packages/backend/server/src/core/features/management.ts

@@ -1,11 +1,11 @@
 import { Injectable, Logger } from '@nestjs/common';
-import { PrismaClient } from '@prisma/client';
 
 import { Config } from '../../fundamentals';
+import { UserService } from '../user/service';
 import { FeatureService } from './service';
 import { FeatureType } from './types';
 
-const STAFF = ['@toeverything.info'];
+const STAFF = ['@toeverything.info', '@affine.pro'];
 
 export enum EarlyAccessType {
   App = 'app',
@@ -18,22 +18,30 @@ export class FeatureManagementService {
 
   constructor(
     private readonly feature: FeatureService,
-    private readonly prisma: PrismaClient,
+    private readonly user: UserService,
     private readonly config: Config
   ) {}
 
   // ======== Admin ========
 
-  // todo(@darkskygit): replace this with abac
   isStaff(email: string) {
     for (const domain of STAFF) {
       if (email.endsWith(domain)) {
         return true;
       }
     }
+
     return false;
   }
 
+  isAdmin(userId: string) {
+    return this.feature.hasUserFeature(userId, FeatureType.Admin);
+  }
+
+  addAdmin(userId: string) {
+    return this.feature.addUserFeature(userId, FeatureType.Admin, 'Admin user');
+  }
+
   // ======== Early Access ========
   async addEarlyAccess(
     userId: string,
@@ -69,31 +77,17 @@ export class FeatureManagementService {
   }
 
   async isEarlyAccessUser(
-    email: string,
+    userId: string,
     type: EarlyAccessType = EarlyAccessType.App
   ) {
-    const user = await this.prisma.user.findFirst({
-      where: {
-        email: {
-          equals: email,
-          mode: 'insensitive',
-        },
-      },
-    });
-
-    if (user) {
-      const canEarlyAccess = await this.feature
-        .hasUserFeature(
-          user.id,
-          type === EarlyAccessType.App
-            ? FeatureType.EarlyAccess
-            : FeatureType.AIEarlyAccess
-        )
-        .catch(() => false);
-
-      return canEarlyAccess;
-    }
-    return false;
+    return await this.feature
+      .hasUserFeature(
+        userId,
+        type === EarlyAccessType.App
+          ? FeatureType.EarlyAccess
+          : FeatureType.AIEarlyAccess
+      )
+      .catch(() => false);
   }
 
   /// check early access by email
@@ -102,7 +96,11 @@ export class FeatureManagementService {
     type: EarlyAccessType = EarlyAccessType.App
   ) {
     if (this.config.featureFlags.earlyAccessPreview && !this.isStaff(email)) {
-      return this.isEarlyAccessUser(email, type);
+      const user = await this.user.findUserByEmail(email);
+      if (!user) {
+        return false;
+      }
+      return this.isEarlyAccessUser(user.id, type);
     } else {
       return true;
     }