punchnet-macos/Tun/Punchnet/Policy/IdentityStore.swift
2026-03-06 15:06:25 +08:00

91 lines
2.9 KiB
Swift

//
// IdentityStore.swift
// punchnet
//
// Created by on 2026/2/5.
//
import Foundation
import NIO
actor IdentityStore {
typealias IdentityID = UInt32
//
nonisolated private let cooldown: Duration = .seconds(5)
// identityId
private var coolingDown: Set<UInt32> = []
// , map[identityId] = version
private var versions: [UInt32: UInt32] = [:]
nonisolated private let alloctor = ByteBufferAllocator()
private let publisher: SnapshotPublisher<IdentitySnapshot>
private var identityMap: [IdentityID: IdentityRuleMap] = [:]
init(publisher: SnapshotPublisher<IdentitySnapshot>) {
self.publisher = publisher
}
//
func policyRequest(using quicClient: SDLQUICClient?, request: inout SDLPolicyRequest) {
let identityId = request.srcIdentityID
guard let quicClient, !coolingDown.contains(identityId) else {
return
}
//
coolingDown.insert(identityId)
let version = self.versions[identityId, default: 1]
request.version = version
//
self.versions[identityId] = version + 1
//
if let queryData = try? request.serializedData() {
quicClient.send(type: .policyRequest, data: queryData)
}
Task {
//
try? await Task.sleep(for: .seconds(5))
self.endCooldown(for: identityId)
}
}
func apply(policyResponse: SDLPolicyResponse) {
let id = policyResponse.srcIdentityID
let version = policyResponse.version
guard self.identityMap[id] == nil || ((self.identityMap[id]?.version ?? 0) < version) else {
return
}
//
var buffer = alloctor.buffer(bytes: policyResponse.rules)
var ruleMap: [UInt8: [UInt16: Bool]] = [:]
while true {
guard let proto = buffer.readInteger(endianness: .big, as: UInt8.self),
let port = buffer.readInteger(endianness: .big, as: UInt16.self) else {
break
}
ruleMap[proto, default: [:]][port] = true
}
self.identityMap[id] = IdentityRuleMap(version: version, ruleMap: ruleMap)
SDLLogger.shared.log("[IdentitySession] get compile Snapshot rules nums: \(self.identityMap[id]?.ruleMap.count), success: \(self.identityMap[id]?.isAllow(proto: 1, port: 80))")
//
let snapshot = compileSnapshot()
publisher.publish(snapshot)
}
private func compileSnapshot() -> IdentitySnapshot {
return IdentitySnapshot(identityMap: identityMap)
}
private func endCooldown(for key: UInt32) {
self.coolingDown.remove(key)
}
}