Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(server): add administrator feature #6995

Merged
merged 1 commit into from
May 27, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 6 additions & 4 deletions packages/backend/server/src/core/auth/guard.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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);
1 change: 1 addition & 0 deletions packages/backend/server/src/core/auth/service.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
52 changes: 52 additions & 0 deletions packages/backend/server/src/core/common/admin-guard.ts
Original file line number Diff line number Diff line change
@@ -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);
};
1 change: 1 addition & 0 deletions packages/backend/server/src/core/common/index.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
export * from './admin-guard';
89 changes: 9 additions & 80 deletions packages/backend/server/src/core/features/feature.ts
Original file line number Diff line number Diff line change
@@ -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}`);
Expand All @@ -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({
Expand All @@ -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);
Expand Down
9 changes: 8 additions & 1 deletion packages/backend/server/src/core/features/index.ts
Original file line number Diff line number Diff line change
@@ -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';

/**
Expand All @@ -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 {}
Expand Down
54 changes: 26 additions & 28 deletions packages/backend/server/src/core/features/management.ts
Original file line number Diff line number Diff line change
@@ -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',
Expand All @@ -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,
Expand Down Expand Up @@ -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
Expand All @@ -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;
}
Expand Down
Loading
Loading