AFFiNE-Mirror/apps/server/src/modules/workspaces/resolver.ts

import type { Storage } from '@affine/storage';
import {
  ForbiddenException,
  Inject,
  InternalServerErrorException,
  Logger,
  NotFoundException,
  UseGuards,
} from '@nestjs/common';
import {
  Args,
  Field,
  Float,
  ID,
  InputType,
  Int,
  Mutation,
  ObjectType,
  OmitType,
  Parent,
  PartialType,
  PickType,
  Query,
  registerEnumType,
  ResolveField,
  Resolver,
} from '@nestjs/graphql';
import type { User, Workspace } from '@prisma/client';
import GraphQLUpload from 'graphql-upload/GraphQLUpload.mjs';
import { applyUpdate, Doc } from 'yjs';

import { PrismaService } from '../../prisma';
import { StorageProvide } from '../../storage';
import { CloudThrottlerGuard, Throttle } from '../../throttler';
import type { FileUpload } from '../../types';
import { Auth, CurrentUser, Public } from '../auth';
import { MailService } from '../auth/mailer';
import { AuthService } from '../auth/service';
import { UsersService } from '../users';
import { UserType } from '../users/resolver';
import { PermissionService } from './permission';
import { Permission } from './types';
import { defaultWorkspaceAvatar } from './utils';

registerEnumType(Permission, {
  name: 'Permission',
  description: 'User permission in workspace',
});

@ObjectType()
export class InviteUserType extends OmitType(
  PartialType(UserType),
  ['id'],
  ObjectType
) {
  @Field(() => ID)
  id!: string;

  @Field(() => Permission, { description: 'User permission in workspace' })
  permission!: Permission;

  @Field({ description: 'Invite id' })
  inviteId!: string;

  @Field({ description: 'User accepted' })
  accepted!: boolean;
}

@ObjectType()
export class WorkspaceType implements Partial<Workspace> {
  @Field(() => ID)
  id!: string;

  @Field({ description: 'is Public workspace' })
  public!: boolean;

  @Field({ description: 'Workspace created date' })
  createdAt!: Date;

  @Field(() => [InviteUserType], {
    description: 'Members of workspace',
  })
  members!: InviteUserType[];
}

@ObjectType()
export class InvitationWorkspaceType {
  @Field(() => ID)
  id!: string;

  @Field({ description: 'Workspace name' })
  name!: string;

  @Field(() => String, {
    // nullable: true,
    description: 'Base64 encoded avatar',
  })
  avatar!: string;
}

@ObjectType()
export class WorkspaceBlobSizes {
  @Field(() => Float)
  size!: number;
}

@ObjectType()
export class InvitationType {
  @Field({ description: 'Workspace information' })
  workspace!: InvitationWorkspaceType;
  @Field({ description: 'User information' })
  user!: UserType;
  @Field({ description: 'Invitee information' })
  invitee!: UserType;
}

@InputType()
export class UpdateWorkspaceInput extends PickType(
  PartialType(WorkspaceType),
  ['public'],
  InputType
) {
  @Field(() => ID)
  id!: string;
}

/**
 * Workspace resolver
 * Public apis rate limit: 10 req/m
 * Other rate limit: 120 req/m
 */
@UseGuards(CloudThrottlerGuard)
@Auth()
@Resolver(() => WorkspaceType)
export class WorkspaceResolver {
  private readonly logger = new Logger('WorkspaceResolver');

  constructor(
    private readonly auth: AuthService,
    private readonly mailer: MailService,
    private readonly prisma: PrismaService,
    private readonly permissions: PermissionService,
    private readonly users: UsersService,
    @Inject(StorageProvide) private readonly storage: Storage
  ) {}

  @ResolveField(() => Permission, {
    description: 'Permission of current signed in user in workspace',
    complexity: 2,
  })
  async permission(
    @CurrentUser() user: UserType,
    @Parent() workspace: WorkspaceType
  ) {
    // may applied in workspaces query
    if ('permission' in workspace) {
      return workspace.permission;
    }

    const permission = await this.permissions.get(workspace.id, user.id);

    if (!permission) {
      throw new ForbiddenException();
    }

    return permission;
  }

