punchnet-macos/Tun/Policy/FlowSessionTable.swift

//
//  FlowSessionTable.swift
//  punchnet
//  tcp/udp Flow流管理
//  Created by 安礼成 on 2026/3/10.
//
import Foundation

// MARK: - 五元组 key
struct FlowSession: Hashable {
    let srcIP: UInt32
    let dstIP: UInt32
    let srcPort: UInt16
    let dstPort: UInt16
    let proto: UInt8

    func hash(into hasher: inout Hasher) {
        // 高效组合 hash
        hasher.combine(srcIP)
        hasher.combine(dstIP)
        hasher.combine(UInt32(srcPort) << 16 | UInt32(dstPort))
        hasher.combine(proto)
    }

    static func ==(lhs: Self, rhs: Self) -> Bool {
        return lhs.srcIP == rhs.srcIP &&
               lhs.dstIP == rhs.dstIP &&
               lhs.srcPort == rhs.srcPort &&
               lhs.dstPort == rhs.dstPort &&
               lhs.proto == rhs.proto
    }

    func reverse() -> FlowSession {
        return FlowSession(
            srcIP: dstIP,
            dstIP: srcIP,
            srcPort: dstPort,
            dstPort: srcPort,
            proto: proto
        )
    }

}

// MARK: - 会话表
final class FlowSessionTable: @unchecked Sendable {
    
    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 tcpPendingTimeout: TimeInterval
    private let tcpEstablishedTimeout: TimeInterval
    private let udpTimeout: TimeInterval
    private let dnsTimeout: TimeInterval

    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 recordOutboundTCP(_ key: FlowSession, flags: TCPFlags) {
        lock.lock()
        defer {
            lock.unlock()
        }

        if flags.contains(.rst) || flags.contains(.fin) {
            sessions.removeValue(forKey: key)
            return
        }

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

    // 删除会话
    func removeSession(_ key: FlowSession) {
        lock.lock()
        defer {
            lock.unlock()
        }
        
        sessions.removeValue(forKey: key)
    }

    func clear() {
        lock.lock()
        defer {
            lock.unlock()
        }

        sessions.removeAll()
    }

    // 清理过期会话
    func cleanupExpiredSessions() {
        lock.lock()
        defer {
            lock.unlock()
        }
        
        let now = Date().timeIntervalSince1970
        self.sessions = self.sessions.filter { $0.value.expiresAt >= now }
    }

    // 返回当前会话数（调试/统计用）
    var count: Int {
        lock.lock()
        defer {
            lock.unlock()
        }
        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, _):
            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)
        default:
            return nil
        }
    }
}