From 40cb56761dfcf90deef6fb2e054e4b5a00fdc93b Mon Sep 17 00:00:00 2001
From: Zhang Minghan <zmh13054618081@dingtalk.com>
Date: Wed, 13 Mar 2024 14:25:02 +0800
Subject: [PATCH] feat: support close relay api (#111 #112)

---
 app/src/admin/api/system.ts     |  4 ++-
 app/src/resources/i18n/cn.json  |  2 ++
 app/src/resources/i18n/en.json  |  4 ++-
 app/src/resources/i18n/ja.json  |  4 ++-
 app/src/resources/i18n/ru.json  |  4 ++-
 app/src/routes/admin/System.tsx | 47 ++++++++++++++++++++++++++++-----
 auth/auth.go                    |  2 +-
 auth/controller.go              |  3 ++-
 channel/system.go               |  8 +++---
 globals/variables.go            |  2 ++
 manager/chat_completions.go     |  5 ++++
 manager/images.go               |  5 ++++
 12 files changed, 73 insertions(+), 17 deletions(-)

diff --git a/app/src/admin/api/system.ts b/app/src/admin/api/system.ts
index decf515..0a428f6 100644
--- a/app/src/admin/api/system.ts
+++ b/app/src/admin/api/system.ts
@@ -33,6 +33,7 @@ export type SearchState = {
 
 export type SiteState = {
   close_register: boolean;
+  close_relay: boolean;
   relay_plan: boolean;
   quota: number;
   buy_link: string;
@@ -123,8 +124,9 @@ export const initialSystemState: SystemProps = {
     pwa_manifest: "",
   },
   site: {
-    relay_plan: false,
     close_register: false,
+    close_relay: false,
+    relay_plan: false,
     quota: 0,
     buy_link: "",
     announcement: "",
diff --git a/app/src/resources/i18n/cn.json b/app/src/resources/i18n/cn.json
index 531fae5..6660f9d 100644
--- a/app/src/resources/i18n/cn.json
+++ b/app/src/resources/i18n/cn.json
@@ -706,6 +706,8 @@
       "searchPlaceholder": "DuckDuckGo 接入点 (格式仅需填写 https://example.com)",
       "closeRegistration": "暂停注册",
       "closeRegistrationTip": "暂停注册，关闭后新用户将无法注册",
+      "closeRelay": "关闭中转 API",
+      "closeRelayTip": "关闭中转 API，关闭后中转 API 将无法使用",
       "relayPlan": "订阅配额支持中转 API",
       "relayPlanTip": "订阅配额支持中转 API，开启后中转 API 计费会优先考虑使用用户订阅配额\n（提示：订阅为次数配额，对 Token 计费的模型可能会影响成本）",
       "quota": "用户初始点数",
diff --git a/app/src/resources/i18n/en.json b/app/src/resources/i18n/en.json
index 1d3a4d0..28466c0 100644
--- a/app/src/resources/i18n/en.json
+++ b/app/src/resources/i18n/en.json
@@ -564,7 +564,9 @@
       "searchPlaceholder": "DuckDuckGo Access Point (Format only https://example.com)",
       "image_store": "Picture storage",
       "image_storeTip": "Images generated by the OpenAI channel DALL-E will be stored on the server to prevent invalidation of the images",
-      "image_storeNoBackend": "No backend domain configured, cannot enable image storage"
+      "image_storeNoBackend": "No backend domain configured, cannot enable image storage",
+      "closeRelay": "Turn off Staging API",
+      "closeRelayTip": "Turn off the staging API, the staging API will not be available after turning off"
     },
     "user": "Users",
     "invitation-code": "Invitation Code",
diff --git a/app/src/resources/i18n/ja.json b/app/src/resources/i18n/ja.json
index 1d63d57..b86f3ee 100644
--- a/app/src/resources/i18n/ja.json
+++ b/app/src/resources/i18n/ja.json
@@ -564,7 +564,9 @@
       "searchPlaceholder": "DuckDuckGoアクセスポイント（フォーマットのみhttps://example.com ）",
       "image_store": "画像ストレージ",
       "image_storeTip": "OpenAIチャンネルDALL - Eによって生成された画像は、画像の無効化を防ぐためにサーバーに保存されます",
-      "image_storeNoBackend": "バックエンドドメインが設定されていません。画像ストレージを有効にできません"
+      "image_storeNoBackend": "バックエンドドメインが設定されていません。画像ストレージを有効にできません",
+      "closeRelay": "ステージングAPIをオフにする",
+      "closeRelayTip": "ステージングAPIをオフにすると、オフにするとステージングAPIは使用できなくなります"
     },
     "user": "ユーザー管理",
     "invitation-code": "招待コード",
diff --git a/app/src/resources/i18n/ru.json b/app/src/resources/i18n/ru.json
index b15206b..b1f754a 100644
--- a/app/src/resources/i18n/ru.json
+++ b/app/src/resources/i18n/ru.json
@@ -564,7 +564,9 @@
       "searchPlaceholder": "Точка доступа DuckDuckGo (только в формате https://example.com)",
       "image_store": "Хранение изображений",
       "image_storeTip": "Изображения, сгенерированные каналом OpenAI DALL-E, будут храниться на сервере, чтобы предотвратить недействительность изображений",
-      "image_storeNoBackend": "Нет настроенного внутреннего домена, невозможно включить хранение изображений"
+      "image_storeNoBackend": "Нет настроенного внутреннего домена, невозможно включить хранение изображений",
+      "closeRelay": "Отключить Staging API",
+      "closeRelayTip": "Отключите промежуточный API, промежуточный API будет недоступен после отключения"
     },
     "user": "Управление пользователями",
     "invitation-code": "Код приглашения",
diff --git a/app/src/routes/admin/System.tsx b/app/src/routes/admin/System.tsx
index bfde227..c16e58d 100644
--- a/app/src/routes/admin/System.tsx
+++ b/app/src/routes/admin/System.tsx
@@ -44,7 +44,7 @@ import {
 } from "@/components/ui/dialog.tsx";
 import { DialogTitle } from "@radix-ui/react-dialog";
 import Require from "@/components/Require.tsx";
-import { Loader2, PencilLine, Settings2 } from "lucide-react";
+import { PencilLine, RotateCw, Save, Settings2 } from "lucide-react";
 import { FlexibleTextarea } from "@/components/ui/textarea.tsx";
 import Tips from "@/components/Tips.tsx";
 import { cn } from "@/components/ui/lib/utils.ts";
@@ -497,6 +497,21 @@ function Site({ data, dispatch, onChange }: CompProps<SiteState>) {
           }}
         />
       </ParagraphItem>
