tcjlj/nezha
1
0
Fork 0
mirror of https://github.com/nezhahq/nezha.git synced 2025-02-08 20:48:13 -05:00
Code Issues Actions Packages Projects Releases Wiki Activity

update: 线程安全

Browse Source
This commit is contained in:
Akkia 2022-05-18 10:28:24 +08:00
parent 990394bf46
commit 73df5fa0de
No known key found for this signature in database
GPG Key ID: DABE9A4AB2DD7EF3
3 changed files with 13 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
cmd/dashboard/controller/member_api.go
View File

@ -62,7 +62,9 @@ type apiResult struct {
// getToken 获取 Token // getToken 获取 Token
func (ma *memberAPI) getToken(c *gin.Context) { func (ma *memberAPI) getToken(c *gin.Context) {
u := c.MustGet(model.CtxKeyAuthorizedUser).(*model.User) u := c.MustGet(model.CtxKeyAuthorizedUser).(*model.User)
singleton.ApiLock.RLock()
tokenList := singleton.UserIDToApiTokenList[u.ID] tokenList := singleton.UserIDToApiTokenList[u.ID]
singleton.ApiLock.RUnlock()
res := make([]*apiResult, len(tokenList)) res := make([]*apiResult, len(tokenList))
for i, token := range tokenList { for i, token := range tokenList {
res[i] = &apiResult{ res[i] = &apiResult{
@ -84,8 +86,12 @@ func (ma *memberAPI) issueNewToken(c *gin.Context) {
Token: utils.MD5(fmt.Sprintf("%d%d%s", time.Now().UnixNano(), u.ID, u.Login)), Token: utils.MD5(fmt.Sprintf("%d%d%s", time.Now().UnixNano(), u.ID, u.Login)),
} }
singleton.DB.Create(token) singleton.DB.Create(token)
singleton.ApiLock.Lock()
singleton.ApiTokenList[token.Token] = token singleton.ApiTokenList[token.Token] = token
singleton.UserIDToApiTokenList[u.ID] = append(singleton.UserIDToApiTokenList[u.ID], token.Token) singleton.UserIDToApiTokenList[u.ID] = append(singleton.UserIDToApiTokenList[u.ID], token.Token)
singleton.ApiLock.Unlock()
c.JSON(http.StatusOK, model.Response{ c.JSON(http.StatusOK, model.Response{
Code: http.StatusOK, Code: http.StatusOK,
Message: "success", Message: "success",
@ -105,6 +111,8 @@ func (ma *memberAPI) deleteToken(c *gin.Context) {
}) })
return return
} }
singleton.ApiLock.Lock()
defer singleton.ApiLock.Unlock()
if _, ok := singleton.ApiTokenList[token]; !ok { if _, ok := singleton.ApiTokenList[token]; !ok {
c.JSON(http.StatusOK, model.Response{ c.JSON(http.StatusOK, model.Response{
Code: http.StatusBadRequest, Code: http.StatusBadRequest,
@ -114,6 +122,7 @@ func (ma *memberAPI) deleteToken(c *gin.Context) {
} }
// 在数据库中删除该Token // 在数据库中删除该Token
singleton.DB.Unscoped().Delete(&model.ApiToken{}, "token = ?", token) singleton.DB.Unscoped().Delete(&model.ApiToken{}, "token = ?", token)
// 在UserIDToApiTokenList中删除该Token // 在UserIDToApiTokenList中删除该Token
for i, t := range singleton.UserIDToApiTokenList[singleton.ApiTokenList[token].UserID] { for i, t := range singleton.UserIDToApiTokenList[singleton.ApiTokenList[token].UserID] {
if t == token { if t == token {

2
pkg/mygin/auth.go
View File

@ -55,10 +55,12 @@ func Authorize(opt AuthorizeOption) func(*gin.Context) {
apiToken := c.GetHeader("Authorization") apiToken := c.GetHeader("Authorization")
if apiToken != "" { if apiToken != "" {
var u model.User var u model.User
singleton.ApiLock.RLock()
if _, ok := singleton.ApiTokenList[apiToken]; ok { if _, ok := singleton.ApiTokenList[apiToken]; ok {
err := singleton.DB.First(&u).Where("id = ?", singleton.ApiTokenList[apiToken].UserID).Error err := singleton.DB.First(&u).Where("id = ?", singleton.ApiTokenList[apiToken].UserID).Error
isLogin = err == nil isLogin = err == nil
} }
singleton.ApiLock.RUnlock()
if isLogin { if isLogin {
c.Set(model.CtxKeyAuthorizedUser, &u) c.Set(model.CtxKeyAuthorizedUser, &u)
c.Set("isAPI", true) c.Set("isAPI", true)

2
service/singleton/api.go
View File

@ -3,11 +3,13 @@ package singleton
import ( import (
"github.com/naiba/nezha/model" "github.com/naiba/nezha/model"
"github.com/naiba/nezha/pkg/utils" "github.com/naiba/nezha/pkg/utils"
"sync"
) )
var ( var (
ApiTokenList = make(map[string]*model.ApiToken) ApiTokenList = make(map[string]*model.ApiToken)
UserIDToApiTokenList = make(map[uint64][]string) UserIDToApiTokenList = make(map[uint64][]string)
ApiLock sync.RWMutex
) )
type ServerAPI struct { type ServerAPI struct {