From 4c08c1faed4c883acf627eb347d783c32bb6aeb1 Mon Sep 17 00:00:00 2001
From: zlwu <zilong.wu@gmail.com>
Date: Wed, 24 Aug 2022 23:11:27 +0800
Subject: [PATCH] add role guard and decorator

---
 src/auth/roles.guard.ts                  | 31 ++++++++++++++++++++++++
 src/common/decorators/roles.decorator.ts |  6 +++++
 src/users/users.resolver.ts              |  6 ++++-
 3 files changed, 42 insertions(+), 1 deletion(-)
 create mode 100644 src/auth/roles.guard.ts
 create mode 100644 src/common/decorators/roles.decorator.ts

diff --git a/src/auth/roles.guard.ts b/src/auth/roles.guard.ts
new file mode 100644
index 00000000..9a91806f
--- /dev/null
+++ b/src/auth/roles.guard.ts
@@ -0,0 +1,31 @@
+import { Injectable, CanActivate, ExecutionContext } from '@nestjs/common';
+import { Reflector } from '@nestjs/core';
+import { GqlExecutionContext } from '@nestjs/graphql';
+import { Role } from '@prisma/client';
+import { PrismaService } from 'nestjs-prisma';
+import { ROLES_KEY } from 'src/common/decorators/roles.decorator';
+
+@Injectable()
+export class RolesGuard implements CanActivate {
+  constructor(private prisma: PrismaService, private reflector: Reflector) {}
+
+  canActivate(context: ExecutionContext): boolean {
+    const requiredRoles = this.reflector.getAllAndOverride<Role[]>(ROLES_KEY, [
+      context.getHandler(),
+      context.getClass(),
+    ]);
+
+    if (!requiredRoles) {
+      return true;
+    }
+
+    const gqlContext = GqlExecutionContext.create(context);
+    const req = gqlContext.getContext().req;
+    const user = req.user;
+    if (user) {
+      return requiredRoles.includes(user.role);
+    } else {
+      return false;
+    }
+  }
+}
diff --git a/src/common/decorators/roles.decorator.ts b/src/common/decorators/roles.decorator.ts
new file mode 100644
index 00000000..7f5fe5c2
--- /dev/null
+++ b/src/common/decorators/roles.decorator.ts
@@ -0,0 +1,6 @@
+import { SetMetadata } from '@nestjs/common';
+import { Role } from '@prisma/client';
+
+export const ROLES_KEY = 'roles';
+
+export const Roles = (...roles: Role[]) => SetMetadata(ROLES_KEY, roles);
diff --git a/src/users/users.resolver.ts b/src/users/users.resolver.ts
index 91ef5c48..3509165b 100644
--- a/src/users/users.resolver.ts
+++ b/src/users/users.resolver.ts
@@ -8,8 +8,11 @@ import {
   ResolveField,
 } from '@nestjs/graphql';
 import { UseGuards } from '@nestjs/common';
+import { Role } from '@prisma/client';
 import { UserEntity } from 'src/common/decorators/user.decorator';
 import { GqlAuthGuard } from 'src/auth/gql-auth.guard';
+import { RolesGuard } from 'src/auth/roles.guard';
+import { Roles } from 'src/common/decorators/roles.decorator';
 import { UsersService } from './users.service';
 import { User } from './models/user.model';
 import { ChangePasswordInput } from './dto/change-password.input';
@@ -28,7 +31,8 @@ export class UsersResolver {
     return user;
   }
 
-  @UseGuards(GqlAuthGuard)
+  @Roles(Role.USER, Role.ADMIN)
+  @UseGuards(GqlAuthGuard, RolesGuard)
   @Mutation(() => User)
   async updateUser(
     @UserEntity() user: User,