refactor: merge token and access code

app/api/auth.ts

@@ -0,0 +1,70 @@
+import { NextRequest } from "next/server";
+import { getServerSideConfig } from "../config/server";
+import md5 from "spark-md5";
+
+const serverConfig = getServerSideConfig();
+
+function getIP(req: NextRequest) {
+  let ip = req.ip ?? req.headers.get("x-real-ip");
+  const forwardedFor = req.headers.get("x-forwarded-for");
+
+  if (!ip && forwardedFor) {
+    ip = forwardedFor.split(",").at(0) ?? "";
+  }
+
+  return ip;
+}
+
+function parseApiKey(bearToken: string) {
+  const token = bearToken.trim().replaceAll("Bearer ", "").trim();
+  const isOpenAiKey = token.startsWith("sk-");
+
+  return {
+    accessCode: isOpenAiKey ? "" : token,
+    apiKey: isOpenAiKey ? token : "",
+  };
+}
+
+export function auth(req: NextRequest) {
+  const authToken = req.headers.get("Authorization") ?? "";
+
+  // check if it is openai api key or user token
+  const { accessCode, apiKey: token } = parseApiKey(authToken);
+
+  const hashedCode = md5.hash(accessCode ?? "").trim();
+
+  console.log("[Auth] allowed hashed codes: ", [...serverConfig.codes]);
+  console.log("[Auth] got access code:", accessCode);
+  console.log("[Auth] hashed access code:", hashedCode);
+  console.log("[User IP] ", getIP(req));
+  console.log("[Time] ", new Date().toLocaleString());
+
+  if (serverConfig.needCode && !serverConfig.codes.has(hashedCode) && !token) {
+    return {
+      error: true,
+      needAccessCode: true,
+      msg: "Please go settings page and fill your access code.",
+    };
+  }
+
+  // if user does not provide an api key, inject system api key
+  if (!token) {
+    const apiKey = serverConfig.apiKey;
+    if (apiKey) {
+      console.log("[Auth] use system api key");
+      req.headers.set("Authorization", `Bearer ${apiKey}`);
+    } else {
+      console.log("[Auth] admin did not provide an api key");
+      return {
+        error: true,
+        msg: "Empty Api Key",
+      };
+    }
+  } else {
+    console.log("[Auth] use user api key");
+  }
+
+  return {
+    error: false,
+  };
+}

app/api/common.ts

@@ -6,8 +6,11 @@ const PROTOCOL = process.env.PROTOCOL ?? DEFAULT_PROTOCOL;
 const BASE_URL = process.env.BASE_URL ?? OPENAI_URL;
 
 export async function requestOpenai(req: NextRequest) {
-  const apiKey = req.headers.get("token");
-  const openaiPath = req.headers.get("path");
+  const authValue = req.headers.get("Authorization") ?? "";
+  const openaiPath = `${req.nextUrl.pathname}${req.nextUrl.search}`.replaceAll(
+    "/api/openai/",
+    "",
+  );
 
   let baseUrl = BASE_URL;
 
@@ -22,10 +25,14 @@ export async function requestOpenai(req: NextRequest) {
     console.log("[Org ID]", process.env.OPENAI_ORG_ID);
   }
 
+  if (!authValue || !authValue.startsWith("Bearer sk-")) {
+    console.error("[OpenAI Request] invlid api key provided", authValue);
+  }
+
   return fetch(`${baseUrl}/${openaiPath}`, {
     headers: {
       "Content-Type": "application/json",
-      Authorization: `Bearer ${apiKey}`,
+      Authorization: authValue,
       ...(process.env.OPENAI_ORG_ID && {
         "OpenAI-Organization": process.env.OPENAI_ORG_ID,
       }),

app/api/config/route.ts

@@ -14,7 +14,7 @@ declare global {
   type DangerConfig = typeof DANGER_CONFIG;
 }
 
-export async function POST(req: NextRequest) {
+export async function POST() {
   return NextResponse.json({
     needCode: serverConfig.needCode,
   });

app/api/openai/[...path]/route.ts

@@ -1,6 +1,7 @@
 import { createParser } from "eventsource-parser";
 import { NextRequest, NextResponse } from "next/server";
-import { requestOpenai } from "../common";
+import { auth } from "../../auth";
+import { requestOpenai } from "../../common";
 
 async function createStream(res: Response) {
   const encoder = new TextEncoder();
@@ -43,7 +44,19 @@ function formatResponse(msg: any) {
   return new Response(jsonMsg);
 }
 
-async function makeRequest(req: NextRequest) {
+async function handle(
+  req: NextRequest,
+  { params }: { params: { path: string[] } },
+) {
+  console.log("[OpenAI Route] params ", params);
+
+  const authResult = auth(req);
+  if (authResult.error) {
+    return NextResponse.json(authResult, {
+      status: 401,
+    });
+  }
+
   try {
     const api = await requestOpenai(req);
 
@@ -52,7 +65,9 @@ async function makeRequest(req: NextRequest) {
     // streaming response
     if (contentType.includes("stream")) {
       const stream = await createStream(api);
-      return new Response(stream);
+      const res = new Response(stream);
+      res.headers.set("Content-Type", contentType);
+      return res;
     }
 
     // try to parse error msg
@@ -80,12 +95,7 @@ async function makeRequest(req: NextRequest) {
   }
 }
 
-export async function POST(req: NextRequest) {
-  return makeRequest(req);
-}
-
-export async function GET(req: NextRequest) {
-  return makeRequest(req);
-}
+export const GET = handle;
+export const POST = handle;
 
 export const runtime = "edge";