114 lines
3.8 KiB
Swift
114 lines
3.8 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
|
||
}
|
||
|
||
// 批量更新, 有外部任务驱动，因为这里依赖于当前的quicClient
|
||
func batUpdatePolicy(using quicClient: SDLQUICClient?, dstIdentityID: UInt32) {
|
||
guard let quicClient else {
|
||
return
|
||
}
|
||
|
||
self.identityMap.keys.forEach { identityId in
|
||
var policyRequest = SDLPolicyRequest()
|
||
policyRequest.srcIdentityID = identityId
|
||
policyRequest.dstIdentityID = dstIdentityID
|
||
policyRequest.version = self.nextVersion(identityId: identityId)
|
||
|
||
// 发送请求
|
||
if let queryData = try? policyRequest.serializedData() {
|
||
quicClient.send(type: .policyRequest, data: queryData)
|
||
}
|
||
}
|
||
}
|
||
|
||
// 提交权限请求
|
||
func policyRequest(srcIdentityId: UInt32, dstIdentityId: UInt32, using quicClient: SDLQUICClient?) {
|
||
guard let quicClient, !coolingDown.contains(srcIdentityId) else {
|
||
return
|
||
}
|
||
|
||
var policyRequest = SDLPolicyRequest()
|
||
policyRequest.srcIdentityID = srcIdentityId
|
||
policyRequest.dstIdentityID = dstIdentityId
|
||
policyRequest.version = self.nextVersion(identityId: srcIdentityId)
|
||
|
||
// 触发一次打洞
|
||
coolingDown.insert(srcIdentityId)
|
||
// 发送请求
|
||
if let queryData = try? policyRequest.serializedData() {
|
||
quicClient.send(type: .policyRequest, data: queryData)
|
||
}
|
||
|
||
Task {
|
||
// 启动冷却期
|
||
try? await Task.sleep(for: .seconds(5))
|
||
self.endCooldown(for: srcIdentityId)
|
||
}
|
||
}
|
||
|
||
// 处理权限的响应
|
||
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)
|
||
}
|
||
|
||
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
|
||
}
|
||
}
|