From 18f38c072a16245b18e813f2db9502450b503285 Mon Sep 17 00:00:00 2001
From: cgsv <cgsv@cgsvdeMacBook-Pro.local>
Date: Tue, 7 Mar 2023 18:01:20 +0800
Subject: [PATCH] add rate limit

---
 .example.env      |  4 +++-
 lib/upstash.ts    | 13 +++++++++++++
 middleware.ts     | 42 ++++++++++++++++++++++++++++++++++++++++++
 package-lock.json | 21 +++++++++++++++++++++
 package.json      |  1 +
 5 files changed, 80 insertions(+), 1 deletion(-)
 create mode 100644 lib/upstash.ts
 create mode 100644 middleware.ts

diff --git a/.example.env b/.example.env
index 7c5ee8d..c62b114 100644
--- a/.example.env
+++ b/.example.env
@@ -1,4 +1,6 @@
 OPENAI_API_KEY=sk-xxx
 BILIBILI_SESSION_TOKEN=
 UPSTASH_REDIS_REST_URL=
-UPSTASH_REDIS_REST_TOKEN=
\ No newline at end of file
+UPSTASH_REDIS_REST_TOKEN=
+UPSTASH_RATE_REDIS_REST_URL=
+UPSTASH_RATE_REDIS_REST_TOKEN=
\ No newline at end of file
diff --git a/lib/upstash.ts b/lib/upstash.ts
new file mode 100644
index 0000000..149d457
--- /dev/null
+++ b/lib/upstash.ts
@@ -0,0 +1,13 @@
+import { Ratelimit } from "@upstash/ratelimit";
+import { Redis } from "@upstash/redis";
+import { RATE_LIMIT_COUNT } from "../utils/constants";
+
+export const ratelimit = new Ratelimit({
+  redis: new Redis({
+    url: process.env.UPSTASH_RATE_REDIS_REST_URL,
+    token: process.env.UPSTASH_RATE_REDIS_REST_TOKEN,
+  }),
+  // 速率限制算法 https://github.com/upstash/ratelimit#ratelimiting-algorithms
+  limiter: Ratelimit.fixedWindow(RATE_LIMIT_COUNT, "1 d"),
+  analytics: true, // <- Enable analytics
+});
diff --git a/middleware.ts b/middleware.ts
new file mode 100644
index 0000000..f2ec71d
--- /dev/null
+++ b/middleware.ts
@@ -0,0 +1,42 @@
+import { Redis } from "@upstash/redis";
+import type { NextFetchEvent, NextRequest } from "next/server";
+import { NextResponse } from "next/server";
+// import { validateLicenseKey } from "./lib/lemon";
+import { checkOpenaiApiKeys } from "./lib/openai/openai";
+import { ratelimit } from "./lib/upstash"
+import { isDev } from "./utils/env";
+
+const redis = Redis.fromEnv();
+
+export async function middleware(req: NextRequest, context: NextFetchEvent) {
+  const { apiKey } = await req.json();
+//   const result = await redis.get<string>(bvId);
+//   if (!isDev && result) {
+//     console.log("hit cache for ", bvId);
+//     return NextResponse.json(result);
+//   }
+
+  // licenseKeys
+  if (apiKey) {
+    if (checkOpenaiApiKeys(apiKey)) {
+      return NextResponse.next();
+    }
+
+    // // 3. something-invalid-sdalkjfasncs-key
+    // if (!(await validateLicenseKey(apiKey, bvId))) {
+    //   return NextResponse.redirect(new URL("/shop", req.url));
+    // }
+  }
+  // TODO: unique to a user (userid, email etc) instead of IP
+  const identifier = req.ip ?? "127.0.0.7";
+  const { success, remaining } = await ratelimit.limit("trans-" + identifier);
+  console.log(`======== ip ${identifier}, remaining: ${remaining} ========`);
+  if (!apiKey && !success) {
+    // return NextResponse.redirect(new URL("/shop", req.url));
+    return NextResponse.error();
+  }
+}
+
+export const config = {
+  matcher: "/api/translate",
+};
diff --git a/package-lock.json b/package-lock.json
index 395eb20..df09ab3 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -11,6 +11,7 @@
         "@types/node": "18.14.4",
         "@types/react": "18.0.28",
         "@types/react-dom": "18.0.11",
+        "@upstash/ratelimit": "^0.3.9",
         "@upstash/redis": "^1.20.1",
         "detect-file-encoding-and-language": "^2.3.2",
         "eslint": "8.35.0",
@@ -1720,6 +1721,26 @@
         "url": "https://opencollective.com/typescript-eslint"
       }
     },
+    "node_modules/@upstash/core-analytics": {
+      "version": "0.0.6",
+      "resolved": "https://mirrors.cloud.tencent.com/npm/@upstash/core-analytics/-/core-analytics-0.0.6.tgz",
+      "integrity": "sha512-cpPSR0XJAJs4Ddz9nq3tINlPS5aLfWVCqhhtHnXt4p7qr5+/Znlt1Es736poB/9rnl1hAHrOsOvVj46NEXcVqA==",
+      "dependencies": {
+        "@upstash/redis": "^1.19.3"
+      },
+      "engines": {
+        "node": ">=16.0.0"
+      }
+    },
+    "node_modules/@upstash/ratelimit": {
+      "version": "0.3.9",
+      "resolved": "https://mirrors.cloud.tencent.com/npm/@upstash/ratelimit/-/ratelimit-0.3.9.tgz",
+      "integrity": "sha512-bowxM/beAw3AMKeVObN0V10a0aeno/wmjf/BY1ydPwl6V8JRcCjMfTqkmFI2N+M9LekMALspM5tC+tJcuuyYEg==",
+      "dependencies": {
+        "@upstash/core-analytics": "0.0.6",
+        "@upstash/redis": "^1.20.1"
+      }
+    },
     "node_modules/@upstash/redis": {
       "version": "1.20.1",
       "resolved": "https://mirrors.cloud.tencent.com/npm/@upstash/redis/-/redis-1.20.1.tgz",
diff --git a/package.json b/package.json
index f91b1dc..4993611 100644
--- a/package.json
+++ b/package.json
@@ -13,6 +13,7 @@
     "@types/node": "18.14.4",
     "@types/react": "18.0.28",
     "@types/react-dom": "18.0.11",
+    "@upstash/ratelimit": "^0.3.9",
     "@upstash/redis": "^1.20.1",
     "detect-file-encoding-and-language": "^2.3.2",
     "eslint": "8.35.0",