rate limiting

app/api/collections/route.ts

@@ -7,11 +7,12 @@ const prisma = new PrismaClient()
 export async function GET(request: Request) {
   const { headers } = request
   const { searchParams } = new URL(request.url)
+
   const userId = searchParams.get('userId') ?? undefined
   const page = searchParams.get('page') ? parseInt(searchParams.get('page') as string) : 0
-
+  
   let is_admin = false
-  if (headers.get('Authorization') ) {
+  if (headers.get('Authorization')) {
     is_admin = await isAdmin(headers) as boolean
   }
 

app/api/route.ts

@@ -1,14 +1,14 @@
+import { NextRequest } from 'next/server'
 import { PrismaClient } from '@prisma/client'
 
 import list_styles from "@/assets/list_styles.json"
 import { UploaderDataset } from '../../utils/uploader'
 import { isTextNSFW } from '@/utils/checker/prompt'
-import { isImageNSFW } from '@/utils/checker/image'
 
 const prisma = new PrismaClient()
 
 export async function POST(
-  request: Request,
+  request: NextRequest,
 ) {
   const global_headers = {
     Authorization: `Bearer ${process.env.HF_TOKEN}`,
@@ -17,8 +17,25 @@ export async function POST(
   }
 
   const { inputs, style, userId } = await request.json()
+  const { headers } = request
+
   const findStyle = list_styles.find((item) => item.name === style)
 
+  const ip_address = headers.get("x-forwarded-for") ?? request.ip 
+
+  if (!headers.get('Authorization')) {
+    const count = await prisma.collection.count({
+      where: {
+        ip_address,
+        createdAt: {
+          gte: new Date(Date.now() - 24 * 60 * 60 * 1000)
+        }
+      }
+    })
+  
+    if (count > 5) return Response.json({ status: 429, ok: false, message: "You have reached the limit of 5 images per day." });
+  }
+
   const textIsNSFW = await isTextNSFW(inputs, global_headers)
   if (textIsNSFW) return Response.json({ status: 401, ok: false, message: "Prompt doesn’t work, try another prompt" });
 

prisma/schema.prisma

@@ -11,6 +11,7 @@ model Collection {
   id         Int      @id @default(autoincrement())
   prompt     String
   file_name  String
+  ip_address String   @default("")
   is_visible Boolean  @default(false)
   createdAt  DateTime @default(now())
   userId     String?  @default("")