From c95d3145be5cae41cde87342439be79ef3ae8972 Mon Sep 17 00:00:00 2001 From: anlicheng <244108715@qq.com> Date: Mon, 25 May 2026 21:54:20 +0800 Subject: [PATCH] =?UTF-8?q?=E5=AE=8C=E5=96=84FlowSession?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- Tun/Inbound/PacketInboundActor.swift | 2 +- Tun/NetworkStack/IPPacket.swift | 23 +++- Tun/Outbound/PacketOutboundActor.swift | 24 +++-- Tun/Outbound/PacketOutboundRouter.swift | 16 --- Tun/Policy/FlowSessionTable.swift | 135 +++++++++++++++++++----- Tun/Policy/PolicyRuntime.swift | 23 +++- Tun/Policy/PolicyService.swift | 12 ++- 7 files changed, 175 insertions(+), 60 deletions(-) diff --git a/Tun/Inbound/PacketInboundActor.swift b/Tun/Inbound/PacketInboundActor.swift index 2423ac3..f33d47c 100644 --- a/Tun/Inbound/PacketInboundActor.swift +++ b/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) diff --git a/Tun/NetworkStack/IPPacket.swift b/Tun/NetworkStack/IPPacket.swift index eef8c6b..62062a9 100644 --- a/Tun/NetworkStack/IPPacket.swift +++ b/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: diff --git a/Tun/Outbound/PacketOutboundActor.swift b/Tun/Outbound/PacketOutboundActor.swift index 3958fdb..54b6c1b 100644 --- a/Tun/Outbound/PacketOutboundActor.swift +++ b/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) { diff --git a/Tun/Outbound/PacketOutboundRouter.swift b/Tun/Outbound/PacketOutboundRouter.swift index 14ac4fc..72ee140 100644 --- a/Tun/Outbound/PacketOutboundRouter.swift +++ b/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 - } - } - -} diff --git a/Tun/Policy/FlowSessionTable.swift b/Tun/Policy/FlowSessionTable.swift index 51ec872..5e1c0be 100644 --- a/Tun/Policy/FlowSessionTable.swift +++ b/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(sessionTimeout: TimeInterval = 300) { - self.sessionTimeout = sessionTimeout + init(tcpPendingTimeout: TimeInterval = 30, + tcpEstablishedTimeout: TimeInterval = 300, + udpTimeout: TimeInterval = 30, + dnsTimeout: TimeInterval = 10) { + self.tcpPendingTimeout = tcpPendingTimeout + self.tcpEstablishedTimeout = tcpEstablishedTimeout + self.udpTimeout = udpTimeout + self.dnsTimeout = dnsTimeout } - // 插入或更新会话 - func updateSession(_ key: FlowSession) { - lock.lock() - defer { - lock.unlock() - } - sessions[key] = Date().timeIntervalSince1970 + sessionTimeout - } - - // 查找会话 - func hasSession(_ key: FlowSession) -> Bool { + func recordOutboundTCP(_ key: FlowSession, flags: TCPFlags) { lock.lock() defer { lock.unlock() } - if let expireTs = sessions[key] { - if expireTs >= Date().timeIntervalSince1970 { - return true - } - self.sessions.removeValue(forKey: key) + if flags.contains(.rst) || flags.contains(.fin) { + sessions.removeValue(forKey: key) + return } - 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) -> 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, + 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) diff --git a/Tun/Policy/PolicyRuntime.swift b/Tun/Policy/PolicyRuntime.swift index 0288a9f..8034450 100644 --- a/Tun/Policy/PolicyRuntime.swift +++ b/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(), - self.flowSessionTable.hasSession(reverseFlowSession) { - self.flowSessionTable.updateSession(reverseFlowSession) + if self.isAllowedBySession(ipPacket: ipPacket) { 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)" diff --git a/Tun/Policy/PolicyService.swift b/Tun/Policy/PolicyService.swift index e936172..40501e0 100644 --- a/Tun/Policy/PolicyService.swift +++ b/Tun/Policy/PolicyService.swift @@ -12,8 +12,7 @@ actor PolicyService { private let policyRuleStore: PolicyRuleStore nonisolated private let snapshotPublisher: SnapshotPublisher - // Flow流会话管理， 过期时间为: 180秒 - nonisolated private let flowSessionTable = FlowSessionTable(sessionTimeout: 180) + nonisolated private let flowSessionTable = FlowSessionTable() // 当前节点的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? {