111 lines
3.4 KiB
Swift
111 lines
3.4 KiB
Swift
//
|
||
// IdentityStore.swift
|
||
// punchnet
|
||
//
|
||
// Created by 安礼成 on 2026/2/5.
|
||
//
|
||
import Foundation
|
||
import NIO
|
||
|
||
final class IdentitySession {
|
||
var version: UInt32
|
||
var totalNum: UInt32
|
||
private var parts: [UInt32: SDLPolicyResponse] = [:]
|
||
|
||
init(part: SDLPolicyResponse) {
|
||
self.version = part.version
|
||
self.totalNum = part.totalNum
|
||
self.parts[part.index] = part
|
||
}
|
||
|
||
func merge(part: SDLPolicyResponse) {
|
||
if part.version < version {
|
||
// 低版本数据丢弃
|
||
} else if part.version == version {
|
||
self.parts[part.index] = part
|
||
} else {
|
||
self.parts.removeAll()
|
||
self.parts[part.index] = part
|
||
}
|
||
}
|
||
|
||
func process() -> Data? {
|
||
// parts是连续的，从0开始，并且数量等于total_num
|
||
let indexs = parts.keys.sorted().map { UInt32($0) }
|
||
guard indexs.count == self.totalNum && isContinuousFromZero(indexs: indexs) else {
|
||
return nil
|
||
}
|
||
|
||
var rulesData: Data = Data()
|
||
for i in 0..<totalNum {
|
||
if let part = self.parts[i] {
|
||
rulesData.append(part.rules)
|
||
}
|
||
}
|
||
SDLLogger.shared.log("[IdentitySession] get a completed rules: \(rulesData.count)")
|
||
|
||
return rulesData
|
||
}
|
||
|
||
private func isContinuousFromZero(indexs: [UInt32]) -> Bool {
|
||
guard !indexs.isEmpty else {
|
||
return false
|
||
}
|
||
|
||
return indexs.enumerated().allSatisfy { idx, value in
|
||
idx == value
|
||
}
|
||
}
|
||
|
||
}
|
||
|
||
actor IdentityStore {
|
||
typealias IdentityID = UInt32
|
||
|
||
nonisolated private let alloctor = ByteBufferAllocator()
|
||
|
||
private let publisher: SnapshotPublisher<IdentitySnapshot>
|
||
private var identityMap: [IdentityID: IdentityRuleMap] = [:]
|
||
private var sessions: [IdentityID: IdentitySession] = [:]
|
||
|
||
init(publisher: SnapshotPublisher<IdentitySnapshot>) {
|
||
self.publisher = publisher
|
||
}
|
||
|
||
func apply(policyResponse: SDLPolicyResponse) {
|
||
let id = policyResponse.srcIdentityID
|
||
let session = self.sessions[id, default: IdentitySession(part: policyResponse)]
|
||
session.merge(part: policyResponse)
|
||
|
||
// 判断一下是否接受完成
|
||
if let rulesData = session.process() {
|
||
var buffer = alloctor.buffer(bytes: rulesData)
|
||
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(ruleMap: ruleMap)
|
||
|
||
// 删除当前的session信息
|
||
self.sessions.removeValue(forKey: id)
|
||
|
||
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)
|
||
} else {
|
||
self.sessions[id] = session
|
||
}
|
||
}
|
||
|
||
private func compileSnapshot() -> IdentitySnapshot {
|
||
return IdentitySnapshot(identityMap: identityMap)
|
||
}
|
||
|
||
}
|