anlicheng/punchnet-macos
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

完善FlowSession

Browse Source
This commit is contained in:
anlicheng 2026-05-25 21:54:20 +08:00
parent 03e8854e79
commit c95d3145be
7 changed files with 175 additions and 60 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
Tun/Inbound/PacketInboundActor.swift
View File

@ -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)

23
Tun/NetworkStack/IPPacket.swift
View File

@ -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:

24
Tun/Outbound/PacketOutboundActor.swift
View File

@ -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) {

16
Tun/Outbound/PacketOutboundRouter.swift
View File

@ -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
}
}
}

135
Tun/Policy/FlowSessionTable.swift
View File

@ -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 }
}
// 返回当前会话数（调试/统计用）
@ -118,13 +167,51 @@ 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
}
}
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)

23
Tun/Policy/PolicyRuntime.swift
View File

@ -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)
@ -58,9 +56,24 @@ struct PolicyRuntime: @unchecked Sendable {
}
}
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)"

12
Tun/Policy/PolicyService.swift
View File

@ -12,8 +12,7 @@ actor PolicyService {
private let policyRuleStore: PolicyRuleStore
nonisolated private let snapshotPublisher: SnapshotPublisher<PolicyRuleSnapshot>
// 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? {