From 73df5fa0debcda1381eec719eb51f4b7b4a34768 Mon Sep 17 00:00:00 2001
From: Akkia <AkkiaS7@Outlook.com>
Date: Wed, 18 May 2022 10:28:24 +0800
Subject: [PATCH] =?UTF-8?q?update:=20=E7=BA=BF=E7=A8=8B=E5=AE=89=E5=85=A8?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 cmd/dashboard/controller/member_api.go | 9 +++++++++
 pkg/mygin/auth.go                      | 2 ++
 service/singleton/api.go               | 2 ++
 3 files changed, 13 insertions(+)

diff --git a/cmd/dashboard/controller/member_api.go b/cmd/dashboard/controller/member_api.go
index 9adc40f..9e73755 100644
--- a/cmd/dashboard/controller/member_api.go
+++ b/cmd/dashboard/controller/member_api.go
@@ -62,7 +62,9 @@ type apiResult struct {
 // getToken 获取 Token
 func (ma *memberAPI) getToken(c *gin.Context) {
 	u := c.MustGet(model.CtxKeyAuthorizedUser).(*model.User)
+	singleton.ApiLock.RLock()
 	tokenList := singleton.UserIDToApiTokenList[u.ID]
+	singleton.ApiLock.RUnlock()
 	res := make([]*apiResult, len(tokenList))
 	for i, token := range tokenList {
 		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)),
 	}
 	singleton.DB.Create(token)
+
+	singleton.ApiLock.Lock()
 	singleton.ApiTokenList[token.Token] = token
 	singleton.UserIDToApiTokenList[u.ID] = append(singleton.UserIDToApiTokenList[u.ID], token.Token)
+	singleton.ApiLock.Unlock()
+
 	c.JSON(http.StatusOK, model.Response{
 		Code:    http.StatusOK,
 		Message: "success",
@@ -105,6 +111,8 @@ func (ma *memberAPI) deleteToken(c *gin.Context) {
 		})
 		return
 	}
+	singleton.ApiLock.Lock()
+	defer singleton.ApiLock.Unlock()
 	if _, ok := singleton.ApiTokenList[token]; !ok {
 		c.JSON(http.StatusOK, model.Response{
 			Code:    http.StatusBadRequest,
@@ -114,6 +122,7 @@ func (ma *memberAPI) deleteToken(c *gin.Context) {
 	}
 	// 在数据库中删除该Token
 	singleton.DB.Unscoped().Delete(&model.ApiToken{}, "token = ?", token)
+
 	// 在UserIDToApiTokenList中删除该Token
 	for i, t := range singleton.UserIDToApiTokenList[singleton.ApiTokenList[token].UserID] {
 		if t == token {
diff --git a/pkg/mygin/auth.go b/pkg/mygin/auth.go
index dd56d92..0f28974 100644
--- a/pkg/mygin/auth.go
+++ b/pkg/mygin/auth.go
@@ -55,10 +55,12 @@ func Authorize(opt AuthorizeOption) func(*gin.Context) {
 			apiToken := c.GetHeader("Authorization")
 			if apiToken != "" {
 				var u model.User
+				singleton.ApiLock.RLock()
 				if _, ok := singleton.ApiTokenList[apiToken]; ok {
 					err := singleton.DB.First(&u).Where("id = ?", singleton.ApiTokenList[apiToken].UserID).Error
 					isLogin = err == nil
 				}
+				singleton.ApiLock.RUnlock()
 				if isLogin {
 					c.Set(model.CtxKeyAuthorizedUser, &u)
 					c.Set("isAPI", true)
diff --git a/service/singleton/api.go b/service/singleton/api.go
index 9cfbb4e..f6da4aa 100644
--- a/service/singleton/api.go
+++ b/service/singleton/api.go
@@ -3,11 +3,13 @@ package singleton
 import (
 	"github.com/naiba/nezha/model"
 	"github.com/naiba/nezha/pkg/utils"
+	"sync"
 )
 
 var (
 	ApiTokenList         = make(map[string]*model.ApiToken)
 	UserIDToApiTokenList = make(map[uint64][]string)
+	ApiLock              sync.RWMutex
 )
 
 type ServerAPI struct {