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";

app/requests.ts

@@ -44,14 +44,21 @@ const makeRequestParam = (
 
 function getHeaders() {
   const accessStore = useAccessStore.getState();
-  let headers: Record<string, string> = {};
+  const headers = {
+    Authorization: "",
+  };
 
-  if (accessStore.enabledAccessControl()) {
-    headers["access-code"] = accessStore.accessCode;
-  }
+  const makeBearer = (token: string) => `Bearer ${token.trim()}`;
+  const validString = (x: string) => x && x.length > 0;
 
-  if (accessStore.token && accessStore.token.length > 0) {
-    headers["token"] = accessStore.token;
+  // use user's api key first
+  if (validString(accessStore.token)) {
+    headers.Authorization = makeBearer(accessStore.token);
+  } else if (
+    accessStore.enabledAccessControl() &&
+    validString(accessStore.accessCode)
+  ) {
+    headers.Authorization = makeBearer(accessStore.accessCode);
   }
 
   return headers;
@@ -59,13 +66,8 @@ function getHeaders() {
 
 export function requestOpenaiClient(path: string) {
   return (body: any, method = "POST") =>
-    fetch("/api/openai", {
+    fetch("/api/openai/" + path, {
       method,
-      headers: {
-        "Content-Type": "application/json",
-        path,
-        ...getHeaders(),
-      },
       body: body && JSON.stringify(body),
     });
 }
@@ -161,16 +163,16 @@ export async function requestChatStream(
   const reqTimeoutId = setTimeout(() => controller.abort(), TIME_OUT_MS);
 
   try {
-    const res = await fetch("/api/openai", {
+    const res = await fetch("/api/openai/v1/chat/completions", {
       method: "POST",
       headers: {
         "Content-Type": "application/json",
-        path: "v1/chat/completions",
         ...getHeaders(),
       },
       body: JSON.stringify(req),
       signal: controller.signal,
     });
+
     clearTimeout(reqTimeoutId);
 
     let responseText = "";

app/store/access.ts

@@ -1,6 +1,7 @@
 import { create } from "zustand";
 import { persist } from "zustand/middleware";
 import { StoreKey } from "../constant";
+import { BOT_HELLO } from "./chat";
 
 export interface AccessControlStore {
   accessCode: string;

middleware.ts

@@ -1,72 +0,0 @@
-import { NextRequest, NextResponse } from "next/server";
-import { getServerSideConfig } from "./app/config/server";
-import md5 from "spark-md5";
-
-export const config = {
-  matcher: ["/api/openai"],
-};
-
-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;
-}
-
-export function middleware(req: NextRequest) {
-  const accessCode = req.headers.get("access-code");
-  const token = req.headers.get("token");
-  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 NextResponse.json(
-      {
-        error: true,
-        needAccessCode: true,
-        msg: "Please go settings page and fill your access code.",
-      },
-      {
-        status: 401,
-      },
-    );
-  }
-
-  // inject api key
-  if (!token) {
-    const apiKey = serverConfig.apiKey;
-    if (apiKey) {
-      console.log("[Auth] set system token");
-      req.headers.set("token", apiKey);
-    } else {
-      return NextResponse.json(
-        {
-          error: true,
-          msg: "Empty Api Key",
-        },
-        {
-          status: 401,
-        },
-      );
-    }
-  } else {
-    console.log("[Auth] set user token");
-  }
-
-  return NextResponse.next({
-    request: {
-      headers: req.headers,
-    },
-  });
-}

next.config.mjs

@@ -15,4 +15,4 @@ const nextConfig = {
   output: "standalone",
 };
 
-module.exports = nextConfig;
+export default nextConfig;