fix: playground group

Calcium-Ion · Oct 10, 2024 · 6a8a4bc · 6a8a4bc
1 parent e298f2e
commit 6a8a4bc
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git 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
@@ -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)