  @ResolveField(() => Int, {
    description: 'member count of workspace',
    complexity: 2,
  })
  memberCount(@Parent() workspace: WorkspaceType) {
    return this.prisma.userWorkspacePermission.count({
      where: {
        workspaceId: workspace.id,
        userId: {
          not: null,
        },
      },
    });
  }

  @ResolveField(() => [String], {
    description: 'Shared pages of workspace',
    complexity: 2,
  })
  async sharedPages(@Parent() workspace: WorkspaceType) {
    const data = await this.prisma.userWorkspacePermission.findMany({
      where: {
        workspaceId: workspace.id,
      },
    });

    return data.map(item => item.subPageId).filter(Boolean);
  }

  @ResolveField(() => UserType, {
    description: 'Owner of workspace',
    complexity: 2,
  })
  async owner(@Parent() workspace: WorkspaceType) {
    const data = await this.permissions.getWorkspaceOwner(workspace.id);

    return data.user;
  }

  @ResolveField(() => [InviteUserType], {
    description: 'Members of workspace',
    complexity: 2,
  })
  async members(
    @Parent() workspace: WorkspaceType,
    @Args('skip', { type: () => Int, nullable: true }) skip?: number,
    @Args('take', { type: () => Int, nullable: true }) take?: number
  ) {
    const data = await this.prisma.userWorkspacePermission.findMany({
      where: {
        workspaceId: workspace.id,
        userId: {
          not: null,
        },
      },
      skip,
      take: take || 8,
      orderBy: [
        {
          createdAt: 'asc',
        },
        {
          type: 'desc',
        },
      ],
      include: {
        user: true,
      },
    });

    return data
      .filter(({ user }) => !!user)
      .map(({ id, accepted, type, user }) => ({
        ...user,
        permission: type,
        inviteId: id,
        accepted,
      }));
  }

  @Query(() => Boolean, {
    description: 'Get is owner of workspace',
    complexity: 2,
  })
  async isOwner(
    @CurrentUser() user: UserType,
    @Args('workspaceId') workspaceId: string
  ) {
    const data = await this.permissions.tryGetWorkspaceOwner(workspaceId);

    return data?.user?.id === user.id;
  }

  @Query(() => [WorkspaceType], {
    description: 'Get all accessible workspaces for current user',
    complexity: 2,
  })
  async workspaces(@CurrentUser() user: User) {
    const data = await this.prisma.userWorkspacePermission.findMany({
      where: {
        userId: user.id,
        accepted: true,
      },
      include: {
        workspace: true,
      },
    });

    return data.map(({ workspace, type }) => {
      return {
        ...workspace,
        permission: type,
      };
    });
  }

  @Throttle(10, 30)
  @Public()
  @Query(() => WorkspaceType, {
    description: 'Get public workspace by id',
  })
  async publicWorkspace(@Args('id') id: string) {
    const workspace = await this.prisma.workspace.findUnique({
      where: { id },
    });

    if (workspace?.public) {
      return workspace;
    }

    throw new NotFoundException("Workspace doesn't exist");
  }

  @Query(() => WorkspaceType, {
    description: 'Get workspace by id',
  })
  async workspace(@CurrentUser() user: UserType, @Args('id') id: string) {
    await this.permissions.check(id, user.id);
    const workspace = await this.prisma.workspace.findUnique({ where: { id } });

    if (!workspace) {
      throw new NotFoundException("Workspace doesn't exist");
    }

    return workspace;
  }

  @Mutation(() => WorkspaceType, {
    description: 'Create a new workspace',
  })
  async createWorkspace(
    @CurrentUser() user: UserType,
    @Args({ name: 'init', type: () => GraphQLUpload })
    update: FileUpload
  ) {
    // convert stream to buffer
    const buffer = await new Promise<Buffer>((resolve, reject) => {
      const stream = update.createReadStream();
      const chunks: Uint8Array[] = [];
      stream.on('data', chunk => {
        chunks.push(chunk);
      });
      stream.on('error', reject);
      stream.on('end', () => {
        resolve(Buffer.concat(chunks));
      });
    });

    const workspace = await this.prisma.workspace.create({
      data: {
        public: false,
        users: {
          create: {
            type: Permission.Owner,
            user: {
              connect: {
                id: user.id,
              },
            },
            accepted: true,
          },
        },
      },
    });

    if (buffer.length) {
      await this.prisma.snapshot.create({
        data: {
          id: workspace.id,
          workspaceId: workspace.id,
          blob: buffer,
        },
      });
    }

    return workspace;
  }

