From 1b874aefbc4d09df1348c66ac19b20e5c0da870e Mon Sep 17 00:00:00 2001
From: Mengxin Liu <liumengxinfly@gmail.com>
Date: Thu, 12 Sep 2024 07:36:46 +0000
Subject: [PATCH] virtual key

Signed-off-by: Mengxin Liu <liumengxinfly@gmail.com>
---
 src/index.ts | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/src/index.ts b/src/index.ts
index cabebcb..6b1c7b7 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -30,6 +30,12 @@ async function handleChat(c: Context) {
   const queryParams = new URLSearchParams(c.req.query()).toString()
   const urlWithQueryParams = `${azureEndpoint}?${queryParams}`
 
+  const apiKey = c.req.header('api-key')
+  const realKey = await c.env.MALACCA_USER.get(apiKey)
+  if (!realKey) {
+    return c.text('Unauthorized', 401)
+  }
+
 
   async function generateCacheKey(urlWithQueryParams: string, body: any): Promise<string> {
     const cacheKey = await crypto.subtle.digest(
@@ -54,10 +60,13 @@ async function handleChat(c: Context) {
     return new Response(responseFromCache)
   }
 
+  const headers = new Headers(c.req.header())
+  headers.set('api-key', realKey)
+
   const response = await fetch(urlWithQueryParams, {
     method: c.req.method,
     body: JSON.stringify(body),
-    headers: c.req.header()
+    headers: headers
   })
 
   const { readable, writable } = new TransformStream()