punchnet-macos/Tun/Punchnet/Policy/IdentityStore.swift
2026-05-19 23:14:06 +08:00

107 lines
3.4 KiB
Swift

//
// IdentityStore.swift
// punchnet
// 1.
// Created by on 2026/2/5.
//
import Foundation
import NIO
actor IdentityStore {
//
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: [UInt32: IdentityRuleMap] = [:]
init(publisher: SnapshotPublisher<IdentitySnapshot>) {
self.publisher = publisher
}
func makeBatchPolicyRequests(dstIdentityID: UInt32) -> [Data] {
return self.identityMap.keys.compactMap { identityId in
var policyRequest = SDLPolicyRequest()
policyRequest.srcIdentityID = identityId
policyRequest.dstIdentityID = dstIdentityID
policyRequest.version = self.nextVersion(identityId: identityId)
return try? policyRequest.serializedData()
}
}
func makePolicyRequest(srcIdentityId: UInt32, dstIdentityId: UInt32) -> Data? {
guard !coolingDown.contains(srcIdentityId) else {
return nil
}
var policyRequest = SDLPolicyRequest()
policyRequest.srcIdentityID = srcIdentityId
policyRequest.dstIdentityID = dstIdentityId
policyRequest.version = self.nextVersion(identityId: srcIdentityId)
coolingDown.insert(srcIdentityId)
Task { [weak self] in
try? await Task.sleep(for: .seconds(5))
await self?.endCooldown(for: srcIdentityId)
}
return try? policyRequest.serializedData()
}
//
func applyPolicyResponse(_ 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)
//
let snapshot = compileSnapshot()
publisher.publish(snapshot)
}
func clear() {
self.coolingDown.removeAll()
self.versions.removeAll()
self.identityMap.removeAll()
self.publisher.publish(IdentitySnapshot.empty())
}
private func compileSnapshot() -> IdentitySnapshot {
return IdentitySnapshot(identityMap: identityMap)
}
private func endCooldown(for key: UInt32) {
self.coolingDown.remove(key)
}
private func nextVersion(identityId: UInt32) -> UInt32 {
let version = self.versions[identityId, default: 1]
//
self.versions[identityId] = version + 1
return version
}
}