diff --git a/server/chat_handler.go b/server/chat_handler.go index c0ab1ca2..48673223 100644 --- a/server/chat_handler.go +++ b/server/chat_handler.go @@ -4,6 +4,7 @@ import ( "bufio" "bytes" "encoding/json" + "errors" "fmt" "github.com/gin-gonic/gin" "github.com/gorilla/websocket" @@ -76,7 +77,14 @@ func (s *Server) sendMessage(session types.ChatSession, role types.ChatRole, pro if user.Status == false { replyMessage(ws, "当前 TOKEN 已经被禁用，如果疑问，请联系管理员！", false) replyMessage(ws, "![](images/wx.png)", true) - return err + return errors.New("当前 TOKEN " + user.Name + "已经被禁用") + } + + if time.Now().Unix() > user.ExpiredTime { + exTime := time.Unix(user.ExpiredTime, 0).Format("2006-01-02 15:04:05") + replyMessage(ws, "当前 TOKEN 已过期，过期时间为："+exTime+"，如果疑问，请联系管理员！", false) + replyMessage(ws, "![](images/wx.png)", true) + return errors.New("当前 TOKEN " + user.Name + "已过期") } if user.MaxCalls > 0 && user.RemainingCalls <= 0 { diff --git a/server/config_handler.go b/server/config_handler.go index f38ab662..f1cf3e0f 100644 --- a/server/config_handler.go +++ b/server/config_handler.go @@ -81,7 +81,7 @@ func (s *Server) SetDebugHandle(c *gin.Context) { } err := json.NewDecoder(c.Request.Body).Decode(&data) if err != nil { - c.JSON(http.StatusOK, types.BizVo{Code: types.Failed, Message: "Invalid args"}) + c.JSON(http.StatusOK, types.BizVo{Code: types.InvalidParams, Message: "Invalid args"}) return } @@ -91,16 +91,21 @@ func (s *Server) SetDebugHandle(c *gin.Context) { // AddUserHandle 添加 Username func (s *Server) AddUserHandle(c *gin.Context) { - var data types.User + var data struct { + Name string `json:"name"` + MaxCalls int `json:"max_calls"` + EnableHistory bool `json:"enable_history"` + Term int `json:"term"` // 有效期 + } err := json.NewDecoder(c.Request.Body).Decode(&data) if err != nil { - c.JSON(http.StatusOK, types.BizVo{Code: types.Failed, Message: "Invalid args"}) + c.JSON(http.StatusOK, types.BizVo{Code: types.InvalidParams, Message: "Invalid args"}) return } // 参数处理 if data.Name == "" || data.MaxCalls < 0 { - c.JSON(http.StatusOK, types.BizVo{Code: types.Failed, Message: "Invalid args"}) + c.JSON(http.StatusOK, types.BizVo{Code: types.InvalidParams, Message: "Invalid args"}) return } @@ -111,7 +116,13 @@ func (s *Server) AddUserHandle(c *gin.Context) { return } - user := types.User{Name: data.Name, MaxCalls: data.MaxCalls, RemainingCalls: data.MaxCalls, EnableHistory: data.EnableHistory, Status: true} + user := types.User{ + Name: data.Name, + MaxCalls: data.MaxCalls, + RemainingCalls: data.MaxCalls, + EnableHistory: data.EnableHistory, + Term: data.Term, + Status: true} err = PutUser(user) if err != nil { c.JSON(http.StatusOK, types.BizVo{Code: types.Failed, Message: "Failed to save configs"}) @@ -127,23 +138,31 @@ func (s *Server) BatchAddUserHandle(c *gin.Context) { Number int `json:"number"` MaxCalls int `json:"max_calls"` EnableHistory bool `json:"enable_history"` + Term int `json:"term"` } err := json.NewDecoder(c.Request.Body).Decode(&data) if err != nil || data.MaxCalls <= 0 { - c.JSON(http.StatusOK, types.BizVo{Code: types.Failed, Message: "Invalid args"}) + c.JSON(http.StatusOK, types.BizVo{Code: types.InvalidParams, Message: "Invalid args"}) return } - var users = make([]string, 0) + var users = make([]types.User, 0) for i := 0; i < data.Number; i++ { name := utils.RandString(12) _, err := GetUser(name) for err == nil { name = utils.RandString(12) } - err = PutUser(types.User{Name: name, MaxCalls: data.MaxCalls, RemainingCalls: data.MaxCalls, EnableHistory: data.EnableHistory, Status: true}) + user := types.User{ + Name: name, + MaxCalls: data.MaxCalls, + RemainingCalls: data.MaxCalls, + EnableHistory: data.EnableHistory, + Term: data.Term, + Status: true} + err = PutUser(user) if err == nil { - users = append(users, name) + users = append(users, user) } } @@ -154,7 +173,7 @@ func (s *Server) SetUserHandle(c *gin.Context) { var data map[string]interface{} err := json.NewDecoder(c.Request.Body).Decode(&data) if err != nil { - c.JSON(http.StatusOK, types.BizVo{Code: types.Failed, Message: "Invalid args"}) + c.JSON(http.StatusOK, types.BizVo{Code: types.InvalidParams, Message: "Invalid args"}) return } @@ -189,6 +208,9 @@ func (s *Server) SetUserHandle(c *gin.Context) { if v, ok := data["remaining_calls"]; ok { user.RemainingCalls = int(v.(float64)) } + if v, ok := data["expired_time"]; ok { + user.ExpiredTime = int64(v.(float64)) + } if v, ok := data["api_key"]; ok { user.ApiKey = v.(string) } diff --git a/server/server.go b/server/server.go index d9f9f3ae..5e2cd9b4 100644 --- a/server/server.go +++ b/server/server.go @@ -271,7 +271,8 @@ func (s *Server) LoginHandle(c *gin.Context) { c.JSON(http.StatusOK, types.BizVo{Code: types.Failed, Message: types.ErrorMsg}) return } - if !utils.ContainUser(GetUsers(), data.Token) { + user, err := GetUser(data.Token) + if err != nil { c.JSON(http.StatusOK, types.BizVo{Code: types.Failed, Message: "Invalid user"}) return } @@ -285,5 +286,20 @@ func (s *Server) LoginHandle(c *gin.Context) { } // 记录客户端 IP 地址 s.ChatSession[sessionId] = types.ChatSession{ClientIP: c.ClientIP(), Username: data.Token, SessionId: sessionId} + // 更新用户激活时间 + user.ActiveTime = time.Now().Unix() + if user.ExpiredTime == 0 { + activeTime := time.Unix(user.ActiveTime, 0) + if user.Term == 0 { + user.Term = 30 // 默认 30 天到期 + } + user.ExpiredTime = activeTime.Add(time.Hour * 24 * time.Duration(user.Term)).Unix() + } + err = PutUser(*user) + if err != nil { + c.JSON(http.StatusOK, types.BizVo{Code: types.Failed, Message: "Save user info failed"}) + return + } + c.JSON(http.StatusOK, types.BizVo{Code: types.Success, Data: sessionId}) } diff --git a/test/test.go b/test/test.go index b8042e94..a3ba164a 100644 --- a/test/test.go +++ b/test/test.go @@ -2,26 +2,9 @@ package main import ( "fmt" - "openai/utils" + "time" ) func main() { - // leveldb 测试 - db, err := utils.NewLevelDB("data") - if err != nil { - panic(err) - } - defer db.Close() - - err = db.Put("name", "xiaoming") - if err != nil { - panic(err) - } - - name, err := db.Get("name") - if err != nil { - panic(err) - } - - fmt.Println("name: ", name) + fmt.Println(time.Unix(1683336167, 0).Format("2006-01-02 15:04:05")) } diff --git a/types/config.go b/types/config.go index 1fddfb60..ff3a32fe 100644 --- a/types/config.go +++ b/types/config.go @@ -15,11 +15,14 @@ type Config struct { type User struct { Name string `json:"name"` - MaxCalls int `json:"max_calls"` // 最多调用次数，如果为 0 则表示不限制 - RemainingCalls int `json:"remaining_calls"` // 剩余调用次数 - EnableHistory bool `json:"enable_history"` // 是否启用聊天记录 - Status bool `json:"status"` // 当前状态 - ApiKey string `json:"api_key"` // OpenAI API KEY + MaxCalls int `json:"max_calls"` // 最多调用次数，如果为 0 则表示不限制 + RemainingCalls int `json:"remaining_calls"` // 剩余调用次数 + EnableHistory bool `json:"enable_history"` // 是否启用聊天记录 + Status bool `json:"status"` // 当前状态 + Term int `json:"term" default:"30"` // 会员有效期，单位：天 + ActiveTime int64 `json:"active_time"` // 激活时间 + ExpiredTime int64 `json:"expired_time"` // 到期时间 + ApiKey string `json:"api_key"` // OpenAI API KEY } // Chat configs struct diff --git a/utils/utils.go b/utils/utils.go index 62b2a6be..cf4d315c 100644 --- a/utils/utils.go +++ b/utils/utils.go @@ -2,7 +2,6 @@ package utils import ( "math/rand" - "openai/types" "strconv" "strings" "time" @@ -40,12 +39,3 @@ func ContainsStr(slice []string, item string) bool { } return false } - -func ContainUser(slice []types.User, user string) bool { - for _, e := range slice { - if e.Name == user { - return true - } - } - return false -}