  @Mutation(() => WorkspaceType, {
    description: 'Update workspace',
  })
  async updateWorkspace(
    @CurrentUser() user: UserType,
    @Args({ name: 'input', type: () => UpdateWorkspaceInput })
    { id, ...updates }: UpdateWorkspaceInput
  ) {
    await this.permissions.check(id, user.id, Permission.Admin);

    return this.prisma.workspace.update({
      where: {
        id,
      },
      data: updates,
    });
  }

  @Mutation(() => Boolean)
  async deleteWorkspace(@CurrentUser() user: UserType, @Args('id') id: string) {
    await this.permissions.check(id, user.id, Permission.Owner);

    await this.prisma.workspace.delete({
      where: {
        id,
      },
    });

    await this.prisma.$transaction([
      this.prisma.update.deleteMany({
        where: {
          workspaceId: id,
        },
      }),
      this.prisma.snapshot.deleteMany({
        where: {
          workspaceId: id,
        },
      }),
    ]);

    return true;
  }

  @Mutation(() => String)
  async invite(
    @CurrentUser() user: UserType,
    @Args('workspaceId') workspaceId: string,
    @Args('email') email: string,
    @Args('permission', { type: () => Permission }) permission: Permission,
    @Args('sendInviteMail', { nullable: true }) sendInviteMail: boolean
  ) {
    await this.permissions.check(workspaceId, user.id, Permission.Admin);

    if (permission === Permission.Owner) {
      throw new ForbiddenException('Cannot change owner');
    }

    const target = await this.users.findUserByEmail(email);

    if (target) {
      const originRecord = await this.prisma.userWorkspacePermission.findFirst({
        where: {
          workspaceId,
          userId: target.id,
        },
      });

      if (originRecord) {
        return originRecord.id;
      }

      const inviteId = await this.permissions.grant(
        workspaceId,
        target.id,
        permission
      );
      if (sendInviteMail) {
        const inviteInfo = await this.getInviteInfo(inviteId);

        try {
          await this.mailer.sendInviteEmail(email, inviteId, {
            workspace: {
              id: inviteInfo.workspace.id,
              name: inviteInfo.workspace.name,
              avatar: inviteInfo.workspace.avatar,
            },
            user: {
              avatar: inviteInfo.user?.avatarUrl || '',
              name: inviteInfo.user?.name || '',
            },
          });
        } catch (e) {
          const ret = await this.permissions.revoke(workspaceId, target.id);

          if (!ret) {
            this.logger.fatal(
              `failed to send ${workspaceId} invite email to ${email} and failed to revoke permission: ${inviteId}, ${e}`
            );
          } else {
            this.logger.warn(
              `failed to send ${workspaceId} invite email to ${email}, but successfully revoked permission: ${e}`
            );
          }

          return new InternalServerErrorException(e);
        }
      }
      return inviteId;
    } else {
      const user = await this.auth.createAnonymousUser(email);
      const inviteId = await this.permissions.grant(
        workspaceId,
        user.id,
        permission
      );
      if (sendInviteMail) {
        const inviteInfo = await this.getInviteInfo(inviteId);

        try {
          await this.mailer.sendInviteEmail(email, inviteId, {
            workspace: {
              id: inviteInfo.workspace.id,
              name: inviteInfo.workspace.name,
              avatar: inviteInfo.workspace.avatar,
            },
            user: {
              avatar: inviteInfo.user?.avatarUrl || '',
              name: inviteInfo.user?.name || '',
            },
          });
        } catch (e) {
          const ret = await this.permissions.revoke(workspaceId, user.id);

          if (!ret) {
            this.logger.fatal(
              `failed to send ${workspaceId} invite email to ${email} and failed to revoke permission: ${inviteId}, ${e}`
            );
          } else {
            this.logger.warn(
              `failed to send ${workspaceId} invite email to ${email}, but successfully revoked permission: ${e}`
            );
          }

          return new InternalServerErrorException(e);
        }
      }
      return inviteId;
    }
  }

