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>(
+    const isPublic = this.reflector.getAllAndOverride<boolean>(
-      'isPublic',
+      PUBLIC_ENTRYPOINT_SYMBOL,
-      context.getHandler()
+      [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 {
+class FeatureConfig<T extends FeatureType> {
-  readonly config: Feature;
+  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 {
+export type FeatureConfigType<F extends FeatureType> = FeatureConfig<F>;
-  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]
->;
 
 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) {
+  if (cachedFeature) {
-    return cachedQuota;
+    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) {
+  const config = new FeatureConfig(feature);
-    throw new Error(`Feature config ${featureId} not found`);
-  }
-
-  const config = new ConfigClass(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({
+    return await this.feature
-      where: {
+      .hasUserFeature(
-        email: {
+        userId,
-          equals: email,
+        type === EarlyAccessType.App
-          mode: 'insensitive',
+          ? FeatureType.EarlyAccess
-        },
+          : FeatureType.AIEarlyAccess
-      },
+      )
-    });
+      .catch(() => false);
-
-    if (user) {
-      const canEarlyAccess = await this.feature
-        .hasUserFeature(
-          user.id,
-          type === EarlyAccessType.App
-            ? FeatureType.EarlyAccess
-            : FeatureType.AIEarlyAccess
-        )
-        .catch(() => false);
-
-      return canEarlyAccess;
-    }
-    return 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

@@ -1,4 +1,4 @@
-import { BadRequestException, ForbiddenException } from '@nestjs/common';
+import { BadRequestException } from '@nestjs/common';
 import {
   Args,
   Context,
@@ -6,35 +6,43 @@ import {
   Mutation,
   Query,
   registerEnumType,
+  ResolveField,
   Resolver,
 } from '@nestjs/graphql';
 
 import { CurrentUser } from '../auth/current-user';
 import { sessionUser } from '../auth/service';
-import { EarlyAccessType, FeatureManagementService } from '../features';
+import { Admin } from '../common';
-import { UserService } from './service';
+import { UserService } from '../user/service';
-import { UserType } from './types';
+import { UserType } from '../user/types';
+import { EarlyAccessType, FeatureManagementService } from './management';
+import { FeatureType } from './types';
 
 registerEnumType(EarlyAccessType, {
   name: 'EarlyAccessType',
 });
 
 @Resolver(() => UserType)
-export class UserManagementResolver {
+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(
-    @CurrentUser() currentUser: CurrentUser,
     @Args('email') email: string,
     @Args({ name: 'type', type: () => EarlyAccessType }) type: EarlyAccessType
   ): Promise<number> {
-    if (!this.feature.isStaff(currentUser.email)) {
-      throw new ForbiddenException('You are not allowed to do this');
-    }
     const user = await this.users.findUserByEmail(email);
     if (user) {
       return this.feature.addEarlyAccess(user.id, type);
@@ -46,14 +54,9 @@ export class UserManagementResolver {
     }
   }
 
+  @Admin()
   @Mutation(() => Int)
-  async removeEarlyAccess(
+  async removeEarlyAccess(@Args('email') email: string): Promise<number> {
-    @CurrentUser() currentUser: CurrentUser,
-    @Args('email') email: string
-  ): Promise<number> {
-    if (!this.feature.isStaff(currentUser.email)) {
-      throw new ForbiddenException('You are not allowed to do this');
-    }
     const user = await this.users.findUserByEmail(email);
     if (!user) {
       throw new BadRequestException(`User ${email} not found`);
@@ -61,18 +64,29 @@ export class UserManagementResolver {
     return this.feature.removeEarlyAccess(user.id);
   }
 
+  @Admin()
   @Query(() => [UserType])
   async earlyAccessUsers(
-    @Context() ctx: { isAdminQuery: boolean },
+    @Context() ctx: { isAdminQuery: boolean }
-    @CurrentUser() user: CurrentUser
   ): Promise<UserType[]> {
-    if (!this.feature.isStaff(user.email)) {
-      throw new ForbiddenException('You are not allowed to do this');
-    }
     // 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
+  async getFeature<F extends FeatureType>(feature: F) {
-  ): Promise<FeatureConfigType<F> | undefined> {
     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,
     ])
   );
 

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

@@ -0,0 +1,68 @@
+import {
+  Field,
+  ObjectType,
+  registerEnumType,
+  ResolveField,
+  Resolver,
+} from '@nestjs/graphql';
+import { SafeIntResolver } from 'graphql-scalars';
+
+import { CurrentUser } from '../auth/current-user';
+import { EarlyAccessType } from '../features';
+import { UserType } from '../user';
+import { QuotaService } from './service';
+
+registerEnumType(EarlyAccessType, {
+  name: 'EarlyAccessType',
+});
+
+@ObjectType('UserQuotaHumanReadable')
+class UserQuotaHumanReadableType {
+  @Field({ name: 'name' })
+  name!: string;
+
+  @Field({ name: 'blobLimit' })
+  blobLimit!: string;
+
+  @Field({ name: 'storageQuota' })
+  storageQuota!: string;
+
+  @Field({ name: 'historyPeriod' })
+  historyPeriod!: string;
+
+  @Field({ name: 'memberLimit' })
+  memberLimit!: string;
+}
+
+@ObjectType('UserQuota')
+class UserQuotaType {
+  @Field({ name: 'name' })
+  name!: string;
+
+  @Field(() => SafeIntResolver, { name: 'blobLimit' })
+  blobLimit!: number;
+
+  @Field(() => SafeIntResolver, { name: 'storageQuota' })
+  storageQuota!: number;
+
+  @Field(() => SafeIntResolver, { name: 'historyPeriod' })
+  historyPeriod!: number;
+
+  @Field({ name: 'memberLimit' })
+  memberLimit!: number;
+
+  @Field({ name: 'humanReadable' })
+  humanReadable!: UserQuotaHumanReadableType;
+}
+
+@Resolver(() => UserType)
+export class FeatureManagementResolver {
+  constructor(private readonly quota: QuotaService) {}
+
+  @ResolveField(() => UserQuotaType, { name: 'quota', nullable: true })
+  async getQuota(@CurrentUser() me: UserType) {
+    const quota = await this.quota.getUserQuota(me.id);
+
+    return quota.feature;
+  }
+}

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

@@ -4,7 +4,8 @@ import { PrismaClient } from '@prisma/client';
 import type { EventPayload } from '../../fundamentals';
 import { OnEvent, PrismaTransaction } from '../../fundamentals';
 import { SubscriptionPlan } from '../../plugins/payment/types';
-import { FeatureKind, FeatureManagementService } from '../features';
+import { FeatureManagementService } from '../features/management';
+import { FeatureKind } from '../features/types';
 import { QuotaConfig } from './quota';
 import { QuotaType } from './types';
 

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

@@ -1,16 +1,13 @@
 import { Module } from '@nestjs/common';
 
-import { FeatureModule } from '../features';
-import { QuotaModule } from '../quota';
 import { StorageModule } from '../storage';
 import { UserAvatarController } from './controller';
-import { UserManagementResolver } from './management';
 import { UserResolver } from './resolver';
 import { UserService } from './service';
 
 @Module({
-  imports: [StorageModule, FeatureModule, QuotaModule],
+  imports: [StorageModule],
-  providers: [UserResolver, UserManagementResolver, UserService],
+  providers: [UserResolver, UserService],
   controllers: [UserAvatarController],
   exports: [UserService],
 })

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

@@ -7,30 +7,23 @@ import {
   ResolveField,
   Resolver,
 } from '@nestjs/graphql';
-import type { User } from '@prisma/client';
 import { PrismaClient } from '@prisma/client';
 import GraphQLUpload from 'graphql-upload/GraphQLUpload.mjs';
 import { isNil, omitBy } from 'lodash-es';
 
 import type { FileUpload } from '../../fundamentals';
-import {
+import { EventEmitter, Throttle } from '../../fundamentals';
-  EventEmitter,
-  PaymentRequiredException,
-  Throttle,
-} from '../../fundamentals';
 import { CurrentUser } from '../auth/current-user';
 import { Public } from '../auth/guard';
 import { sessionUser } from '../auth/service';
-import { FeatureManagementService, FeatureType } from '../features';
-import { QuotaService } from '../quota';
 import { AvatarStorage } from '../storage';
+import { validators } from '../utils/validators';
 import { UserService } from './service';
 import {
   DeleteAccount,
   RemoveAvatar,
   UpdateUserInput,
   UserOrLimitedUser,
-  UserQuotaType,
   UserType,
 } from './types';
 
@@ -40,8 +33,6 @@ export class UserResolver {
     private readonly prisma: PrismaClient,
     private readonly storage: AvatarStorage,
     private readonly users: UserService,
-    private readonly feature: FeatureManagementService,
-    private readonly quota: QuotaService,
     private readonly event: EventEmitter
   ) {}
 
@@ -53,14 +44,10 @@ export class UserResolver {
   })
   @Public()
   async user(
-    @CurrentUser() currentUser?: CurrentUser,
+    @Args('email') email: string,
-    @Args('email') email?: string
+    @CurrentUser() currentUser?: CurrentUser
   ): Promise<typeof UserOrLimitedUser | null> {
-    if (!email || !(await this.feature.canEarlyAccess(email))) {
+    validators.assertValidEmail(email);
-      throw new PaymentRequiredException(
-        `You don't have early access permission\nVisit https://community.affine.pro/c/insider-general/ for more information`
-      );
-    }
 
     // TODO: need to limit a user can only get another user witch is in the same workspace
     const user = await this.users.findUserWithHashedPasswordByEmail(email);
@@ -79,13 +66,6 @@ export class UserResolver {
     };
   }
 
-  @ResolveField(() => UserQuotaType, { name: 'quota', nullable: true })
-  async getQuota(@CurrentUser() me: User) {
-    const quota = await this.quota.getUserQuota(me.id);
-
-    return quota.feature;
-  }
-
   @ResolveField(() => Int, {
     name: 'invoiceCount',
     description: 'Get user invoice count',
@@ -96,14 +76,6 @@ export class UserResolver {
     });
   }
 
-  @ResolveField(() => [FeatureType], {
-    name: 'features',
-    description: 'Enabled features of a user',
-  })
-  async userFeatures(@CurrentUser() user: CurrentUser) {
-    return this.feature.getActivatedUserFeatures(user.id);
-  }
-
   @Mutation(() => UserType, {
     name: 'uploadAvatar',
     description: 'Upload user avatar',

packages/backend/server/src/core/user/types.ts

@@ -6,49 +6,9 @@ import {
   ObjectType,
 } from '@nestjs/graphql';
 import type { User } from '@prisma/client';
-import { SafeIntResolver } from 'graphql-scalars';
 
 import { CurrentUser } from '../auth/current-user';
 
-@ObjectType('UserQuotaHumanReadable')
-export class UserQuotaHumanReadableType {
-  @Field({ name: 'name' })
-  name!: string;
-
-  @Field({ name: 'blobLimit' })
-  blobLimit!: string;
-
-  @Field({ name: 'storageQuota' })
-  storageQuota!: string;
-
-  @Field({ name: 'historyPeriod' })
-  historyPeriod!: string;
-
-  @Field({ name: 'memberLimit' })
-  memberLimit!: string;
-}
-
-@ObjectType('UserQuota')
-export class UserQuotaType {
-  @Field({ name: 'name' })
-  name!: string;
-
-  @Field(() => SafeIntResolver, { name: 'blobLimit' })
-  blobLimit!: number;
-
-  @Field(() => SafeIntResolver, { name: 'storageQuota' })
-  storageQuota!: number;
-
-  @Field(() => SafeIntResolver, { name: 'historyPeriod' })
-  historyPeriod!: number;
-
-  @Field({ name: 'memberLimit' })
-  memberLimit!: number;
-
-  @Field({ name: 'humanReadable' })
-  humanReadable!: UserQuotaHumanReadableType;
-}
-
 @ObjectType()
 export class UserType implements CurrentUser {
   @Field(() => ID)

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

@@ -10,6 +10,7 @@ import {
 } from '@nestjs/graphql';
 
 import { CurrentUser } from '../auth';
+import { Admin } from '../common';
 import { FeatureManagementService, FeatureType } from '../features';
 import { PermissionService } from './permission';
 import { WorkspaceType } from './types';
@@ -21,41 +22,29 @@ export class WorkspaceManagementResolver {
     private readonly permission: PermissionService
   ) {}
 
+  @Admin()
   @Mutation(() => Int)
   async addWorkspaceFeature(
-    @CurrentUser() currentUser: CurrentUser,
     @Args('workspaceId') workspaceId: string,
     @Args('feature', { type: () => FeatureType }) feature: FeatureType
   ): Promise<number> {
-    if (!this.feature.isStaff(currentUser.email)) {
-      throw new ForbiddenException('You are not allowed to do this');
-    }
-
     return this.feature.addWorkspaceFeatures(workspaceId, feature);
   }
 
+  @Admin()
   @Mutation(() => Int)
   async removeWorkspaceFeature(
-    @CurrentUser() currentUser: CurrentUser,
     @Args('workspaceId') workspaceId: string,
     @Args('feature', { type: () => FeatureType }) feature: FeatureType
   ): Promise<boolean> {
-    if (!this.feature.isStaff(currentUser.email)) {
-      throw new ForbiddenException('You are not allowed to do this');
-    }
-
     return this.feature.removeWorkspaceFeature(workspaceId, feature);
   }
 
+  @Admin()
   @Query(() => [WorkspaceType])
   async listWorkspaceFeatures(
-    @CurrentUser() user: CurrentUser,
     @Args('feature', { type: () => FeatureType }) feature: FeatureType
   ): Promise<WorkspaceType[]> {
-    if (!this.feature.isStaff(user.email)) {
-      throw new ForbiddenException('You are not allowed to do this');
-    }
-
     return this.feature.listFeatureWorkspaces(feature);
   }
 

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

@@ -1,16 +1,18 @@
 import { ModuleRef } from '@nestjs/core';
 import { PrismaClient } from '@prisma/client';
 
+import { FeatureManagementService } from '../../core/features';
 import { UserService } from '../../core/user';
 import { Config, CryptoHelper } from '../../fundamentals';
 
-export class SelfHostAdmin99999999 {
+export class SelfHostAdmin1 {
   // 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) {
+      const crypto = ref.get(CryptoHelper, { strict: false });
+      const user = ref.get(UserService, { strict: false });
+      const feature = ref.get(FeatureManagementService, { strict: false });
       if (
         !process.env.AFFINE_ADMIN_EMAIL ||
         !process.env.AFFINE_ADMIN_PASSWORD
@@ -19,6 +21,7 @@ export class SelfHostAdmin99999999 {
           'You have to set AFFINE_ADMIN_EMAIL and AFFINE_ADMIN_PASSWORD environment variables to generate the initial user for self-hosted AFFiNE Server.'
         );
       }
+
       await user.findOrCreateUser(process.env.AFFINE_ADMIN_EMAIL, {
         name: 'AFFINE First User',
         emailVerifiedAt: new Date(),
@@ -26,6 +29,15 @@ export class SelfHostAdmin99999999 {
           process.env.AFFINE_ADMIN_PASSWORD
         ),
       });
+
+      const firstUser = await db.user.findFirst({
+        orderBy: {
+          createdAt: 'asc',
+        },
+      });
+      if (firstUser) {
+        await feature.addAdmin(firstUser.id);
+      }
     }
   }
 

packages/backend/server/src/data/migrations/1716195522794-administrator-feature.ts

@@ -0,0 +1,14 @@
+import { PrismaClient } from '@prisma/client';
+
+import { FeatureType } from '../../core/features';
+import { upsertLatestFeatureVersion } from './utils/user-features';
+
+export class AdministratorFeature1716195522794 {
+  // do the migration
+  static async up(db: PrismaClient) {
+    await upsertLatestFeatureVersion(db, FeatureType.Admin);
+  }
+
+  // revert the migration
+  static async down(_db: PrismaClient) {}
+}

packages/backend/server/src/plugins/payment/service.ts

@@ -14,7 +14,7 @@ import Stripe from 'stripe';
 
 import { CurrentUser } from '../../core/auth';
 import { EarlyAccessType, FeatureManagementService } from '../../core/features';
-import { EventEmitter } from '../../fundamentals';
+import { Config, EventEmitter } from '../../fundamentals';
 import { ScheduleManager } from './schedule';
 import {
   InvoiceStatus,
@@ -66,6 +66,7 @@ export class SubscriptionService {
   private readonly logger = new Logger(SubscriptionService.name);
 
   constructor(
+    private readonly config: Config,
     private readonly stripe: Stripe,
     private readonly db: PrismaClient,
     private readonly scheduleManager: ScheduleManager,
@@ -78,10 +79,10 @@ export class SubscriptionService {
     let canHaveAIEarlyAccessDiscount = false;
     if (user) {
       canHaveEarlyAccessDiscount = await this.features.isEarlyAccessUser(
-        user.email
+        user.id
       );
       canHaveAIEarlyAccessDiscount = await this.features.isEarlyAccessUser(
-        user.email,
+        user.id,
         EarlyAccessType.AI
       );
 
@@ -154,6 +155,14 @@ export class SubscriptionService {
     redirectUrl: string;
     idempotencyKey: string;
   }) {
+    if (
+      this.config.deploy &&
+      this.config.affine.canary &&
+      !this.features.isStaff(user.email)
+    ) {
+      throw new BadRequestException('You are not allowed to do this.');
+    }
+
     const currentSubscription = await this.db.userSubscription.findFirst({
       where: {
         userId: user.id,
@@ -631,7 +640,7 @@ export class SubscriptionService {
   private async getOrCreateCustomer(
     idempotencyKey: string,
     user: CurrentUser
-  ): Promise<UserStripeCustomer & { email: string }> {
+  ): Promise<UserStripeCustomer> {
     let customer = await this.db.userStripeCustomer.findUnique({
       where: {
         userId: user.id,
@@ -662,10 +671,7 @@ export class SubscriptionService {
       });
     }
 
-    return {
+    return customer;
-      ...customer,
-      email: user.email,
-    };
   }
 
   private async retrieveUserFromCustomer(customerId: string) {
@@ -737,11 +743,11 @@ export class SubscriptionService {
    * Get available for different plans with special early-access price and coupon
    */
   private async getAvailablePrice(
-    customer: UserStripeCustomer & { email: string },
+    customer: UserStripeCustomer,
     plan: SubscriptionPlan,
     recurring: SubscriptionRecurring
   ): Promise<{ price: string; coupon?: string }> {
-    const isEaUser = await this.features.isEarlyAccessUser(customer.email);
+    const isEaUser = await this.features.isEarlyAccessUser(customer.userId);
     const oldSubscriptions = await this.stripe.subscriptions.list({
       customer: customer.stripeCustomerId,
       status: 'all',
@@ -771,7 +777,7 @@ export class SubscriptionService {
       };
     } else {
       const isAIEaUser = await this.features.isEarlyAccessUser(
-        customer.email,
+        customer.userId,
         EarlyAccessType.AI
       );
 

packages/backend/server/src/schema.gql

@@ -96,6 +96,7 @@ enum EarlyAccessType {
 """The type of workspace feature"""
 enum FeatureType {
   AIEarlyAccess
+  Admin
   Copilot
   EarlyAccess
   UnlimitedCopilot
@@ -184,6 +185,7 @@ type LimitedUserType {
 
 type Mutation {
   acceptInviteById(inviteId: String!, sendAcceptMail: Boolean, workspaceId: String!): Boolean!
+  addAdminister(email: String!): Boolean!
   addToEarlyAccess(email: String!, type: EarlyAccessType!): Int!
   addWorkspaceFeature(feature: FeatureType!, workspaceId: String!): Int!
   cancelSubscription(idempotencyKey: String!, plan: SubscriptionPlan = Pro): UserSubscription!
@@ -428,23 +430,6 @@ type UserInvoice {
 
 union UserOrLimitedUser = LimitedUserType | UserType
 
-type UserQuota {
-  blobLimit: SafeInt!
-  historyPeriod: SafeInt!
-  humanReadable: UserQuotaHumanReadable!
-  memberLimit: Int!
-  name: String!
-  storageQuota: SafeInt!
-}
-
-type UserQuotaHumanReadable {
-  blobLimit: String!
-  historyPeriod: String!
-  memberLimit: String!
-  name: String!
-  storageQuota: String!
-}
-
 type UserSubscription {
   canceledAt: DateTime
   createdAt: DateTime!
@@ -492,7 +477,6 @@ type UserType {
 
   """User name"""
   name: String!
-  quota: UserQuota
   subscription(plan: SubscriptionPlan = Pro): UserSubscription @deprecated(reason: "use `UserType.subscriptions`")
   subscriptions: [UserSubscription!]!
   token: tokenType! @deprecated(reason: "use [/api/auth/authorize]")

packages/backend/server/tests/payment/service.spec.ts

@@ -178,10 +178,8 @@ test('should list normal price for unauthenticated user', async t => {
 test('should list normal prices for authenticated user', async t => {
   const { feature, service, u1, stripe } = t.context;
 
-  feature.isEarlyAccessUser.withArgs(u1.email).resolves(false);
+  feature.isEarlyAccessUser.withArgs(u1.id).resolves(false);
-  feature.isEarlyAccessUser
+  feature.isEarlyAccessUser.withArgs(u1.id, EarlyAccessType.AI).resolves(false);
-    .withArgs(u1.email, EarlyAccessType.AI)
-    .resolves(false);
 
   // @ts-expect-error stub
   Sinon.stub(stripe.subscriptions, 'list').resolves({ data: [] });
@@ -200,10 +198,8 @@ test('should list normal prices for authenticated user', async t => {
 test('should list early access prices for pro ea user', async t => {
   const { feature, service, u1, stripe } = t.context;
 
-  feature.isEarlyAccessUser.withArgs(u1.email).resolves(true);
+  feature.isEarlyAccessUser.withArgs(u1.id).resolves(true);
-  feature.isEarlyAccessUser
+  feature.isEarlyAccessUser.withArgs(u1.id, EarlyAccessType.AI).resolves(false);
-    .withArgs(u1.email, EarlyAccessType.AI)
-    .resolves(false);
 
   // @ts-expect-error stub
   Sinon.stub(stripe.subscriptions, 'list').resolves({ data: [] });
@@ -222,10 +218,8 @@ test('should list early access prices for pro ea user', async t => {
 test('should list normal prices for pro ea user with old subscriptions', async t => {
   const { feature, service, u1, stripe } = t.context;
 
-  feature.isEarlyAccessUser.withArgs(u1.email).resolves(true);
+  feature.isEarlyAccessUser.withArgs(u1.id).resolves(true);
-  feature.isEarlyAccessUser
+  feature.isEarlyAccessUser.withArgs(u1.id, EarlyAccessType.AI).resolves(false);
-    .withArgs(u1.email, EarlyAccessType.AI)
-    .resolves(false);
 
   Sinon.stub(stripe.subscriptions, 'list').resolves({
     data: [
@@ -260,10 +254,8 @@ test('should list normal prices for pro ea user with old subscriptions', async t
 test('should list early access prices for ai ea user', async t => {
   const { feature, service, u1, stripe } = t.context;
 
-  feature.isEarlyAccessUser.withArgs(u1.email).resolves(false);
+  feature.isEarlyAccessUser.withArgs(u1.id).resolves(false);
-  feature.isEarlyAccessUser
+  feature.isEarlyAccessUser.withArgs(u1.id, EarlyAccessType.AI).resolves(true);
-    .withArgs(u1.email, EarlyAccessType.AI)
-    .resolves(true);
 
   // @ts-expect-error stub
   Sinon.stub(stripe.subscriptions, 'list').resolves({ data: [] });
@@ -282,10 +274,8 @@ test('should list early access prices for ai ea user', async t => {
 test('should list early access prices for pro and ai ea user', async t => {
   const { feature, service, u1, stripe } = t.context;
 
-  feature.isEarlyAccessUser.withArgs(u1.email).resolves(true);
+  feature.isEarlyAccessUser.withArgs(u1.id).resolves(true);
-  feature.isEarlyAccessUser
+  feature.isEarlyAccessUser.withArgs(u1.id, EarlyAccessType.AI).resolves(true);
-    .withArgs(u1.email, EarlyAccessType.AI)
-    .resolves(true);
 
   // @ts-expect-error stub
   Sinon.stub(stripe.subscriptions, 'list').resolves({ data: [] });
@@ -304,10 +294,8 @@ test('should list early access prices for pro and ai ea user', async t => {
 test('should list normal prices for ai ea user with old subscriptions', async t => {
   const { feature, service, u1, stripe } = t.context;
 
-  feature.isEarlyAccessUser.withArgs(u1.email).resolves(false);
+  feature.isEarlyAccessUser.withArgs(u1.id).resolves(false);
-  feature.isEarlyAccessUser
+  feature.isEarlyAccessUser.withArgs(u1.id, EarlyAccessType.AI).resolves(true);
-    .withArgs(u1.email, EarlyAccessType.AI)
-    .resolves(true);
 
   Sinon.stub(stripe.subscriptions, 'list').resolves({
     data: [
@@ -555,9 +543,9 @@ test('should get correct ai plan price for checking out', async t => {
 
   // pro ea user
   {
-    feature.isEarlyAccessUser.withArgs(u1.email).resolves(true);
+    feature.isEarlyAccessUser.withArgs(u1.id).resolves(true);
     feature.isEarlyAccessUser
-      .withArgs(u1.email, EarlyAccessType.AI)
+      .withArgs(u1.id, EarlyAccessType.AI)
       .resolves(false);
     // @ts-expect-error stub
     subListStub.resolves({ data: [] });
@@ -574,9 +562,9 @@ test('should get correct ai plan price for checking out', async t => {
 
   // pro ea user, but has old subscription
   {
-    feature.isEarlyAccessUser.withArgs(u1.email).resolves(true);
+    feature.isEarlyAccessUser.withArgs(u1.id).resolves(true);
     feature.isEarlyAccessUser
-      .withArgs(u1.email, EarlyAccessType.AI)
+      .withArgs(u1.id, EarlyAccessType.AI)
       .resolves(false);
     subListStub.resolves({
       data: [