From d21e48a1a3ea40445fa05be7b4d04b40f45b4b49 Mon Sep 17 00:00:00 2001
From: diced <pranaco2@gmail.com>
Date: Mon, 20 Jun 2022 10:49:24 -0700
Subject: [PATCH] fix(api): improve ratelimits

---
 .../migration.sql                             |  9 ++++
 prisma/schema.prisma                          |  2 +-
 src/lib/middleware/withZipline.ts             |  6 +--
 src/pages/api/upload.ts                       | 44 ++++++++++++-------
 4 files changed, 41 insertions(+), 20 deletions(-)
 create mode 100644 prisma/migrations/20220620173040_ratelimit_date/migration.sql

diff --git a/prisma/migrations/20220620173040_ratelimit_date/migration.sql b/prisma/migrations/20220620173040_ratelimit_date/migration.sql
new file mode 100644
index 0000000..b30e047
--- /dev/null
+++ b/prisma/migrations/20220620173040_ratelimit_date/migration.sql
@@ -0,0 +1,9 @@
+/*
+  Warnings:
+
+  - You are about to drop the column `ratelimited` on the `User` table. All the data in the column will be lost.
+
+*/
+-- AlterTable
+ALTER TABLE "User" DROP COLUMN "ratelimited",
+ADD COLUMN     "ratelimit" TIMESTAMP(3);
diff --git a/prisma/schema.prisma b/prisma/schema.prisma
index 78d8af5..c06a012 100644
--- a/prisma/schema.prisma
+++ b/prisma/schema.prisma
@@ -17,7 +17,7 @@ model User {
   embedTitle    String?
   embedColor    String    @default("#2f3136")
   embedSiteName String?   @default("{image.file} • {user.name}")
-  ratelimited   Boolean   @default(false)
+  ratelimit     DateTime?
   domains       String[]
   images        Image[]
   urls          Url[]
diff --git a/src/lib/middleware/withZipline.ts b/src/lib/middleware/withZipline.ts
index d757dd8..fddf776 100644
--- a/src/lib/middleware/withZipline.ts
+++ b/src/lib/middleware/withZipline.ts
@@ -37,7 +37,7 @@ export type NextApiRes = NextApiResponse & {
   forbid: (message: string, extra?: any) => void;
   bad: (message: string) => void;
   json: (json: any) => void;
-  ratelimited: () => void;
+  ratelimited: (remaining: number) => void;
   setCookie: (name: string, value: unknown, options: CookieSerializeOptions) => void;
 }
 
@@ -69,9 +69,9 @@ export const withZipline = (handler: (req: NextApiRequest, res: NextApiResponse)
     });
   };
 
-  res.ratelimited = () => {
+  res.ratelimited = (remaining: number) => {
     res.status(429);
-
+    res.setHeader('X-Ratelimit-Remaining', Math.floor(remaining / 1000));
     res.json({
       error: '429: ratelimited',
     });
diff --git a/src/pages/api/upload.ts b/src/pages/api/upload.ts
index 8bc2c6a..bcfed65 100644
--- a/src/pages/api/upload.ts
+++ b/src/pages/api/upload.ts
@@ -22,7 +22,21 @@ async function handler(req: NextApiReq, res: NextApiRes) {
   });
 
   if (!user) return res.forbid('authorization incorect');
-  if (user.ratelimited) return res.ratelimited();
+  if (user.ratelimit) {
+    const remaining = user.ratelimit.getTime() - Date.now();
+    if (remaining <= 0) {
+      await prisma.user.update({
+        where: {
+          id: user.id,
+        },
+        data: {
+          ratelimit: null,
+        },
+      });
+    } else {
+      return res.ratelimited(remaining);
+    }
+  }
 
   await run(uploader.array('file'))(req, res);
 
@@ -85,28 +99,26 @@ async function handler(req: NextApiReq, res: NextApiRes) {
     }
   }
 
-  if (user.administrator && zconfig.ratelimit.admin !== 0) {
+  if (user.administrator && zconfig.ratelimit.admin > 0) {
     await prisma.user.update({
       where: {
         id: user.id,
       },
       data: {
-        ratelimited: true,
+        ratelimit: new Date(Date.now() + (zconfig.ratelimit.admin * 1000)),
       },
     });
-    setTimeout(async () => await prisma.user.update({ where: { id: user.id }, data: { ratelimited: false } }), zconfig.ratelimit.admin * 1000).unref();
-  }
-
-  if (!user.administrator && zconfig.ratelimit.user !== 0) {
-    await prisma.user.update({
-      where: {
-        id: user.id,
-      },
-      data: {
-        ratelimited: true,
-      },
-    });
-    setTimeout(async () => await prisma.user.update({ where: { id: user.id }, data: { ratelimited: false } }), zconfig.ratelimit.user * 1000).unref();
+  } else if (!user.administrator && zconfig.ratelimit.user > 0) {
+    if (user.administrator && zconfig.ratelimit.user > 0) {
+      await prisma.user.update({
+        where: {
+          id: user.id,
+        },
+        data: {
+          ratelimit: new Date(Date.now() + (zconfig.ratelimit.user * 1000)),
+        },
+      });
+    }
   }
 
   return res.json({ files });