From 9867fc0b19fe50cda3d64f46dedbd6f752c8cdc1 Mon Sep 17 00:00:00 2001
From: zengl <craftleon@hotmail.com>
Date: Mon, 14 Oct 2024 09:52:06 +0800
Subject: [PATCH] fix token mutex issue

---
 ac/msghandler.go     |  5 +++--
 ac/tokenstore.go     | 28 +++++++++++++++-------------
 ac/udpac.go          |  7 +++++--
 server/tokenstore.go | 28 +++++++++++++++-------------
 4 files changed, 38 insertions(+), 30 deletions(-)

diff --git a/ac/msghandler.go b/ac/msghandler.go
index fd8ad425..b515cdeb 100644
--- a/ac/msghandler.go
+++ b/ac/msghandler.go
@@ -22,7 +22,6 @@ const (
 )
 
 func (a *UdpAC) HandleUdpACOperations(ppd *core.PacketParserData) (err error) {
-	a.wg.Add(1)
 	defer a.wg.Done()
 
 	acId := a.config.ACId
@@ -49,7 +48,7 @@ func (a *UdpAC) HandleUdpACOperations(ppd *core.PacketParserData) (err error) {
 	}
 	artMsg, err = a.HandleAccessControl(agentUser, srcAddrs, dstAddrs, openTimeSec, artMsg)
 	if err != nil {
-		log.Error("ac(%s#%d)[HandleUdpACOperations] HandleAccessControl failed, err: %v", acId, err)
+		log.Error("ac(%s#%d)[HandleUdpACOperations] HandleAccessControl failed, err: %v", acId, transactionId, err)
 	}
 
 	// generate ac token and save user and access information
@@ -60,6 +59,7 @@ func (a *UdpAC) HandleUdpACOperations(ppd *core.PacketParserData) (err error) {
 		OpenTime: openTimeSec,
 	}
 	artMsg.ACToken = a.GenerateAccessToken(entry)
+	//log.Info("generate knock token: %s", artMsg.ACToken)
 
 	// send ac result
 	artBytes, _ := json.Marshal(artMsg)
@@ -70,6 +70,7 @@ func (a *UdpAC) HandleUdpACOperations(ppd *core.PacketParserData) (err error) {
 		PrevParserData: ppd,
 		Message:        artBytes,
 	}
+	//log.Info("ART result: %s", string(artBytes))
 
 	// forward to a specific transaction
 	transaction := ppd.ConnData.FindRemoteTransaction(transactionId)
diff --git a/ac/tokenstore.go b/ac/tokenstore.go
index 2f966218..f81e5ca2 100644
--- a/ac/tokenstore.go
+++ b/ac/tokenstore.go
@@ -77,21 +77,23 @@ func (a *UdpAC) tokenStoreRefreshRoutine() {
 			return
 
 		case <-time.After(TokenStoreRefreshInterval * time.Second):
-			a.tokenStoreMutex.Lock()
-			defer a.tokenStoreMutex.Unlock()
-
-			now := time.Now()
-			for head, tokenMap := range a.tokenStore {
-				for token, entry := range tokenMap {
-					if now.After(entry.ExpireTime) {
-						log.Info("[TokenStore] token %s expired, remove", token)
-						delete(tokenMap, token)
+			func() {
+				a.tokenStoreMutex.Lock()
+				defer a.tokenStoreMutex.Unlock()
+
+				now := time.Now()
+				for head, tokenMap := range a.tokenStore {
+					for token, entry := range tokenMap {
+						if now.After(entry.ExpireTime) {
+							log.Info("[TokenStore] token %s expired, remove", token)
+							delete(tokenMap, token)
+						}
+					}
+					if len(tokenMap) == 0 {
+						delete(a.tokenStore, head)
 					}
 				}
-				if len(tokenMap) == 0 {
-					delete(a.tokenStore, head)
-				}
-			}
+			}()
 		}
 	}
 }
diff --git a/ac/udpac.go b/ac/udpac.go
index affbbc83..d9e1e2a6 100644
--- a/ac/udpac.go
+++ b/ac/udpac.go
@@ -65,8 +65,10 @@ type UdpConn struct {
 }
 
 func (c *UdpConn) Close() {
-	c.netConn.Close()
-	c.ConnData.Close()
+	if c.netConn != nil {
+		c.netConn.Close()
+		c.ConnData.Close()
+	}
 }
 
 /*
@@ -438,6 +440,7 @@ func (a *UdpAC) recvMessageRoutine() {
 			switch ppd.HeaderType {
 			case core.NHP_AOP:
 				// deal with NHP_AOP message
+				a.wg.Add(1)
 				go a.HandleUdpACOperations(ppd)
 			}
 		}
diff --git a/server/tokenstore.go b/server/tokenstore.go
index 9320945a..53a64ef9 100644
--- a/server/tokenstore.go
+++ b/server/tokenstore.go
@@ -77,21 +77,23 @@ func (s *UdpServer) tokenStoreRefreshRoutine() {
 			return
 
 		case <-time.After(TokenStoreRefreshInterval * time.Second):
-			s.tokenStoreMutex.Lock()
-			defer s.tokenStoreMutex.Unlock()
-
-			now := time.Now()
-			for head, tokenMap := range s.tokenStore {
-				for token, entry := range tokenMap {
-					if now.After(entry.ExpireTime) {
-						log.Info("[TokenStore] token %s expired, remove", token)
-						delete(tokenMap, token)
+			func() {
+				s.tokenStoreMutex.Lock()
+				defer s.tokenStoreMutex.Unlock()
+
+				now := time.Now()
+				for head, tokenMap := range s.tokenStore {
+					for token, entry := range tokenMap {
+						if now.After(entry.ExpireTime) {
+							log.Info("[TokenStore] token %s expired, remove", token)
+							delete(tokenMap, token)
+						}
+					}
+					if len(tokenMap) == 0 {
+						delete(s.tokenStore, head)
 					}
 				}
-				if len(tokenMap) == 0 {
-					delete(s.tokenStore, head)
-				}
-			}
+			}()
 		}
 	}
 }