fix token mutex issue

OpenNHP · Oct 14, 2024 · 9867fc0 · 9867fc0
1 parent 1e8fad9
commit 9867fc0
Show file tree

Hide file tree

Showing 4 changed files with 38 additions and 30 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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)
-				}
-			}
+			}()
 		}
 	}
 }