From 873f9cf6f8a0a77954cf7320da5eac2c2cdd7bc7 Mon Sep 17 00:00:00 2001
From: kwaroran <kwaworan@gmail.com>
Date: Sat, 22 Jul 2023 01:04:29 +0900
Subject: [PATCH] [feat] added tokens for reverse proxy

---
 src/ts/process/request.ts   |  5 +++--
 src/ts/storage/globalApi.ts | 41 ++++++++++++++++++++++++++-----------
 2 files changed, 32 insertions(+), 14 deletions(-)

diff --git a/src/ts/process/request.ts b/src/ts/process/request.ts
index 1e466158..d41e9f40 100644
--- a/src/ts/process/request.ts
+++ b/src/ts/process/request.ts
@@ -201,7 +201,8 @@ export async function requestChatDataMain(arg:requestDataArgument, model:'model'
                         headers: {
                             "risu-header": encodeURIComponent(JSON.stringify(headers)),
                             "risu-url": encodeURIComponent(replacerURL),
-                            "Content-Type": "application/json"
+                            "Content-Type": "application/json",
+                            "x-risu-tk": "use"
                         },
                         method: "POST",
                         signal: abortSignal
@@ -268,7 +269,7 @@ export async function requestChatDataMain(arg:requestDataArgument, model:'model'
             const res = await globalFetch(replacerURL, {
                 body: body,
                 headers: headers,
-                abortSignal
+                abortSignal,
             })
 
             const dat = res.data as any
diff --git a/src/ts/storage/globalApi.ts b/src/ts/storage/globalApi.ts
index 5cd358df..ef55d3a3 100644
--- a/src/ts/storage/globalApi.ts
+++ b/src/ts/storage/globalApi.ts
@@ -437,10 +437,18 @@ export async function loadData() {
 
 const knownHostes = ["localhost","127.0.0.1"]
 
-export async function globalFetch(url:string, arg:{plainFetchForce?:boolean,body?:any,headers?:{[key:string]:string}, rawResponse?:boolean, method?:"POST"|"GET", abortSignal?:AbortSignal} = {}): Promise<{
+export async function globalFetch(url:string, arg:{
+    plainFetchForce?:boolean,
+    body?:any,
+    headers?:{[key:string]:string},
+    rawResponse?:boolean,
+    method?:"POST"|"GET",
+    abortSignal?:AbortSignal,
+    useRisuToken?:boolean
+} = {}): Promise<{
     ok: boolean;
     data: any;
-    headers:{[key:string]:string}
+    headers:{[key:string]:string},
 }> {
     try {
         const db = get(DataBase)
@@ -625,13 +633,18 @@ export async function globalFetch(url:string, arg:{plainFetchForce?:boolean,body
                 if(arg.rawResponse){
                     const furl = ((!isTauri) && (!isNodeServer)) ? `${hubURL}/proxy2` : `/proxy2`
                 
+                    let headers = {
+                        "risu-header": encodeURIComponent(JSON.stringify(arg.headers)),
+                        "risu-url": encodeURIComponent(url),
+                        "Content-Type": "application/json",
+                    }
+                    if(arg.useRisuToken){
+                        headers["x-risu-tk"] = "use"
+                    }
+
                     const da = await fetch(furl, {
                         body: body,
-                        headers: {
-                            "risu-header": encodeURIComponent(JSON.stringify(arg.headers)),
-                            "risu-url": encodeURIComponent(url),
-                            "Content-Type": "application/json"
-                        },
+                        headers: headers,
                         method: method,
                         signal: arg.abortSignal
                     })
@@ -646,13 +659,17 @@ export async function globalFetch(url:string, arg:{plainFetchForce?:boolean,body
                 else{
                     const furl = ((!isTauri) && (!isNodeServer)) ? `${hubURL}/proxy2` : `/proxy2`
 
+                    let headers = {
+                        "risu-header": encodeURIComponent(JSON.stringify(arg.headers)),
+                        "risu-url": encodeURIComponent(url),
+                        "Content-Type": "application/json"
+                    }
+                    if(arg.useRisuToken){
+                        headers["x-risu-tk"] = "use"
+                    }
                     const da = await fetch(furl, {
                         body: body,
-                        headers: {
-                            "risu-header": encodeURIComponent(JSON.stringify(arg.headers)),
-                            "risu-url": encodeURIComponent(url),
-                            "Content-Type": "application/json"
-                        },
+                        headers: headers,
                         method: method
                     })
                     const daText = await da.text()