+      <ParagraphItem>
+        <Label>
+          {t("admin.system.closeRelay")}
+          <Tips
+            className={`inline-block`}
+            content={t("admin.system.closeRelayTip")}
+          />
+        </Label>
+        <Switch
+          checked={data.close_relay}
+          onCheckedChange={(value) => {
+            dispatch({ type: "update:site.close_relay", value });
+          }}
+        />
+      </ParagraphItem>
       <ParagraphItem>
         <Label>
           {t("admin.system.relayPlan")}
@@ -843,7 +858,7 @@ function System() {
     if (doToast !== false) toastState(toast, t, res, true);
   };
 
-  useEffectAsync(async () => {
+  const doRefresh = async () => {
     setLoading(true);
     const res = await getConfig();
     setLoading(false);
@@ -851,18 +866,36 @@ function System() {
     if (res.status) {
       setData({ type: "set", value: res.data });
     }
-  }, []);
+  };
+
+  useEffectAsync(doRefresh, []);
 
   return (
     <div className={`system`}>
       <Card className={`admin-card system-card`}>
         <CardHeader className={`select-none`}>
-          <CardTitle>
-            {t("admin.settings")}
-            {loading && <Loader2 className={`h-4 w-4 ml-2 inline-block`} />}
-          </CardTitle>
+          <CardTitle>{t("admin.settings")}</CardTitle>
         </CardHeader>
         <CardContent className={`flex flex-col gap-1`}>
+          <div className={`system-actions flex flex-row`}>
+            <div className={`grow`} />
+            <Button
+              size={`icon`}
+              variant={`outline`}
+              loading={true}
+              className={`mr-2`}
+              onClick={async () => await doRefresh()}
+            >
+              <RotateCw className={cn(loading && `animate-spin`, `h-4 w-4`)} />
+            </Button>
+            <Button
+              size={`icon`}
+              loading={true}
+              onClick={async () => await doSaving()}
+            >
+              <Save className={`h-4 w-4`} />
+            </Button>
+          </div>
           <General
             form={data}
             data={data.general}
diff --git a/auth/auth.go b/auth/auth.go
index 19f61d9..79e8a95 100644
--- a/auth/auth.go
+++ b/auth/auth.go
@@ -196,7 +196,7 @@ func DeepLogin(c *gin.Context, token string) (string, error) {
 
 	db := utils.GetDBFromContext(c)
 	if !IsUserExist(db, user.Username) {
-		if channel.SystemInstance.IsCloseRegister() {
+		if globals.CloseRegistration {
 			return "", errors.New("this site is not open for registration")
 		}
 
diff --git a/auth/controller.go b/auth/controller.go
index 298c3f8..de14af6 100644
--- a/auth/controller.go
+++ b/auth/controller.go
@@ -2,6 +2,7 @@ package auth
 
 import (
 	"chat/channel"
+	"chat/globals"
 	"chat/utils"
 	"github.com/gin-gonic/gin"
 	"net/http"
@@ -143,7 +144,7 @@ func RegisterAPI(c *gin.Context) {
 		return
 	}
 
-	if channel.SystemInstance.IsCloseRegister() {
+	if globals.CloseRegistration {
 		c.JSON(http.StatusOK, gin.H{
 			"status": false,
 			"error":  "this site is not open for registration",
diff --git a/channel/system.go b/channel/system.go
index 3ae0d6f..04551b0 100644
--- a/channel/system.go
+++ b/channel/system.go
@@ -35,6 +35,7 @@ type generalState struct {
 
 type siteState struct {
 	CloseRegister bool    `json:"close_register" mapstructure:"closeregister"`
+	CloseRelay    bool    `json:"close_relay" mapstructure:"closerelay"`
 	RelayPlan     bool    `json:"relay_plan" mapstructure:"relayplan"`
 	Quota         float64 `json:"quota" mapstructure:"quota"`
 	BuyLink       string  `json:"buy_link" mapstructure:"buylink"`
@@ -94,6 +95,9 @@ func NewSystemConfig() *SystemConfig {
 func (c *SystemConfig) Load() {
 	globals.NotifyUrl = c.GetBackend()
 
+	globals.CloseRegistration = c.Site.CloseRegister
+	globals.CloseRelay = c.Site.CloseRelay
+
 	globals.ArticlePermissionGroup = c.Common.Article
 	globals.GenerationPermissionGroup = c.Common.Generation
 	globals.CacheAcceptedModels = c.Common.Cache
@@ -282,10 +286,6 @@ func (c *SystemConfig) AcceptImageStore() bool {
 	return c.Common.ImageStore
 }
 
-func (c *SystemConfig) IsCloseRegister() bool {
-	return c.Site.CloseRegister
-}
-
 func (c *SystemConfig) SupportRelayPlan() bool {
 	return c.Site.RelayPlan
 }
diff --git a/globals/variables.go b/globals/variables.go
index b79922a..c0a4722 100644
--- a/globals/variables.go
+++ b/globals/variables.go
@@ -18,6 +18,8 @@ var CacheAcceptedModels []string
 var CacheAcceptedExpire int64
 var CacheAcceptedSize int64
 var AcceptImageStore bool
+var CloseRegistration bool
+var CloseRelay bool
 
 func OriginIsAllowed(uri string) bool {
 	if len(AllowedOrigins) == 0 {
diff --git a/manager/chat_completions.go b/manager/chat_completions.go
index 03cfefd..139fe8e 100644
--- a/manager/chat_completions.go
+++ b/manager/chat_completions.go
@@ -26,6 +26,11 @@ func supportRelayPlan() bool {
 }
 
 func ChatRelayAPI(c *gin.Context) {
+	if globals.CloseRelay {
+		abortWithErrorResponse(c, fmt.Errorf("relay api is denied of access"), "access_denied_error")
+		return
+	}
+
 	username := utils.GetUserFromContext(c)
 	if username == "" {
 		abortWithErrorResponse(c, fmt.Errorf("access denied for invalid api key"), "authentication_error")
diff --git a/manager/images.go b/manager/images.go
index bf1003f..228d21e 100644
--- a/manager/images.go
+++ b/manager/images.go
@@ -15,6 +15,11 @@ import (
 )
 
 func ImagesRelayAPI(c *gin.Context) {
+	if globals.CloseRelay {
+		abortWithErrorResponse(c, fmt.Errorf("relay api is denied of access"), "access_denied_error")
+		return
+	}
+
 	username := utils.GetUserFromContext(c)
 	if username == "" {
 		abortWithErrorResponse(c, fmt.Errorf("access denied for invalid api key"), "authentication_error")