完善FlowSession

Tun/Inbound/PacketInboundActor.swift

@@ -188,7 +188,7 @@ actor PacketInboundActor {
     private func makePolicyPacketContext(identityID: UInt32, ipPacket: IPPacketView) -> PolicyPacketContext {
         let ports: (UInt16?, UInt16?)
         switch ipPacket.transportPacket {
-        case .tcp(let srcPort, let dstPort):
+        case .tcp(let srcPort, let dstPort, _):
             ports = (srcPort, dstPort)
         case .udp(let srcPort, let dstPort, _):
             ports = (srcPort, dstPort)

Tun/NetworkStack/IPPacket.swift

@@ -38,6 +38,19 @@ struct IPHeader {
     }
 }
 
+struct TCPFlags: OptionSet {
+    let rawValue: UInt16
+
+    static let fin = TCPFlags(rawValue: 1 << 0)
+    static let syn = TCPFlags(rawValue: 1 << 1)
+    static let rst = TCPFlags(rawValue: 1 << 2)
+    static let psh = TCPFlags(rawValue: 1 << 3)
+    static let ack = TCPFlags(rawValue: 1 << 4)
+    static let urg = TCPFlags(rawValue: 1 << 5)
+    static let ece = TCPFlags(rawValue: 1 << 6)
+    static let cwr = TCPFlags(rawValue: 1 << 7)
+}
+
 // MARK: - Lightweight IP Packet View
 
 struct IPPacketView {
@@ -46,7 +59,7 @@ struct IPPacketView {
     let transportPacket: TransportPacket
 
     enum TransportPacket {
-        case tcp(srcPort: UInt16, dstPort: UInt16)
+        case tcp(srcPort: UInt16, dstPort: UInt16, flags: TCPFlags)
         case udp(srcPort: UInt16, dstPort: UInt16, payloadOffset: Int)
         case icmp
         case unsupported(UInt8)
@@ -95,9 +108,15 @@ struct IPPacketView {
             guard data.count >= offset + 20 else {
                 return .malformed
             }
+            let offsetAndFlags = UInt16(bytes: (Self.byte(at: offset + 12, in: data), Self.byte(at: offset + 13, in: data)))
+            let dataOffset = Int(offsetAndFlags >> 12) * 4
+            guard dataOffset >= 20, data.count >= offset + dataOffset else {
+                return .malformed
+            }
             return .tcp(
                 srcPort: UInt16(bytes: (Self.byte(at: offset, in: data), Self.byte(at: offset + 1, in: data))),
-                dstPort: UInt16(bytes: (Self.byte(at: offset + 2, in: data), Self.byte(at: offset + 3, in: data)))
+                dstPort: UInt16(bytes: (Self.byte(at: offset + 2, in: data), Self.byte(at: offset + 3, in: data))),
+                flags: TCPFlags(rawValue: offsetAndFlags & 0x01FF)
             )
 
         case .udp:

Tun/Outbound/PacketOutboundActor.swift

@@ -121,10 +121,6 @@ actor PacketOutboundActor {
         let router = PacketOutboundRouter(networkAddress: self.networkAddress, exitNode: self.exitNode)
         let decision = router.route(packet: packet)
         
-        if decision.shouldTrackFlow {
-            self.policyService.recordOutboundFlow(ipPacket: packet)
-        }
-        
         switch decision {
         case .loopback(let ipPacketData):
             let nePacket = NEPacket(data: ipPacketData, protocolFamily: 2)
@@ -136,13 +132,17 @@ actor PacketOutboundActor {
             SDLLogger.log("[PacketOutboundActor] get local dns request: \(name)")
             await self.dnsService?.queryLocal(tracker: tracker, dnsPayload: payload)
         case .forwardToNextHop(let ip, let type, let data, let kind):
-            await self.forwardPacketToNextHop(ip: ip, type: type, data: data, kind: kind)
+            await self.forwardPacketToNextHop(ip: ip, type: type, data: data, kind: kind, originalPacket: packet)
         case .drop(let reason):
             SDLLogger.log("[PacketOutboundActor] drop tun packet, reason: \(reason.rawValue)", for: .trace)
         }
     }
     
-    private func forwardPacketToNextHop(ip: UInt32, type: LayerPacket.PacketType, data: Data, kind: PacketOutboundRouter.ForwardKind) async {
+    private func forwardPacketToNextHop(ip: UInt32,
+                                        type: LayerPacket.PacketType,
+                                        data: Data,
+                                        kind: PacketOutboundRouter.ForwardKind,
+                                        originalPacket: IPPacketView) async {
         switch kind {
         case .sameNetwork:
             SDLLogger.log("[PacketOutboundActor] dstIp: \(SDLUtil.int32ToIp(ip)) same network", for: .trace)
@@ -152,7 +152,10 @@ actor PacketOutboundActor {
 
         if let dstMac = self.arpResolver.snapshot().lookup(ip) {
             SDLLogger.log("[PacketOutboundActor] dstIp: \(SDLUtil.int32ToIp(ip)), dst_mac is: \(SDLUtil.formatMacAddress(mac: dstMac))", for: .trace)
-            await self.routeLayerPacket(dstMac: dstMac, type: type, data: data)
+            let didSend = await self.routeLayerPacket(dstMac: dstMac, type: type, data: data)
+            if didSend {
+                self.policyService.recordOutboundFlow(ipPacket: originalPacket)
+            }
         } else {
             SDLLogger.log("[PacketOutboundActor] dstIp: \(SDLUtil.int32ToIp(ip)) arp query not found, broadcast", for: .trace)
             if let arpRequest = try? await self.arpResolver.makeArpRequest(targetIp: ip) {
@@ -161,9 +164,10 @@ actor PacketOutboundActor {
         }
     }
     
-    func routeLayerPacket(dstMac: Data, type: LayerPacket.PacketType, data: Data) async {
+    @discardableResult
+    func routeLayerPacket(dstMac: Data, type: LayerPacket.PacketType, data: Data) async -> Bool {
         guard let plan = try? self.makeDeliveryPlan(dstMac: dstMac, type: type, data: data) else {
-            return
+            return false
         }
 
         switch plan {
@@ -182,6 +186,8 @@ actor PacketOutboundActor {
                 await self.superServiceProxy.send(type: .queryInfo, data: queryData)
             }
         }
+
+        return true
     }
 
     private func finishPacketReader(generation: UInt64) {

Tun/Outbound/PacketOutboundRouter.swift

@@ -91,19 +91,3 @@ struct PacketOutboundRouter {
         return .localDNS(name: query.name, payload: dnsPayload, tracker: tracker)
     }
 }
-
-extension PacketOutboundRouter.RouteDecision {
-    
-    var shouldTrackFlow: Bool {
-        switch self {
-        case .forwardToNextHop(_, _, _, let kind):
-            switch kind {
-            case .sameNetwork, .exitNode, .dnsExitNode:
-                return true
-            }
-        default:
-            return false
-        }
-    }
-    
-}

Tun/Policy/FlowSessionTable.swift

@@ -44,39 +44,88 @@ struct FlowSession: Hashable {
 
 // MARK: - 会话表
 final class FlowSessionTable: @unchecked Sendable {
-    private var sessions: [FlowSession: TimeInterval] = [:]
+    private enum SessionState: Hashable {
+        case tcpPending
+        case tcpEstablished
+        case udp
+    }
+
+    private struct SessionEntry {
+        let state: SessionState
+        let expiresAt: TimeInterval
+    }
+
+    private var sessions: [FlowSession: SessionEntry] = [:]
     private let lock = NSLock()
-    private let sessionTimeout: TimeInterval
+    private let tcpPendingTimeout: TimeInterval
+    private let tcpEstablishedTimeout: TimeInterval
+    private let udpTimeout: TimeInterval
+    private let dnsTimeout: TimeInterval
 
-    /// - Parameter sessionTimeout: 会话闲置多久（秒）被清理
+    init(tcpPendingTimeout: TimeInterval = 30,
-    init(sessionTimeout: TimeInterval = 300) {
+         tcpEstablishedTimeout: TimeInterval = 300,
-        self.sessionTimeout = sessionTimeout
+         udpTimeout: TimeInterval = 30,
+         dnsTimeout: TimeInterval = 10) {
+        self.tcpPendingTimeout = tcpPendingTimeout
+        self.tcpEstablishedTimeout = tcpEstablishedTimeout
+        self.udpTimeout = udpTimeout
+        self.dnsTimeout = dnsTimeout
     }
 
-    // 插入或更新会话
+    func recordOutboundTCP(_ key: FlowSession, flags: TCPFlags) {
-    func updateSession(_ key: FlowSession) {
-        lock.lock()
-        defer {
-            lock.unlock()
-        }
-        sessions[key] = Date().timeIntervalSince1970 + sessionTimeout
-    }
-    
-    // 查找会话
-    func hasSession(_ key: FlowSession) -> Bool {
         lock.lock()
         defer {
             lock.unlock()
         }
 
-        if let expireTs = sessions[key] {
+        if flags.contains(.rst) || flags.contains(.fin) {
-            if expireTs >= Date().timeIntervalSince1970 {
+            sessions.removeValue(forKey: key)
-                return true
+            return
-            }
-            self.sessions.removeValue(forKey: key)
         }
 
-        return false
+        if flags.contains(.syn) && !flags.contains(.ack) {
+            sessions[key] = SessionEntry(state: .tcpPending, expiresAt: Date().timeIntervalSince1970 + tcpPendingTimeout)
+            return
+        }
+
+        self.touchIfValidLocked(key, allowedStates: [.tcpEstablished], timeout: tcpEstablishedTimeout)
+    }
+
+    func recordOutboundUDP(_ key: FlowSession, isDNS: Bool) {
+        lock.lock()
+        defer {
+            lock.unlock()
+        }
+
+        sessions[key] = SessionEntry(state: .udp, expiresAt: Date().timeIntervalSince1970 + (isDNS ? dnsTimeout : udpTimeout))
+    }
+
+    func allowInboundTCP(_ key: FlowSession, flags: TCPFlags) -> Bool {
+        lock.lock()
+        defer {
+            lock.unlock()
+        }
+
+        if flags.contains(.rst) || flags.contains(.fin) {
+            let existed = self.hasValidLocked(key, allowedStates: [.tcpPending, .tcpEstablished])
+            sessions.removeValue(forKey: key)
+            return existed
+        }
+
+        if flags.contains(.syn) && flags.contains(.ack) {
+            return self.touchIfValidLocked(key, allowedStates: [.tcpPending, .tcpEstablished], timeout: tcpEstablishedTimeout, nextState: .tcpEstablished)
+        }
+
+        return self.touchIfValidLocked(key, allowedStates: [.tcpEstablished], timeout: tcpEstablishedTimeout)
+    }
+
+    func allowInboundUDP(_ key: FlowSession, isDNS: Bool) -> Bool {
+        lock.lock()
+        defer {
+            lock.unlock()
+        }
+
+        return self.touchIfValidLocked(key, allowedStates: [.udp], timeout: isDNS ? dnsTimeout : udpTimeout)
     }
 
     // 删除会话
@@ -106,7 +155,7 @@ final class FlowSessionTable: @unchecked Sendable {
         }
         
         let now = Date().timeIntervalSince1970
-        self.sessions = self.sessions.filter { $0.value >= now }
+        self.sessions = self.sessions.filter { $0.value.expiresAt >= now }
     }
 
     // 返回当前会话数（调试/统计用）
@@ -117,6 +166,44 @@ final class FlowSessionTable: @unchecked Sendable {
         }
         return sessions.count
     }
+
+    private func hasValidLocked(_ key: FlowSession, allowedStates: Set<SessionState>) -> Bool {
+        guard let entry = sessions[key] else {
+            return false
+        }
+
+        guard entry.expiresAt >= Date().timeIntervalSince1970 else {
+            sessions.removeValue(forKey: key)
+            return false
+        }
+
+        return allowedStates.contains(entry.state)
+    }
+
+    @discardableResult
+    private func touchIfValidLocked(_ key: FlowSession,
+                                    allowedStates: Set<SessionState>,
+                                    timeout: TimeInterval,
+                                    nextState: SessionState? = nil) -> Bool {
+        guard let entry = sessions[key] else {
+            return false
+        }
+
+        guard entry.expiresAt >= Date().timeIntervalSince1970 else {
+            sessions.removeValue(forKey: key)
+            return false
+        }
+
+        guard allowedStates.contains(entry.state) else {
+            return false
+        }
+
+        sessions[key] = SessionEntry(
+            state: nextState ?? entry.state,
+            expiresAt: Date().timeIntervalSince1970 + timeout
+        )
+        return true
+    }
     
 }
 
@@ -124,7 +211,7 @@ extension IPPacketView {
 
     func flowSession() -> FlowSession? {
         switch self.transportPacket {
-        case .tcp(let srcPort, let dstPort):
+        case .tcp(let srcPort, let dstPort, _):
             return FlowSession(srcIP: header.source, dstIP: header.destination, srcPort: srcPort, dstPort: dstPort, proto: header.proto)
         case .udp(let srcPort, let dstPort, _):
             return FlowSession(srcIP: header.source, dstIP: header.destination, srcPort: srcPort, dstPort: dstPort, proto: header.proto)

Tun/Policy/PolicyRuntime.swift

@@ -23,9 +23,7 @@ struct PolicyRuntime: @unchecked Sendable {
     }
 
     func evaluateInbound(srcIdentityID: UInt32, ipPacket: IPPacketView) -> InboundDecision {
-        if let reverseFlowSession = ipPacket.flowSession()?.reverse(),
+        if self.isAllowedBySession(ipPacket: ipPacket) {
-           self.flowSessionTable.hasSession(reverseFlowSession) {
-            self.flowSessionTable.updateSession(reverseFlowSession)
             SDLLogger.log("[PolicyRuntime] session hit, src_identify_id: \(srcIdentityID), allow: \(debugInfo(ipPacket: ipPacket))")
             return .allow
         }
@@ -49,7 +47,7 @@ struct PolicyRuntime: @unchecked Sendable {
         let proto = ipPacket.header.proto
 
         switch ipPacket.transportPacket {
-        case .tcp(_, let dstPort):
+        case .tcp(_, let dstPort, _):
             return ruleMap.isAllow(proto: proto, port: dstPort)
         case .udp(_, let dstPort, _):
             return ruleMap.isAllow(proto: proto, port: dstPort)
@@ -57,10 +55,25 @@ struct PolicyRuntime: @unchecked Sendable {
             return false
         }
     }
+
+    private func isAllowedBySession(ipPacket: IPPacketView) -> Bool {
+        guard let reverseFlowSession = ipPacket.flowSession()?.reverse() else {
+            return false
+        }
+
+        switch ipPacket.transportPacket {
+        case .tcp(_, _, let flags):
+            return self.flowSessionTable.allowInboundTCP(reverseFlowSession, flags: flags)
+        case .udp(let srcPort, _, _):
+            return self.flowSessionTable.allowInboundUDP(reverseFlowSession, isDNS: srcPort == 53)
+        default:
+            return false
+        }
+    }
     
     private func debugInfo(ipPacket: IPPacketView) -> String {
         switch ipPacket.transportPacket {
-        case .tcp(_, let dstPort):
+        case .tcp(_, let dstPort, _):
             return "tcp: \(dstPort)"
         case .udp(_, let dstPort, _):
             return "udp: \(dstPort)"

Tun/Policy/PolicyService.swift

@@ -12,8 +12,7 @@ actor PolicyService {
     private let policyRuleStore: PolicyRuleStore
     nonisolated private let snapshotPublisher: SnapshotPublisher<PolicyRuleSnapshot>
     
-    // Flow流会话管理， 过期时间为: 180秒
+    nonisolated private let flowSessionTable = FlowSessionTable()
-    nonisolated private let flowSessionTable = FlowSessionTable(sessionTimeout: 180)
     
     // 当前节点的identityId值
     let identityId: UInt32
@@ -35,7 +34,14 @@ actor PolicyService {
             return
         }
 
-        self.flowSessionTable.updateSession(flowSession)
+        switch ipPacket.transportPacket {
+        case .tcp(_, _, let flags):
+            self.flowSessionTable.recordOutboundTCP(flowSession, flags: flags)
+        case .udp(_, let dstPort, _):
+            self.flowSessionTable.recordOutboundUDP(flowSession, isDNS: dstPort == 53)
+        default:
+            return
+        }
     }
 
     func makePolicyRequest(srcIdentityID: UInt32) async -> Data? {