feat: 添加头部 secret key

Author: DavidKk

.env.example

@@ -0,0 +1 @@
+SECRET_KEY=

README.md

@@ -28,7 +28,15 @@ Due to network conditions and geographical locations, access to the Gemini API c
 ```bash
 $ curl "http://$YOU_SERVER_HOST:$PORT/api/v1beta/models/gemini-pro:streamGenerateContent?key=$GEMINI_API_TOKEN" \
   -H "Content-Type: application/json" \
+  -H "x-gemini-proxy-secret: $SECRET_KEY" \
   -H 'cache-control: no-cache' \
   --data-raw '{"contents":[{"role":"user","parts":[{"text":"Hello Gemini"}]}]}' \
   --compressed
 ```
+
+### Parameters
+
+**SECRET_KEY**: User Limitation
+
+- Generate SECRET_KEY using `pnpm secret` and write it to the environment variables.
+- If not set, there will be no user limitation.

README.zh-CN.md

@@ -28,7 +28,15 @@ Vercel Edge 上的 Gemini API 代理是一项代理服务，旨在解决某些
 ```bash
 $ curl "http://$YOU_SERVER_HOST:$PORT/v1/models/gemini-pro:generateContent?key=$GEMINI_API_TOKEN" \
   -H "Content-Type: application/json" \
+  -H "x-gemini-proxy-secret: $SECRET_KEY" \
   -H 'cache-control: no-cache' \
   --data-raw '{"contents":[{"role":"user","parts":[{"text":"你好 Gemini"}]}]}' \
   --compressed
 ```
+
+### 参数
+
+**SECRET_KEY**: 用户限制
+
+- 使用 `pnpm secret` 生成 SECRET_KEY，并将其写入环境变量。
+- 如果不设置，则不会有用户限制。

__tests__/handleRequest.spec.ts

@@ -15,7 +15,7 @@ describe('test handleRequest', () => {
       ],
     }
 
-    const request = new Request('https://example.com/v1/projects/my-project/locations/us-central1/agents/my-agent/sessions/1234567890:detectIntent?key=123', {
+    const request = new Request('https://example.com/api/v1/projects/my-project/locations/us-central1/agents/my-agent/sessions/1234567890:detectIntent?key=123', {
       method: 'POST',
       headers: {
         'Content-Type': 'application/json',

__tests__/utils/pickHeaders.spec.ts

@@ -5,7 +5,7 @@ describe('test utils/pickHeaders', () => {
     const headers = new Headers()
     const keys = ['content-type', 'content-length']
     const picked = pickHeaders(headers, keys)
-    expect(picked).toEqual(new Headers())
+    expect(picked).toEqual({})
   })
 
   it('should return a Headers object with the specified headers', () => {
@@ -15,12 +15,10 @@ describe('test utils/pickHeaders', () => {
     headers.set('x-custom-header', 'value')
     const keys = ['content-type', 'content-length']
     const picked = pickHeaders(headers, keys)
-    expect(picked).toEqual(
-      new Headers({
-        'content-type': 'application/json',
-        'content-length': '100',
-      })
-    )
+    expect(picked).toEqual({
+      'content-type': 'application/json',
+      'content-length': '100',
+    })
   })
 
   it('should return a Headers object with the specified headers using a RegExp key', () => {
@@ -30,12 +28,10 @@ describe('test utils/pickHeaders', () => {
     headers.set('x-custom-header', 'value')
     const keys = [/^content-/]
     const picked = pickHeaders(headers, keys)
-    expect(picked).toEqual(
-      new Headers({
-        'content-type': 'application/json',
-        'content-length': '100',
-      })
-    )
+    expect(picked).toEqual({
+      'content-type': 'application/json',
+      'content-length': '100',
+    })
   })
 
   it('should return a Headers object with the specified headers, including non-string values', () => {
@@ -45,11 +41,9 @@ describe('test utils/pickHeaders', () => {
     headers.set('x-custom-header', 'value')
     const keys = ['content-type', 'content-length']
     const picked = pickHeaders(headers, keys)
-    expect(picked).toEqual(
-      new Headers({
-        'content-type': 'application/json',
-        'content-length': '100',
-      })
-    )
+    expect(picked).toEqual({
+      'content-type': 'application/json',
+      'content-length': '100',
+    })
   })
 })

package.json

@@ -8,6 +8,7 @@
     "ci:coverage": "pnpm test",
     "postinstall": "node --experimental-modules ./scripts/postinstall.mjs",
     "preinstall": "node --experimental-modules ./scripts/preinstall.mjs",
+    "secret": "node ./scripts/secret.mjs",
     "clear": "rimraf --glob ./libs",
     "format": "prettier --config .prettierrc.js --write \"**/*.{js,jsx,ts,tsx,d.ts,vue,md,json,yml,yaml}\"",
     "lint": "pnpm lint:ts",
@@ -20,12 +21,12 @@
     "ok": "pnpm format && pnpm lint && pnpm build && pnpm test"
   },
   "dependencies": {
+    "dotenv": "^16.4.5",
     "next": "canary",
     "react": "canary",
     "react-dom": "canary"
   },
   "devDependencies": {
-    "json5": "^2.2.3",
     "@commitlint/cli": "^19.3.0",
     "@commitlint/config-conventional": "^19.2.2",
     "@jest/types": "^29.6.3",
@@ -48,6 +49,7 @@
     "jest": "^29.7.0",
     "jest-junit": "^16.0.0",
     "jest-runner-tsd": "^6.0.0",
+    "json5": "^2.2.3",
     "lint-staged": "^15.2.7",
     "prettier": "^3.3.3",
     "pretty-error": "^4.0.0",

pnpm-lock.yaml

@@ -0,0 +1,5 @@
+import crypto from 'crypto'
+
+const secretKey = crypto.randomBytes(16).toString('hex')
+// eslint-disable-next-line no-console
+console.log('Secret Key:', secretKey)

scripts/secret.mjs

@@ -4,6 +4,13 @@ import { handleRequest } from './handleRequest'
 import { createErrorResponse } from './libs/response'
 
 export default async function proxy(request: Request & { nextUrl?: URL }) {
+  if (process.env.GEMINI_PROXY_SECRET) {
+    const secret = request.headers.get('x-gemini-proxy-secret')
+    if (secret !== process.env.GEMINI_PROXY_SECRET) {
+      return createErrorResponse('Unauthorized', 401, { statusText: 'Unauthorized' })
+    }
+  }
+
   const context = createContext(request)
   const { logger } = context