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

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;
     }

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

@@ -0,0 +1,92 @@
+import { BadRequestException } from '@nestjs/common';
+import {
+  Args,
+  Context,
+  Int,
+  Mutation,
+  Query,
+  registerEnumType,
+  ResolveField,
+  Resolver,
+} from '@nestjs/graphql';
+
+import { CurrentUser } from '../auth/current-user';
+import { sessionUser } from '../auth/service';
+import { Admin } from '../common';
+import { UserService } from '../user/service';
+import { UserType } from '../user/types';
+import { EarlyAccessType, FeatureManagementService } from './management';
+import { FeatureType } from './types';
+
+registerEnumType(EarlyAccessType, {
+  name: 'EarlyAccessType',
+});
+
+@Resolver(() => UserType)
+export class FeatureManagementResolver {
+  constructor(
+    private readonly users: UserService,
+    private readonly feature: FeatureManagementService
+  ) {}
+
+  @ResolveField(() => [FeatureType], {
+    name: 'features',
+    description: 'Enabled features of a user',
+  })
+  async userFeatures(@CurrentUser() user: CurrentUser) {
+    return this.feature.getActivatedUserFeatures(user.id);
+  }
+
+  @Admin()
+  @Mutation(() => Int)
+  async addToEarlyAccess(
+    @Args('email') email: string,
+    @Args({ name: 'type', type: () => EarlyAccessType }) type: EarlyAccessType
+  ): Promise<number> {
+    const user = await this.users.findUserByEmail(email);
+    if (user) {
+      return this.feature.addEarlyAccess(user.id, type);
+    } else {
+      const user = await this.users.createAnonymousUser(email, {
+        registered: false,
+      });
+      return this.feature.addEarlyAccess(user.id, type);
+    }
+  }
+
+  @Admin()
+  @Mutation(() => Int)
+  async removeEarlyAccess(@Args('email') email: string): Promise<number> {
+    const user = await this.users.findUserByEmail(email);
+    if (!user) {
+      throw new BadRequestException(`User ${email} not found`);
+    }
+    return this.feature.removeEarlyAccess(user.id);
+  }
+
+  @Admin()
+  @Query(() => [UserType])
+  async earlyAccessUsers(
+    @Context() ctx: { isAdminQuery: boolean }
+  ): Promise<UserType[]> {
+    // allow query other user's subscription
+    ctx.isAdminQuery = true;
+    return this.feature.listEarlyAccess().then(users => {
+      return users.map(sessionUser);
+    });
+  }
+
+  @Admin()
+  @Mutation(() => Boolean)
+  async addAdminister(@Args('email') email: string): Promise<boolean> {
+    const user = await this.users.findUserByEmail(email);
+
+    if (!user) {
+      throw new BadRequestException(`User ${email} not found`);
+    }
+
+    await this.feature.addAdmin(user.id);
+
+    return true;
+  }
+}

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

@@ -8,9 +8,8 @@ import { FeatureKind, FeatureType } from './types';
 @Injectable()
 export class FeatureService {
   constructor(private readonly prisma: PrismaClient) {}
-  async getFeature<F extends FeatureType>(
-    feature: F
-  ): Promise<FeatureConfigType<F> | undefined> {
+
+  async getFeature<F extends FeatureType>(feature: F) {
     const data = await this.prisma.features.findFirst({
       where: {
         feature,
@@ -21,8 +20,9 @@ export class FeatureService {
         version: 'desc',
       },
     });
+
     if (data) {
-      return getFeature(this.prisma, data.id) as FeatureConfigType<F>;
+      return getFeature(this.prisma, data.id) as Promise<FeatureConfigType<F>>;
     }
     return undefined;
   }

packages/backend/server/src/core/features/types/admin.ts

@@ -0,0 +1,8 @@
+import { z } from 'zod';
+
+import { FeatureType } from './common';
+
+export const featureAdministrator = z.object({
+  feature: z.literal(FeatureType.Admin),
+  configs: z.object({}),
+});

packages/backend/server/src/core/features/types/common.ts

@@ -2,6 +2,7 @@ import { registerEnumType } from '@nestjs/graphql';
 
 export enum FeatureType {
   // user feature
+  Admin = 'administrator',
   EarlyAccess = 'early_access',
   AIEarlyAccess = 'ai_early_access',
   UnlimitedCopilot = 'unlimited_copilot',

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

@@ -1,5 +1,6 @@
 import { z } from 'zod';
 
+import { featureAdministrator } from './admin';
 import { FeatureType } from './common';
 import { featureCopilot } from './copilot';
 import { featureAIEarlyAccess, featureEarlyAccess } from './early-access';
@@ -65,6 +66,12 @@ export const Features: Feature[] = [
     version: 1,
     configs: {},
   },
+  {
+    feature: FeatureType.Admin,
+    type: FeatureKind.Feature,
+    version: 1,
+    configs: {},
+  },
 ];
 
 /// ======== schema infer ========
@@ -80,6 +87,7 @@ export const FeatureSchema = commonFeatureSchema
       featureAIEarlyAccess,
       featureUnlimitedWorkspace,
       featureUnlimitedCopilot,
+      featureAdministrator,
     ])
   );