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

---
 Tun/Punchnet/Actors/SDLUDPHoleActor.swift |  2 +-
 Tun/Punchnet/SDLContext.swift             | 53 ++++++++++++++++-------
 Tun/Punchnet/SDLMessage.swift             |  2 -
 3 files changed, 39 insertions(+), 18 deletions(-)

diff --git a/Tun/Punchnet/Actors/SDLUDPHoleActor.swift b/Tun/Punchnet/Actors/SDLUDPHoleActor.swift
index e09b6a7..f6106aa 100644
--- a/Tun/Punchnet/Actors/SDLUDPHoleActor.swift
+++ b/Tun/Punchnet/Actors/SDLUDPHoleActor.swift
@@ -8,6 +8,7 @@
 import Foundation
 import NIOCore
 import NIOPosix
+import SwiftProtobuf
 
 // 处理和sn-server服务器之间的通讯
 actor SDLUDPHoleActor {
@@ -204,7 +205,6 @@ actor SDLUDPHoleActor {
             }
             
             return .peerInfo(peerInfo)
-            
         case .event:
             guard let eventVal = buffer.readInteger(as: UInt8.self),
                   let event = SDLEventType(rawValue: eventVal),
diff --git a/Tun/Punchnet/SDLContext.swift b/Tun/Punchnet/SDLContext.swift
index 1e75135..6a69b73 100644
--- a/Tun/Punchnet/SDLContext.swift
+++ b/Tun/Punchnet/SDLContext.swift
@@ -149,6 +149,7 @@ public class SDLContext {
     
     private func startUDPHole() async throws {
         self.udpHoleActor = try await SDLUDPHoleActor(logger: self.logger)
+        
         try await withThrowingTaskGroup(of: Void.self) { group in
             group.addTask {
                 try await self.udpHoleActor?.start()
@@ -160,20 +161,19 @@ public class SDLContext {
                     try await Task.sleep(nanoseconds: 5 * 1_000_000_000)
                     try Task.checkCancellation()
 
-                    // TODO
-//                    if let udpHoleActor = self.udpHoleActor {
-//                        let cookie = await udpHoleActor.getCookieId()
-//                        var stunRequest = SDLStunRequest()
-//                        stunRequest.clientID = self.config.clientId
-//                        stunRequest.networkID = self.devAddr.networkID
-//                        stunRequest.ip = self.devAddr.netAddr
-//                        stunRequest.mac = self.devAddr.mac
-//                        stunRequest.natType = UInt32(self.natType.rawValue)
-//                        
-//                        let remoteAddress = self.config.stunSocketAddress
-//                        await udpHoleActor.send(type: .stunRequest, data: try stunRequest.serializedData(), remoteAddress: remoteAddress)
-//                        self.lastCookie = cookie
-//                    }
+                    if let udpHoleActor = self.udpHoleActor {
+                        let cookie = await udpHoleActor.getCookieId()
+                        var stunRequest = SDLStunRequest()
+                        stunRequest.clientID = self.config.clientId
+                        stunRequest.networkID = self.config.networkAddress.networkId
+                        stunRequest.ip = self.config.networkAddress.ip
+                        stunRequest.mac = self.config.networkAddress.mac
+                        stunRequest.natType = UInt32(self.natType.rawValue)
+                        
+                        let remoteAddress = self.config.stunSocketAddress
+                        await udpHoleActor.send(type: .stunRequest, data: try stunRequest.serializedData(), remoteAddress: remoteAddress)
+                        self.lastCookie = cookie
+                    }
                 }
             }
             
@@ -181,6 +181,28 @@ public class SDLContext {
                 if let messageStream = self.udpHoleActor?.messageStream {
                     for try await (remoteAddress, message) in messageStream {
                         try Task.checkCancellation()
+                        switch message {
+                        case .registerSuperAck(let registerSuperAck):
+                            await self.handleRegisterSuperAck(registerSuperAck: registerSuperAck)
+                        case .registerSuperNak(let registerSuperNak):
+                            await self.handleRegisterSuperNak(nakPacket: registerSuperNak)
+                        case .peerInfo(let peerInfo):
+                            ()
+                        case .event(let event):
+                            try await self.handleEvent(event: event)
+                        case .stunReply(let stunReply):
+                            await self.handleStunReply(stunReply: stunReply)
+                        case .stunProbeReply(_):
+                            ()
+                        case .data(let data):
+                            try await self.handleData(data: data)
+                        case .register(let register):
+                            try await self.handleRegister(remoteAddress: remoteAddress, register: register)
+                        case .registerAck(let registerAck):
+                            await self.handleRegisterAck(remoteAddress: remoteAddress, registerAck: registerAck)
+                        default:
+                            self.logger.log("get unknown message: \(message)", level: .error)
+                        }
                     }
                 }
             }
@@ -254,6 +276,7 @@ public class SDLContext {
 ////                return
 ////            }
     ///
+    
     private func handleRegisterSuperAck(registerSuperAck: SDLRegisterSuperAck) async {
         // 需要对数据通过rsa的私钥解码
         let aesKey = try! self.rsaCipher.decode(data: Data(registerSuperAck.aesKey))
@@ -345,7 +368,7 @@ public class SDLContext {
         }
     }
     
-    private func handleRegisterAck(remoteAddress: SocketAddress, registerAck: SDLRegister) async {
+    private func handleRegisterAck(remoteAddress: SocketAddress, registerAck: SDLRegisterAck) async {
         // 判断目标地址是否是tun的网卡地址, 并且是在同一个网络下
         let networkAddr = config.networkAddress
         if registerAck.dstMac == networkAddr.mac && registerAck.networkID == networkAddr.networkId {
diff --git a/Tun/Punchnet/SDLMessage.swift b/Tun/Punchnet/SDLMessage.swift
index 811e664..0a67f0c 100644
--- a/Tun/Punchnet/SDLMessage.swift
+++ b/Tun/Punchnet/SDLMessage.swift
@@ -104,8 +104,6 @@ enum SDLHoleInboundMessage {
     case registerSuperNak(SDLRegisterSuperNak)
     
     case peerInfo(SDLPeerInfo)
-    case pong
-    
     case event(SDLEvent)
     
     case stunReply(SDLStunReply)