fix Tun/Punchnet/SDLIPV6AssistClient.swift

Tun/Punchnet/Protobuf/SDLMessage.pb.swift

@@ -52,12 +52,12 @@ struct SDLV6Info: @unchecked Sendable {
 }
 
 /// ipv6 assist相关
-struct SDLV6AssistProbe: @unchecked Sendable {
+struct SDLV6AssistProbe: Sendable {
   // SwiftProtobuf.Message conformance is added in an extension below. See the
   // `Message` and `Message+*Additions` files in the SwiftProtobuf library for
   // methods supported on all messages.
 
-  var assistToken: Data = Data()
+  var pktID: UInt32 = 0
 
   var unknownFields = SwiftProtobuf.UnknownStorage()
 
@@ -69,6 +69,8 @@ struct SDLV6AssistProbeReply: Sendable {
   // `Message` and `Message+*Additions` files in the SwiftProtobuf library for
   // methods supported on all messages.
 
+  var pktID: UInt32 = 0
+
   var v6Info: SDLV6Info {
     get {return _v6Info ?? SDLV6Info()}
     set {_v6Info = newValue}
@@ -85,7 +87,7 @@ struct SDLV6AssistProbeReply: Sendable {
   fileprivate var _v6Info: SDLV6Info? = nil
 }
 
-struct SDLWelcome: @unchecked Sendable {
+struct SDLWelcome: Sendable {
   // SwiftProtobuf.Message conformance is added in an extension below. See the
   // `Message` and `Message+*Additions` files in the SwiftProtobuf library for
   // methods supported on all messages.
@@ -111,8 +113,6 @@ struct SDLWelcome: @unchecked Sendable {
   /// Clears the value of `ipv6Assist`. Subsequent reads from it will return its default value.
   mutating func clearIpv6Assist() {self._ipv6Assist = nil}
 
-  var assistToken: Data = Data()
-
   var unknownFields = SwiftProtobuf.UnknownStorage()
 
   init() {}
@@ -699,7 +699,7 @@ extension SDLV6Info: SwiftProtobuf.Message, SwiftProtobuf._MessageImplementation
 extension SDLV6AssistProbe: SwiftProtobuf.Message, SwiftProtobuf._MessageImplementationBase, SwiftProtobuf._ProtoNameProviding {
   static let protoMessageName: String = "SDLV6AssistProbe"
   static let _protobuf_nameMap: SwiftProtobuf._NameMap = [
-    1: .standard(proto: "assist_token"),
+    1: .standard(proto: "pkt_id"),
   ]
 
   mutating func decodeMessage<D: SwiftProtobuf.Decoder>(decoder: inout D) throws {
@@ -708,21 +708,21 @@ extension SDLV6AssistProbe: SwiftProtobuf.Message, SwiftProtobuf._MessageImpleme
       // allocates stack space for every case branch when no optimizations are
       // enabled. https://github.com/apple/swift-protobuf/issues/1034
       switch fieldNumber {
-      case 1: try { try decoder.decodeSingularBytesField(value: &self.assistToken) }()
+      case 1: try { try decoder.decodeSingularUInt32Field(value: &self.pktID) }()
       default: break
       }
     }
   }
 
   func traverse<V: SwiftProtobuf.Visitor>(visitor: inout V) throws {
-    if !self.assistToken.isEmpty {
+    if self.pktID != 0 {
-      try visitor.visitSingularBytesField(value: self.assistToken, fieldNumber: 1)
+      try visitor.visitSingularUInt32Field(value: self.pktID, fieldNumber: 1)
     }
     try unknownFields.traverse(visitor: &visitor)
   }
 
   static func ==(lhs: SDLV6AssistProbe, rhs: SDLV6AssistProbe) -> Bool {
-    if lhs.assistToken != rhs.assistToken {return false}
+    if lhs.pktID != rhs.pktID {return false}
     if lhs.unknownFields != rhs.unknownFields {return false}
     return true
   }
@@ -731,7 +731,8 @@ extension SDLV6AssistProbe: SwiftProtobuf.Message, SwiftProtobuf._MessageImpleme
 extension SDLV6AssistProbeReply: SwiftProtobuf.Message, SwiftProtobuf._MessageImplementationBase, SwiftProtobuf._ProtoNameProviding {
   static let protoMessageName: String = "SDLV6AssistProbeReply"
   static let _protobuf_nameMap: SwiftProtobuf._NameMap = [
-    1: .standard(proto: "v6_info"),
+    1: .standard(proto: "pkt_id"),
+    2: .standard(proto: "v6_info"),
   ]
 
   mutating func decodeMessage<D: SwiftProtobuf.Decoder>(decoder: inout D) throws {
@@ -740,7 +741,8 @@ extension SDLV6AssistProbeReply: SwiftProtobuf.Message, SwiftProtobuf._MessageIm
       // allocates stack space for every case branch when no optimizations are
       // enabled. https://github.com/apple/swift-protobuf/issues/1034
       switch fieldNumber {
-      case 1: try { try decoder.decodeSingularMessageField(value: &self._v6Info) }()
+      case 1: try { try decoder.decodeSingularUInt32Field(value: &self.pktID) }()
+      case 2: try { try decoder.decodeSingularMessageField(value: &self._v6Info) }()
       default: break
       }
     }
@@ -751,13 +753,17 @@ extension SDLV6AssistProbeReply: SwiftProtobuf.Message, SwiftProtobuf._MessageIm
     // allocates stack space for every if/case branch local when no optimizations
     // are enabled. https://github.com/apple/swift-protobuf/issues/1034 and
     // https://github.com/apple/swift-protobuf/issues/1182
+    if self.pktID != 0 {
+      try visitor.visitSingularUInt32Field(value: self.pktID, fieldNumber: 1)
+    }
     try { if let v = self._v6Info {
-      try visitor.visitSingularMessageField(value: v, fieldNumber: 1)
+      try visitor.visitSingularMessageField(value: v, fieldNumber: 2)
     } }()
     try unknownFields.traverse(visitor: &visitor)
   }
 
   static func ==(lhs: SDLV6AssistProbeReply, rhs: SDLV6AssistProbeReply) -> Bool {
+    if lhs.pktID != rhs.pktID {return false}
     if lhs._v6Info != rhs._v6Info {return false}
     if lhs.unknownFields != rhs.unknownFields {return false}
     return true
@@ -772,7 +778,6 @@ extension SDLWelcome: SwiftProtobuf.Message, SwiftProtobuf._MessageImplementatio
     3: .standard(proto: "max_packet_size"),
     4: .standard(proto: "heartbeat_sec"),
     5: .standard(proto: "ipv6_assist"),
-    6: .standard(proto: "assist_token"),
   ]
 
   mutating func decodeMessage<D: SwiftProtobuf.Decoder>(decoder: inout D) throws {
@@ -786,7 +791,6 @@ extension SDLWelcome: SwiftProtobuf.Message, SwiftProtobuf._MessageImplementatio
       case 3: try { try decoder.decodeSingularUInt32Field(value: &self.maxPacketSize) }()
       case 4: try { try decoder.decodeSingularUInt32Field(value: &self.heartbeatSec) }()
       case 5: try { try decoder.decodeSingularMessageField(value: &self._ipv6Assist) }()
-      case 6: try { try decoder.decodeSingularBytesField(value: &self.assistToken) }()
       default: break
       }
     }
@@ -812,9 +816,6 @@ extension SDLWelcome: SwiftProtobuf.Message, SwiftProtobuf._MessageImplementatio
     try { if let v = self._ipv6Assist {
       try visitor.visitSingularMessageField(value: v, fieldNumber: 5)
     } }()
-    if !self.assistToken.isEmpty {
-      try visitor.visitSingularBytesField(value: self.assistToken, fieldNumber: 6)
-    }
     try unknownFields.traverse(visitor: &visitor)
   }
 
@@ -824,7 +825,6 @@ extension SDLWelcome: SwiftProtobuf.Message, SwiftProtobuf._MessageImplementatio
     if lhs.maxPacketSize != rhs.maxPacketSize {return false}
     if lhs.heartbeatSec != rhs.heartbeatSec {return false}
     if lhs._ipv6Assist != rhs._ipv6Assist {return false}
-    if lhs.assistToken != rhs.assistToken {return false}
     if lhs.unknownFields != rhs.unknownFields {return false}
     return true
   }

Tun/Punchnet/SDLIPV6AssistClient.swift

@@ -7,6 +7,10 @@
 import Foundation
 import Network
 
+enum SDLIPV6AssistError: Error {
+    case lostConnection
+}
+
 actor SDLIPV6AssistClient {
     
     private enum State {
@@ -20,10 +24,8 @@ actor SDLIPV6AssistClient {
     private var receiveTask: Task<Void, Never>?
     private let assistServerAddress: NWEndpoint
     
-    // 用于对外输出收到的原始 IP 响应包
+    private var packetId: UInt32 = 1
-    let packetFlow: AsyncStream<SDLV6AssistProbeReply>
+    private var pendingRequests: [UInt32: CheckedContinuation<SDLV6AssistProbeReply, Error>] = [:]
-    private let packetContinuation: AsyncStream<SDLV6AssistProbeReply>.Continuation
-    private var didFinishPacketFlow = false
     
     // 用来处理关闭事件
     private let closeStream: AsyncStream<Void>
@@ -35,10 +37,6 @@ actor SDLIPV6AssistClient {
             return nil
         }
         
-        let (packetStream, packetContinuation) = AsyncStream.makeStream(of: SDLV6AssistProbeReply.self, bufferingPolicy: .bufferingNewest(256))
-        self.packetFlow = packetStream
-        self.packetContinuation = packetContinuation
-        
         let (closeStream, closeContinuation) = AsyncStream.makeStream(of: Void.self, bufferingPolicy: .bufferingNewest(1))
         self.closeStream = closeStream
         self.closeContinuation = closeContinuation
@@ -106,15 +104,36 @@ actor SDLIPV6AssistClient {
         }
     }
     
-    func probe() {
+    func probe() async throws -> SDLV6AssistProbeReply {
         guard case .running = self.state, let connection = self.connection, connection.state == .ready else {
-            return
+            throw SDLIPV6AssistError.lostConnection
         }
-        var assistProbe = SDLV6AssistProbe()
-        assistProbe.assistToken = Data()
         
-        if let data = try? assistProbe.serializedData() {
+        return try await withCheckedThrowingContinuation { cont in
-            connection.send(content: data, completion: .contentProcessed { _ in})
+            let pktId = self.nextPacketId()
+            var assistProbe = SDLV6AssistProbe()
+            assistProbe.pktID = pktId
+            
+            do {
+                let data = try assistProbe.serializedData()
+                connection.send(content: data, completion: .contentProcessed { error in
+                    if let error {
+                        Task {
+                            await self.handleProcesseError(packetId: pktId, error: error)
+                        }
+                    }
+                })
+                self.pendingRequests[pktId] = cont
+            } catch let err {
+                cont.resume(throwing: err)
+            }
+        }
+        
+    }
+    
+    private func handleProcesseError(packetId: UInt32, error: NWError) {
+        if let cont = self.pendingRequests.removeValue(forKey: packetId) {
+            cont.resume(throwing: error)
         }
     }
     
@@ -128,7 +147,6 @@ actor SDLIPV6AssistClient {
         self.receiveTask = nil
         self.connection?.cancel()
         self.connection = nil
-        self.finishPacketFlowIfNeeded()
         self.finishCloseStreamIfNeeded()
     }
     
@@ -170,13 +188,12 @@ actor SDLIPV6AssistClient {
     }
     
     private func handleReceivedPacket(_ data: Data) {
-        guard case .running = self.state else {
-            return
-        }
-        
         do {
             let packet = try SDLV6AssistProbeReply(serializedBytes: data)
-            self.packetContinuation.yield(packet)
+            let pktId = packet.pktID
+            if let cont = self.pendingRequests.removeValue(forKey: pktId) {
+                cont.resume(returning: packet)
+            }
         } catch {
             SDLLogger.log("[SDLIPV6AssistClient] Receive error: \(error)", for: .debug)
         }
@@ -194,15 +211,6 @@ actor SDLIPV6AssistClient {
         }
     }
     
-    private func finishPacketFlowIfNeeded() {
-        guard !self.didFinishPacketFlow else {
-            return
-        }
-        
-        self.didFinishPacketFlow = true
-        self.packetContinuation.finish()
-    }
-    
     private func finishCloseStreamIfNeeded() {
         guard !self.didFinishCloseStream else {
             return
@@ -212,7 +220,15 @@ actor SDLIPV6AssistClient {
         self.closeContinuation.finish()
     }
     
+    private func nextPacketId() -> UInt32 {
+        let packetId = self.packetId
+        self.packetId &+= 1
+        
+        return packetId
+    }
+    
     deinit {
         self.connection?.cancel()
     }
+    
 }