From 4d9f85d6a87e35b894cb3d3f8f8c98b2fcf3a9d2 Mon Sep 17 00:00:00 2001
From: shl0501 <shl050105010501@gmail.com>
Date: Tue, 11 Feb 2025 17:51:38 +0900
Subject: [PATCH] feat: implement ai request-restriction guard

---
 apps/server/src/ai/ai.controller.ts           |  2 ++
 apps/server/src/ai/ai.module.ts               |  3 ++-
 .../src/ai/guards/restrict-request.guard.ts   | 24 +++++++++++++++++++
 .../src/ai/guards/restrict-request.module.ts  |  9 +++++++
 apps/server/src/common/redis.module.ts        | 12 +++++++++-
 5 files changed, 48 insertions(+), 2 deletions(-)
 create mode 100644 apps/server/src/ai/guards/restrict-request.guard.ts
 create mode 100644 apps/server/src/ai/guards/restrict-request.module.ts

diff --git a/apps/server/src/ai/ai.controller.ts b/apps/server/src/ai/ai.controller.ts
index 81a7a04..1b318a7 100644
--- a/apps/server/src/ai/ai.controller.ts
+++ b/apps/server/src/ai/ai.controller.ts
@@ -5,6 +5,7 @@ import { AiService } from './ai.service';
 import { CreateHistoryDto } from './dto/create-history.dto';
 import { ImproveQuestionDto } from './dto/improve-question.dto';
 import { ShortenQuestionDto } from './dto/shorten-question.dto';
+import { AiRequestValidationGuard } from './guards/restrict-request.guard';
 import { CreateHistorySwagger } from './swagger/create-history.swagger';
 import { ImproveQuestionSwagger } from './swagger/improve-question.swagger';
 import { ShortenQuestionSwagger } from './swagger/shorten-question.swagger';
@@ -12,6 +13,7 @@ import { ShortenQuestionSwagger } from './swagger/shorten-question.swagger';
 import { SessionTokenValidationGuard } from '@common/guards/session-token-validation.guard';
 
 @Controller('ai')
+@UseGuards(AiRequestValidationGuard)
 export class AiController {
   constructor(private readonly aiService: AiService) {}
 
diff --git a/apps/server/src/ai/ai.module.ts b/apps/server/src/ai/ai.module.ts
index 1c2e666..904fc24 100644
--- a/apps/server/src/ai/ai.module.ts
+++ b/apps/server/src/ai/ai.module.ts
@@ -2,12 +2,13 @@ import { Module } from '@nestjs/common';
 
 import { AiController } from './ai.controller';
 import { AiService } from './ai.service';
+import { AiRequestValidationModule } from './guards/restrict-request.module';
 
 import { SessionTokenModule } from '@common/guards/session-token.module';
 import { PrismaModule } from '@prisma-alias/prisma.module';
 
 @Module({
-  imports: [PrismaModule, SessionTokenModule],
+  imports: [PrismaModule, SessionTokenModule, AiRequestValidationModule],
   providers: [AiService],
   controllers: [AiController],
 })
diff --git a/apps/server/src/ai/guards/restrict-request.guard.ts b/apps/server/src/ai/guards/restrict-request.guard.ts
new file mode 100644
index 0000000..877e644
--- /dev/null
+++ b/apps/server/src/ai/guards/restrict-request.guard.ts
@@ -0,0 +1,24 @@
+import { CanActivate, ExecutionContext, ForbiddenException, Inject, Injectable } from '@nestjs/common';
+import Redis from 'ioredis';
+
+@Injectable()
+export class AiRequestValidationGuard implements CanActivate {
+  private readonly ttl = 3;
+
+  constructor(@Inject('REDIS_RESTRICT_TOKEN') private readonly restrictRedisClient: Redis) {}
+
+  async validateAiRequest(token: string) {
+    const exists = await this.restrictRedisClient.get(token);
+    if (Number(exists) === 1) {
+      throw new ForbiddenException('저희의 토큰은 소중합니다.');
+    }
+    await this.restrictRedisClient.set(token, 1, 'EX', this.ttl);
+  }
+  async canActivate(context: ExecutionContext) {
+    const request = context.switchToHttp().getRequest();
+    const token = request.body?.token || request.query?.token;
+
+    await this.validateAiRequest(token);
+    return true;
+  }
+}
diff --git a/apps/server/src/ai/guards/restrict-request.module.ts b/apps/server/src/ai/guards/restrict-request.module.ts
new file mode 100644
index 0000000..86859b9
--- /dev/null
+++ b/apps/server/src/ai/guards/restrict-request.module.ts
@@ -0,0 +1,9 @@
+import { Module } from '@nestjs/common';
+
+import { AiRequestValidationGuard } from './restrict-request.guard';
+
+@Module({
+  providers: [AiRequestValidationGuard],
+  exports: [AiRequestValidationGuard],
+})
+export class AiRequestValidationModule {}
diff --git a/apps/server/src/common/redis.module.ts b/apps/server/src/common/redis.module.ts
index 8b6d30f..7c0b9d1 100644
--- a/apps/server/src/common/redis.module.ts
+++ b/apps/server/src/common/redis.module.ts
@@ -34,7 +34,17 @@ import Redis from 'ioredis';
         });
       },
     },
+    {
+      provide: 'REDIS_RESTRICT_TOKEN',
+      useFactory: () => {
+        return new Redis({
+          host: process.env.REDIS_HOST,
+          port: Number(process.env.REDIS_PORT),
+          db: 3,
+        });
+      },
+    },
   ],
-  exports: ['REDIS_SESSION', 'REDIS_TOKEN', 'REDIS_REFRESH_TOKEN'],
+  exports: ['REDIS_SESSION', 'REDIS_TOKEN', 'REDIS_REFRESH_TOKEN', 'REDIS_RESTRICT_TOKEN'],
 })
 export class RedisModule {}