  @Throttle(10, 30)
  @Public()
  @Query(() => InvitationType, {
    description: 'Update workspace',
  })
  async getInviteInfo(@Args('inviteId') inviteId: string) {
    const workspaceId = await this.prisma.userWorkspacePermission
      .findUniqueOrThrow({
        where: {
          id: inviteId,
        },
        select: {
          workspaceId: true,
        },
      })
      .then(({ workspaceId }) => workspaceId);

    const snapshot = await this.prisma.snapshot.findFirstOrThrow({
      where: {
        id: workspaceId,
        workspaceId,
      },
    });

    const doc = new Doc();

    applyUpdate(doc, new Uint8Array(snapshot.blob));
    const metaJSON = doc.getMap('meta').toJSON();

    const owner = await this.permissions.getWorkspaceOwner(workspaceId);
    const invitee = await this.permissions.getInvitationById(
      inviteId,
      workspaceId
    );

    let avatar = '';

    if (metaJSON.avatar) {
      const avatarBlob = await this.storage.getBlob(
        workspaceId,
        metaJSON.avatar
      );
      avatar = avatarBlob?.data.toString('base64') || '';
    }

    return {
      workspace: {
        name: metaJSON.name || '',
        avatar: avatar || defaultWorkspaceAvatar,
        id: workspaceId,
      },
      user: owner.user,
      invitee: invitee.user,
    };
  }

  @Mutation(() => Boolean)
  async revoke(
    @CurrentUser() user: UserType,
    @Args('workspaceId') workspaceId: string,
    @Args('userId') userId: string
  ) {
    await this.permissions.check(workspaceId, user.id, Permission.Admin);

    return this.permissions.revoke(workspaceId, userId);
  }

  @Mutation(() => Boolean)
  @Public()
  async acceptInviteById(
    @Args('workspaceId') workspaceId: string,
    @Args('inviteId') inviteId: string,
    @Args('sendAcceptMail', { nullable: true }) sendAcceptMail: boolean
  ) {
    const {
      invitee,
      user: inviter,
      workspace,
    } = await this.getInviteInfo(inviteId);

    if (!inviter || !invitee) {
      throw new ForbiddenException(
        `can not find inviter/invitee by inviteId: ${inviteId}`
      );
    }

    if (sendAcceptMail) {
      await this.mailer.sendAcceptedEmail(inviter.email, {
        inviteeName: invitee.name,
        workspaceName: workspace.name,
      });
    }

    return this.permissions.acceptById(workspaceId, inviteId);
  }

  @Mutation(() => Boolean)
  async acceptInvite(
    @CurrentUser() user: UserType,
    @Args('workspaceId') workspaceId: string
  ) {
    return this.permissions.accept(workspaceId, user.id);
  }

  @Mutation(() => Boolean)
  async leaveWorkspace(
    @CurrentUser() user: UserType,
    @Args('workspaceId') workspaceId: string,
    @Args('workspaceName') workspaceName: string,
    @Args('sendLeaveMail', { nullable: true }) sendLeaveMail: boolean
  ) {
    await this.permissions.check(workspaceId, user.id);

    const owner = await this.permissions.getWorkspaceOwner(workspaceId);

    if (!owner.user) {
      throw new ForbiddenException(
        `can not find owner by workspaceId: ${workspaceId}`
      );
    }

    if (sendLeaveMail) {
      await this.mailer.sendLeaveWorkspaceEmail(owner.user.email, {
        workspaceName,
        inviteeName: user.name,
      });
    }

    return this.permissions.revoke(workspaceId, user.id);
  }

  @Mutation(() => Boolean)
  async sharePage(
    @CurrentUser() user: UserType,
    @Args('workspaceId') workspaceId: string,
    @Args('pageId') pageId: string
  ) {
    const userWorkspace = await this.prisma.userWorkspacePermission.findFirst({
      where: {
        userId: user.id,
        workspaceId,
      },
    });
    if (!userWorkspace?.accepted) {
      throw new ForbiddenException('Permission denied');
    }
    return this.permissions.grantPage(workspaceId, pageId);
  }

