From 0d4611052e75cbe9b2dc9309b60435178dcab663 Mon Sep 17 00:00:00 2001
From: Yidadaa <yidadaa@qq.com>
Date: Tue, 13 Jun 2023 00:39:29 +0800
Subject: [PATCH] feat: white url list for openai security

---
 app/api/openai/[...path]/route.ts | 18 ++++++++++++++++++
 app/client/platforms/openai.ts    | 13 +++++--------
 app/constant.ts                   |  6 ++++++
 3 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/app/api/openai/[...path]/route.ts b/app/api/openai/[...path]/route.ts
index 981749e7..04f3b6da 100644
--- a/app/api/openai/[...path]/route.ts
+++ b/app/api/openai/[...path]/route.ts
@@ -1,14 +1,32 @@
+import { OpenaiPath } from "@/app/constant";
 import { prettyObject } from "@/app/utils/format";
 import { NextRequest, NextResponse } from "next/server";
 import { auth } from "../../auth";
 import { requestOpenai } from "../../common";
 
+const ALLOWD_PATH = new Set(Object.values(OpenaiPath));
+
 async function handle(
   req: NextRequest,
   { params }: { params: { path: string[] } },
 ) {
   console.log("[OpenAI Route] params ", params);
 
+  const subpath = params.path.join("/");
+
+  if (!ALLOWD_PATH.has(subpath)) {
+    console.log("[OpenAI Route] forbidden path ", subpath);
+    return NextResponse.json(
+      {
+        error: true,
+        msg: "you are not allowed to request " + subpath,
+      },
+      {
+        status: 403,
+      },
+    );
+  }
+
   const authResult = auth(req);
   if (authResult.error) {
     return NextResponse.json(authResult, {
diff --git a/app/client/platforms/openai.ts b/app/client/platforms/openai.ts
index 84c4a2df..fd4c3365 100644
--- a/app/client/platforms/openai.ts
+++ b/app/client/platforms/openai.ts
@@ -1,4 +1,4 @@
-import { REQUEST_TIMEOUT_MS } from "@/app/constant";
+import { OpenaiPath, REQUEST_TIMEOUT_MS } from "@/app/constant";
 import { useAccessStore, useAppConfig, useChatStore } from "@/app/store";
 
 import { ChatOptions, getHeaders, LLMApi, LLMUsage } from "../api";
@@ -10,10 +10,6 @@ import {
 import { prettyObject } from "@/app/utils/format";
 
 export class ChatGPTApi implements LLMApi {
-  public ChatPath = "v1/chat/completions";
-  public UsagePath = "dashboard/billing/usage";
-  public SubsPath = "dashboard/billing/subscription";
-
   path(path: string): string {
     let openaiUrl = useAccessStore.getState().openaiUrl;
     if (openaiUrl.endsWith("/")) {
@@ -55,7 +51,7 @@ export class ChatGPTApi implements LLMApi {
     options.onController?.(controller);
 
     try {
-      const chatPath = this.path(this.ChatPath);
+      const chatPath = this.path(OpenaiPath.ChatPath);
       const chatPayload = {
         method: "POST",
         body: JSON.stringify(requestPayload),
@@ -177,14 +173,14 @@ export class ChatGPTApi implements LLMApi {
     const [used, subs] = await Promise.all([
       fetch(
         this.path(
-          `${this.UsagePath}?start_date=${startDate}&end_date=${endDate}`,
+          `${OpenaiPath.UsagePath}?start_date=${startDate}&end_date=${endDate}`,
         ),
         {
           method: "GET",
           headers: getHeaders(),
         },
       ),
-      fetch(this.path(this.SubsPath), {
+      fetch(this.path(OpenaiPath.SubsPath), {
         method: "GET",
         headers: getHeaders(),
       }),
@@ -228,3 +224,4 @@ export class ChatGPTApi implements LLMApi {
     } as LLMUsage;
   }
 }
+export { OpenaiPath };
diff --git a/app/constant.ts b/app/constant.ts
index b640919e..0f805275 100644
--- a/app/constant.ts
+++ b/app/constant.ts
@@ -45,3 +45,9 @@ export const LAST_INPUT_KEY = "last-input";
 export const REQUEST_TIMEOUT_MS = 60000;
 
 export const EXPORT_MESSAGE_CLASS_NAME = "export-markdown";
+
+export const OpenaiPath = {
+  ChatPath: "v1/chat/completions",
+  UsagePath: "dashboard/billing/usage",
+  SubsPath: "dashboard/billing/subscription",
+};