127 lines
3.1 KiB
Swift
127 lines
3.1 KiB
Swift
//
|
||
// FiveTuple.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 SDLFlowSessionManager {
|
||
private var sessions: [FlowSession: TimeInterval] = [:]
|
||
private let lock = NSLock()
|
||
private let sessionTimeout: TimeInterval
|
||
|
||
/// - Parameter sessionTimeout: 会话闲置多久（秒）被清理
|
||
init(sessionTimeout: TimeInterval = 300) {
|
||
self.sessionTimeout = sessionTimeout
|
||
}
|
||
|
||
// 插入或更新会话
|
||
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 expireTs >= Date().timeIntervalSince1970 {
|
||
return true
|
||
}
|
||
self.sessions.removeValue(forKey: key)
|
||
}
|
||
|
||
return false
|
||
}
|
||
|
||
// 删除会话
|
||
func removeSession(_ key: FlowSession) {
|
||
lock.lock()
|
||
defer {
|
||
lock.unlock()
|
||
}
|
||
|
||
sessions.removeValue(forKey: key)
|
||
}
|
||
|
||
// 清理过期会话
|
||
func cleanupExpiredSessions() {
|
||
lock.lock()
|
||
defer {
|
||
lock.unlock()
|
||
}
|
||
|
||
let now = Date().timeIntervalSince1970
|
||
self.sessions = self.sessions.filter { $0.value >= now }
|
||
}
|
||
|
||
// 返回当前会话数（调试/统计用）
|
||
var count: Int {
|
||
lock.lock()
|
||
defer {
|
||
lock.unlock()
|
||
}
|
||
return sessions.count
|
||
}
|
||
|
||
}
|
||
|
||
extension IPPacket {
|
||
|
||
func flowSession() -> FlowSession? {
|
||
switch self.transportPacket {
|
||
case .tcp(let tcpPacket):
|
||
return FlowSession(srcIP: header.source, dstIP: header.destination, srcPort: tcpPacket.header.srcPort, dstPort: tcpPacket.header.dstPort, proto: header.proto)
|
||
case .udp(let udpPacket):
|
||
return FlowSession(srcIP: header.source, dstIP: header.destination, srcPort: udpPacket.srcPort, dstPort: udpPacket.dstPort, proto: header.proto)
|
||
default:
|
||
return nil
|
||
}
|
||
}
|
||
}
|