  @Mutation(() => Boolean)
  async revokePage(
    @CurrentUser() user: UserType,
    @Args('workspaceId') workspaceId: string,
    @Args('pageId') pageId: string
  ) {
    await this.permissions.check(workspaceId, user.id, Permission.Admin);

    return this.permissions.revokePage(workspaceId, pageId);
  }

  @Query(() => [String], {
    description: 'List blobs of workspace',
  })
  async listBlobs(
    @CurrentUser() user: UserType,
    @Args('workspaceId') workspaceId: string
  ) {
    await this.permissions.check(workspaceId, user.id);

    return this.storage.listBlobs(workspaceId);
  }

  @Query(() => WorkspaceBlobSizes)
  async collectBlobSizes(
    @CurrentUser() user: UserType,
    @Args('workspaceId') workspaceId: string
  ) {
    await this.permissions.check(workspaceId, user.id);

    return this.storage.blobsSize([workspaceId]).then(size => ({ size }));
  }

  @Query(() => WorkspaceBlobSizes)
  async collectAllBlobSizes(@CurrentUser() user: UserType) {
    const workspaces = await this.prisma.userWorkspacePermission
      .findMany({
        where: {
          userId: user.id,
          accepted: true,
          type: Permission.Owner,
        },
        select: {
          workspace: {
            select: {
              id: true,
            },
          },
        },
      })
      .then(data => data.map(({ workspace }) => workspace.id));

    const size = await this.storage.blobsSize(workspaces);
    return { size };
  }

  @Query(() => WorkspaceBlobSizes)
  async checkBlobSize(
    @CurrentUser() user: UserType,
    @Args('workspaceId') workspaceId: string,
    @Args('size', { type: () => Float }) size: number
  ) {
    const canWrite = await this.permissions.tryCheck(
      workspaceId,
      user.id,
      Permission.Write
    );
    if (canWrite) {
      const { user } = await this.permissions.getWorkspaceOwner(workspaceId);
      if (user) {
        const quota = await this.users.getStorageQuotaById(user.id);
        const { size: currentSize } = await this.collectAllBlobSizes(user);

        return { size: quota - (size + currentSize) };
      }
    }
    return false;
  }

  @Mutation(() => String)
  async setBlob(
    @CurrentUser() user: UserType,
    @Args('workspaceId') workspaceId: string,
    @Args({ name: 'blob', type: () => GraphQLUpload })
    blob: FileUpload
  ) {
    await this.permissions.check(workspaceId, user.id, Permission.Write);

    // quota was apply to owner's account
    const { user: owner } =
      await this.permissions.getWorkspaceOwner(workspaceId);
    if (!owner) return new NotFoundException('Workspace owner not found');
    const quota = await this.users.getStorageQuotaById(owner.id);
    const { size } = await this.collectAllBlobSizes(owner);

    const checkExceeded = (recvSize: number) => {
      if (size + recvSize > quota) {
        this.logger.log(
          `storage size limit exceeded: ${size + recvSize} > ${quota}`
        );
        return true;
      } else {
        return false;
      }
    };

    if (checkExceeded(0)) {
      throw new ForbiddenException('storage size limit exceeded');
    }
    const buffer = await new Promise<Buffer>((resolve, reject) => {
      const stream = blob.createReadStream();
      const chunks: Uint8Array[] = [];
      stream.on('data', chunk => {
        chunks.push(chunk);

        // check size after receive each chunk to avoid unnecessary memory usage
        const bufferSize = chunks.reduce((acc, cur) => acc + cur.length, 0);
        if (checkExceeded(bufferSize)) {
          reject(new ForbiddenException('storage size limit exceeded'));
        }
      });
      stream.on('error', reject);
      stream.on('end', () => {
        const buffer = Buffer.concat(chunks);

        if (checkExceeded(buffer.length)) {
          reject(new ForbiddenException('storage size limit exceeded'));
        } else {
          resolve(buffer);
        }
      });
    });

    return this.storage.uploadBlob(workspaceId, buffer);
  }

  @Mutation(() => Boolean)
  async deleteBlob(
    @CurrentUser() user: UserType,
    @Args('workspaceId') workspaceId: string,
    @Args('hash') hash: string
  ) {
    await this.permissions.check(workspaceId, user.id);

    return this.storage.deleteBlob(workspaceId, hash);
  }
}