From 48ebd74859042aaa3d8cdbb74a61bd7391400b89 Mon Sep 17 00:00:00 2001 From: Yidadaa Date: Wed, 3 May 2023 23:08:37 +0800 Subject: [PATCH] refactor: merge token and access code --- app/api/auth.ts | 70 ++++++++++++++++++++++++ app/api/common.ts | 13 +++-- app/api/config/route.ts | 2 +- app/api/openai/{ => [...path]}/route.ts | 30 +++++++---- app/requests.ts | 30 ++++++----- app/store/access.ts | 1 + middleware.ts | 72 ------------------------- next.config.js => next.config.mjs | 2 +- 8 files changed, 119 insertions(+), 101 deletions(-) create mode 100644 app/api/auth.ts rename app/api/openai/{ => [...path]}/route.ts (81%) delete mode 100644 middleware.ts rename next.config.js => next.config.mjs (90%) diff --git a/app/api/auth.ts b/app/api/auth.ts new file mode 100644 index 00000000..9a834832 --- /dev/null +++ b/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, + }; +} diff --git a/app/api/common.ts b/app/api/common.ts index a86d6861..861caf3b 100644 --- a/app/api/common.ts +++ b/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, }), diff --git a/app/api/config/route.ts b/app/api/config/route.ts index 65290a47..4936c109 100644 --- a/app/api/config/route.ts +++ b/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, }); diff --git a/app/api/openai/route.ts b/app/api/openai/[...path]/route.ts similarity index 81% rename from app/api/openai/route.ts rename to app/api/openai/[...path]/route.ts index 15d56e54..1ca103c6 100644 --- a/app/api/openai/route.ts +++ b/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"; diff --git a/app/requests.ts b/app/requests.ts index 790157d8..497c8737 100644 --- a/app/requests.ts +++ b/app/requests.ts @@ -44,14 +44,21 @@ const makeRequestParam = ( function getHeaders() { const accessStore = useAccessStore.getState(); - let headers: Record = {}; + 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 = ""; diff --git a/app/store/access.ts b/app/store/access.ts index e72052b4..ba0faf93 100644 --- a/app/store/access.ts +++ b/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; diff --git a/middleware.ts b/middleware.ts deleted file mode 100644 index 5cc5473a..00000000 --- a/middleware.ts +++ /dev/null @@ -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, - }, - }); -} diff --git a/next.config.js b/next.config.mjs similarity index 90% rename from next.config.js rename to next.config.mjs index f7d5ff08..82647954 100644 --- a/next.config.js +++ b/next.config.mjs @@ -15,4 +15,4 @@ const nextConfig = { output: "standalone", }; -module.exports = nextConfig; +export default nextConfig;