punchnet-macos/Tun/Policy/PolicyRuleStore.swift

//
//  PolicyRuleStore.swift
//  punchnet
//  1. 需要增加规则基于轮训更新的逻辑
//  Created by 安礼成 on 2026/2/5.
//
import Foundation
import NIO

actor PolicyRuleStore {

    private struct PolicyEntry {
        let ruleMap: PolicyRuleMap
        let expiresAt: TimeInterval
    }

    private struct PendingRequest {
        let version: UInt32
        let retryCount: Int
        let nextRetryAt: TimeInterval
    }

    private let policyTTL: TimeInterval = 120
    private let refreshLeadTime: TimeInterval = 15
    private let baseRetryDelay: TimeInterval = 2
    private let maxRetryDelay: TimeInterval = 60

    // 处理各个请求的版本问题, map[identityId] = version
    private var versions: [UInt32: UInt32] = [:]
    private var pendingByIdentity: [UInt32: PendingRequest] = [:]
    
    nonisolated private let alloctor = ByteBufferAllocator()
    
    private let publisher: SnapshotPublisher<PolicyRuleSnapshot>
    private var policyByIdentity: [UInt32: PolicyEntry] = [:]
    
    init(publisher: SnapshotPublisher<PolicyRuleSnapshot>) {
        self.publisher = publisher
    }

    func makeBatchPolicyRequests(dstIdentityID: UInt32) -> [Data] {
        let now = Date().timeIntervalSince1970
        let identities = Set(self.policyByIdentity.keys).union(self.pendingByIdentity.keys)

        return identities.compactMap { identityId in
            self.makePolicyRequestIfDue(srcIdentityId: identityId, dstIdentityId: dstIdentityID, now: now, forceMissing: false)
        }
    }

    func makePolicyRequest(srcIdentityId: UInt32, dstIdentityId: UInt32) -> Data? {
        guard self.policyByIdentity[srcIdentityId] == nil else {
            return nil
        }

        return self.makePolicyRequestIfDue(
            srcIdentityId: srcIdentityId,
            dstIdentityId: dstIdentityId,
            now: Date().timeIntervalSince1970,
            forceMissing: true
        )
    }
    
    // 处理权限的响应
    func applyPolicyResponse(_ policyResponse: SDLPolicyResponse) {
        let id = policyResponse.srcIdentityID
        let version = policyResponse.version
        
        if let pending = self.pendingByIdentity[id], pending.version <= version {
            self.pendingByIdentity.removeValue(forKey: id)
        }

        guard self.policyByIdentity[id] == nil || ((self.policyByIdentity[id]?.ruleMap.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
        }
        let now = Date().timeIntervalSince1970
        self.policyByIdentity[id] = PolicyEntry(
            ruleMap: PolicyRuleMap(version: version, ruleMap: ruleMap),
            expiresAt: now + self.policyTTL
        )
        SDLLogger.log("[PolicyRuleStore] apply policy response, srcIdentityID: \(id), version: \(version), rulesCount: \(ruleMap.reduce(0) { $0 + $1.value.count })", for: .debug)
        
        // 发布新的快照信息
        let snapshot = compileSnapshot()
        publisher.publish(snapshot)
    }

    func clear() {
        self.versions.removeAll()
        self.pendingByIdentity.removeAll()
        self.policyByIdentity.removeAll()
        self.publisher.publish(PolicyRuleSnapshot.empty())
    }
    
    private func compileSnapshot() -> PolicyRuleSnapshot {
        return PolicyRuleSnapshot(ruleMapByIdentity: self.policyByIdentity.mapValues(\.ruleMap))
    }
    
    private func nextVersion(identityId: UInt32) -> UInt32 {
        let version = self.versions[identityId, default: 1]
        // 更新请求的版本问题
        self.versions[identityId] = version + 1
        
        return version
    }

    private func makePolicyRequestIfDue(srcIdentityId: UInt32, dstIdentityId: UInt32, now: TimeInterval, forceMissing: Bool) -> Data? {
        if let pending = self.pendingByIdentity[srcIdentityId], pending.nextRetryAt > now {
            return nil
        }

        if !forceMissing,
           let entry = self.policyByIdentity[srcIdentityId],
           entry.expiresAt - self.refreshLeadTime > now {
            return nil
        }

        let version = self.nextVersion(identityId: srcIdentityId)
        guard let data = self.makePolicyRequestData(srcIdentityId: srcIdentityId, dstIdentityId: dstIdentityId, version: version) else {
            return nil
        }

        let retryCount = self.pendingByIdentity[srcIdentityId].map { $0.retryCount + 1 } ?? 0
        let retryDelay = self.retryDelay(for: retryCount, identityId: srcIdentityId)
        self.pendingByIdentity[srcIdentityId] = PendingRequest(
            version: version,
            retryCount: retryCount,
            nextRetryAt: now + retryDelay
        )

        return data
    }

    private func makePolicyRequestData(srcIdentityId: UInt32, dstIdentityId: UInt32, version: UInt32) -> Data? {
        var policyRequest = SDLPolicyRequest()
        policyRequest.srcIdentityID = srcIdentityId
        policyRequest.dstIdentityID = dstIdentityId
        policyRequest.version = version
        return try? policyRequest.serializedData()
    }

    private func retryDelay(for retryCount: Int, identityId: UInt32) -> TimeInterval {
        let cappedRetryCount = min(retryCount, 5)
        let multiplier = Double(1 << cappedRetryCount)
        let jitter = Double(identityId % 1_000) / 1_000
        return min(self.baseRetryDelay * multiplier + jitter, self.maxRetryDelay)
    }
}