From cd4c977b8338b095a4635e343cabc2f3f9dfb4f0 Mon Sep 17 00:00:00 2001
From: anlicheng <244108715@qq.com>
Date: Tue, 27 Jan 2026 23:47:46 +0800
Subject: [PATCH] fix

---
 Tun/Punchnet/Actors/SDLPuncherActor.swift | 73 ++++++++++++++---------
 Tun/Punchnet/SDLContext.swift             |  2 +-
 2 files changed, 46 insertions(+), 29 deletions(-)

diff --git a/Tun/Punchnet/Actors/SDLPuncherActor.swift b/Tun/Punchnet/Actors/SDLPuncherActor.swift
index 117a328..e1b1bdb 100644
--- a/Tun/Punchnet/Actors/SDLPuncherActor.swift
+++ b/Tun/Punchnet/Actors/SDLPuncherActor.swift
@@ -6,6 +6,7 @@
 //
 
 import Foundation
+import NIOCore
 
 actor SDLPuncherActor {
     // dstMac
@@ -14,8 +15,13 @@ actor SDLPuncherActor {
     
     private var udpHoleActor: SDLUDPHoleActor?
     
+    private var pktId: UInt32 = 1
+    // 提交后还没有响应的请求
+    private var pendingRequests: [UInt32:RegisterRequest] = [:]
+    
     // 处理holer
     private var logger: SDLLogger
+    private var querySocketAddress: SocketAddress
     
     struct RegisterRequest {
         let srcMac: Data
@@ -23,7 +29,8 @@ actor SDLPuncherActor {
         let networkId: UInt32
     }
     
-    init(logger: SDLLogger) {
+    init(querySocketAddress: SocketAddress, logger: SDLLogger) {
+        self.querySocketAddress = querySocketAddress
         self.logger = logger
     }
     
@@ -41,11 +48,35 @@ actor SDLPuncherActor {
         // 触发一次打洞
         coolingDown.insert(dstMac)
         
+        let pktId = self.pktId
+        self.pktId &+= 1
+        if self.pktId == 0 {
+            self.pktId = 1
+        }
+        
         Task {
-            await self.tryHole(request: request)
+            await self.tryHole(pktId: pktId, request: request)
             // 启动冷却期
             try? await Task.sleep(for: .seconds(5))
             self.endCooldown(for: dstMac)
+            self.removePendingRequest(for: pktId)
+        }
+    }
+    
+    func handlePeerInfo(peerInfo: SDLPeerInfo) async {
+        if let request = pendingRequests.removeValue(forKey: peerInfo.pktID) {
+            if let remoteAddress = peerInfo.v4Info.socketAddress() {
+                self.logger.log("[SDLContext] hole sock address: \(remoteAddress)", level: .debug)
+                // 发送register包
+                var register = SDLRegister()
+                register.networkID = request.networkId
+                register.srcMac = request.srcMac
+                register.dstMac = request.dstMac
+                
+                await self.udpHoleActor?.send(type: .register, data: try! register.serializedData(), remoteAddress: remoteAddress)
+            } else {
+                self.logger.log("[SDLContext] hole sock address is invalid: \(peerInfo.v4Info)", level: .warning)
+            }
         }
     }
     
@@ -53,32 +84,18 @@ actor SDLPuncherActor {
         self.coolingDown.remove(key)
     }
     
-    private func tryHole(request: RegisterRequest) async {
-//        var queryInfo = SDLQueryInfo()
-//        queryInfo.dstMac = request.dstMac
-//        guard let message = try? await self.superClientActor?.request(type: .queryInfo, data: try queryInfo.serializedData()) else {
-//            return
-//        }
-//        
-//        switch message.packet {
-//        case .empty:
-//            self.logger.log("[SDLContext] hole query_info get empty: \(message)", level: .debug)
-//        case .peerInfo(let peerInfo):
-//            if let remoteAddress = peerInfo.v4Info.socketAddress() {
-//                self.logger.log("[SDLContext] hole sock address: \(remoteAddress)", level: .debug)
-//                // 发送register包
-//                var register = SDLRegister()
-//                register.networkID = request.networkId
-//                register.srcMac = request.srcMac
-//                register.dstMac = request.dstMac
-//                
-//                await self.udpHoleActor?.send(type: .register, data: try! register.serializedData(), remoteAddress: remoteAddress)
-//            } else {
-//                self.logger.log("[SDLContext] hole sock address is invalid: \(peerInfo.v4Info)", level: .warning)
-//            }
-//        default:
-//            self.logger.log("[SDLContext] hole query_info is packet: \(message)", level: .warning)
-//        }
+    private func removePendingRequest(for pktId: UInt32) {
+        self.pendingRequests.removeValue(forKey: pktId)
     }
     
+    private func tryHole(pktId: UInt32, request: RegisterRequest) async {
+        var queryInfo = SDLQueryInfo()
+        queryInfo.pktID = pktId
+        queryInfo.dstMac = request.dstMac
+        self.pendingRequests[pktId] = request
+        
+        if let queryData = try? queryInfo.serializedData() {
+            await self.udpHoleActor?.send(type: .queryInfo, data: queryData, remoteAddress: self.querySocketAddress)
+        }
+    }
 }
diff --git a/Tun/Punchnet/SDLContext.swift b/Tun/Punchnet/SDLContext.swift
index 6a69b73..7c5249e 100644
--- a/Tun/Punchnet/SDLContext.swift
+++ b/Tun/Punchnet/SDLContext.swift
@@ -83,7 +83,7 @@ public class SDLContext {
         self.sessionManager = SessionManager()
         self.arpServer = ArpServer(known_macs: [:])
         self.providerActor = SDLTunnelProviderActor(provider: provider, logger: logger)
-        self.puncherActor = SDLPuncherActor(logger: logger)
+        self.puncherActor = SDLPuncherActor(querySocketAddress: config.stunSocketAddress, logger: logger)
     }
     
     public func start() async throws {