From cefc0edfa0bdb867fb91f832dbb3206f975bbb47 Mon Sep 17 00:00:00 2001
From: PandaX <18610165312@163.com>
Date: Thu, 26 Oct 2023 10:13:17 +0800
Subject: [PATCH] =?UTF-8?q?[fix]udp=E8=AE=A4=E8=AF=81=E9=80=BB=E8=BE=91?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 iothub/server/udpserver/hook.go | 44 ++++++++++++++++-----------------
 1 file changed, 22 insertions(+), 22 deletions(-)

diff --git a/iothub/server/udpserver/hook.go b/iothub/server/udpserver/hook.go
index 87fb57d..85af874 100644
--- a/iothub/server/udpserver/hook.go
+++ b/iothub/server/udpserver/hook.go
@@ -29,7 +29,7 @@ func InitUdpHook(addr string, hs *hook_message_work.HookService) {
 		global.Log.Infof("UDP IOTHUB HOOK Start SUCCESS, Server listen: %s", addr)
 	}
 	buffer := make([]byte, 1024)
-	isAuth := false
+	authMap := make(map[string]bool)
 	etoken := &global_model.DeviceAuth{}
 	hhs := &HookUdpService{
 		HookService: hs,
@@ -41,33 +41,15 @@ func InitUdpHook(addr string, hs *hook_message_work.HookService) {
 			global.Log.Error("Error accepting connection:", err)
 			_ = server.listener.Close()
 			//设置断开连接
-			if isAuth {
+			if isAuth, ok := authMap[client.AddrPort().String()]; ok && isAuth {
 				data := netbase.CreateConnectionInfo(message.DisConnectMes, "udp", client.IP.String(), client.AddrPort().String(), etoken)
 				hhs.HookService.MessageCh <- data
 			}
 			delete(udpclient.UdpClient, etoken.DeviceId)
-			isAuth = false
+			delete(authMap, client.AddrPort().String())
 			continue
 		}
-		if !isAuth {
-			token := string(buffer[:n])
-			etoken.GetDeviceToken(token)
-			auth := netbase.Auth(token)
-			// 认证成功，创建连接记录
-			if auth {
-				global.Log.Infof("UDP协议 设备%s,认证成功", etoken.DeviceId)
-				data := netbase.CreateConnectionInfo(message.ConnectMes, "udp", client.IP.String(), client.AddrPort().String(), etoken)
-				hhs.HookService.MessageCh <- data
-				isAuth = true
-				udpclient.UdpClient[etoken.DeviceId] = &udpclient.UdpClientT{
-					Conn: server.listener,
-					Addr: client,
-				}
-				hhs.Send(client, "success")
-			} else {
-				hhs.Send(client, "fail")
-			}
-		} else {
+		if isAuth, ok := authMap[client.AddrPort().String()]; ok && isAuth {
 			data := &netbase.DeviceEventInfo{
 				DeviceId:   etoken.DeviceId,
 				DeviceAuth: etoken,
@@ -81,6 +63,24 @@ func InitUdpHook(addr string, hs *hook_message_work.HookService) {
 
 			// etoken中添加设备标识
 			hhs.HookService.MessageCh <- data
+		} else {
+			token := string(buffer[:n])
+			etoken.GetDeviceToken(token)
+			auth := netbase.Auth(token)
+			// 认证成功，创建连接记录
+			if auth {
+				global.Log.Infof("UDP协议 设备%s,认证成功", etoken.DeviceId)
+				data := netbase.CreateConnectionInfo(message.ConnectMes, "udp", client.IP.String(), client.AddrPort().String(), etoken)
+				hhs.HookService.MessageCh <- data
+				authMap[client.AddrPort().String()] = true
+				udpclient.UdpClient[etoken.DeviceId] = &udpclient.UdpClientT{
+					Conn: server.listener,
+					Addr: client,
+				}
+				hhs.Send(client, "success")
+			} else {
+				hhs.Send(client, "fail")
+			}
 		}
 	}
 }