From 6a8a4bcf65a87577d0919af659161f5a1a4d2d4e Mon Sep 17 00:00:00 2001
From: "1808837298@qq.com" <1808837298@qq.com>
Date: Thu, 10 Oct 2024 13:39:09 +0800
Subject: [PATCH] fix: playground group

---
 common/user_groups.go | 5 +++++
 controller/relay.go   | 8 +++++++-
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/common/user_groups.go b/common/user_groups.go
index 67c3e7912..e705394f1 100644
--- a/common/user_groups.go
+++ b/common/user_groups.go
@@ -21,3 +21,8 @@ func UpdateUserUsableGroupsByJSONString(jsonStr string) error {
 	UserUsableGroups = make(map[string]string)
 	return json.Unmarshal([]byte(jsonStr), &UserUsableGroups)
 }
+
+func GroupInUserUsableGroups(groupName string) bool {
+	_, ok := UserUsableGroups[groupName]
+	return ok
+}
diff --git a/controller/relay.go b/controller/relay.go
index ef0423286..4a49d51a4 100644
--- a/controller/relay.go
+++ b/controller/relay.go
@@ -68,9 +68,15 @@ func Playground(c *gin.Context) {
 	}
 	c.Set("original_model", playgroundRequest.Model)
 	group := playgroundRequest.Group
+	userGroup := c.GetString("group")
+
 	if group == "" {
-		group = c.GetString("group")
+		group = userGroup
 	} else {
+		if !common.GroupInUserUsableGroups(group) && group != userGroup {
+			openaiErr = service.OpenAIErrorWrapperLocal(errors.New("无权访问该分组"), "group_not_allowed", http.StatusForbidden)
+			return
+		}
 		c.Set("group", group)
 	}
 	c.Set("token_name", "playground-"+group)