devplacepy/CheckEngine.swift at c4490e783a2ffafff79026dabb43675e184795df

 // retoor <retoor@molodetz.nl>
 import Foundation
 import AsyncHTTPClient
 import NIOCore
 import NIOPosix
 import Logging
 actor CheckEngine {
     private var checks: [String: CheckTarget] = [:]
     private var activeTasks: [String: Task<Void, Never>] = [:]
     private let metrics = MetricsStore()
     private let httpClient: HTTPClient
     private let logger = Logger(label: "molouptime.engine")
     enum CheckError: Error, LocalizedError {
         case engineStopped
         case privateTarget(String)
         case checkFailed(String)
         var errorDescription: String? {
             switch self {
             case .engineStopped: return "Engine is stopped"
             case .privateTarget(let t): return "Target \(t) is a private address"
             case .checkFailed(let m): return m
             }
         }
     }
     init() {
         let configuration = HTTPClient.Configuration(
             timeout: .init(connect: .seconds(10), read: .seconds(15))
         )
         self.httpClient = HTTPClient(
             eventLoopGroupProvider: .shared(MultiThreadedEventLoopGroup(numberOfThreads: 4)),
             configuration: configuration
         )
     }
     deinit {
         try? httpClient.syncShutdown()
     }
     var onResult: ((MetricSample) -> Void)?
     var onStateChange: ((String, String, String) -> Void)?
     func addCheck(_ check: CheckTarget) {
         checks[check.uid] = check
         startCheck(check)
         logger.info("Added check \(check.uid) for \(check.target)")
     }
     func removeCheck(_ uid: String) {
         checks[uid] = nil
         activeTasks[uid]?.cancel()
         activeTasks[uid] = nil
         logger.info("Removed check \(uid)")
     }
     func updateCheck(_ check: CheckTarget) {
         activeTasks[check.uid]?.cancel()
         checks[check.uid] = check
         startCheck(check)
         logger.info("Updated check \(check.uid)")
     }
     func listChecks() -> [CheckTarget] {
         return Array(checks.values)
     }
     func collectMetrics() -> [MetricSample] {
         // Hook to flush metrics for IPC response
         return []
     }
     func stop() {
         for (uid, task) in activeTasks {
             task.cancel()
             activeTasks[uid] = nil
         }
         checks.removeAll()
         try? httpClient.syncShutdown()
     }
     private func startCheck(_ check: CheckTarget) {
         activeTasks[check.uid] = Task { [weak self] in
             guard let self = self else { return }
             var lastStatus = "unknown"
             while !Task.isCancelled {
                 let sample = await self.performCheck(check)
                 let statusChanged = sample.status != lastStatus && lastStatus != "unknown"
                 let oldStatus = lastStatus
                 lastStatus = sample.status
                 await self.metrics.record(sample)
                 self.onResult?(sample)
                 if statusChanged {
                     self.onStateChange?(check.uid, oldStatus, sample.status)
                 }
                 do {
                     try await Task.sleep(nanoseconds: UInt64(check.intervalSeconds) * 1_000_000_000)
                 } catch {
                     break
                 }
             }
         }
     }
     private func performCheck(_ check: CheckTarget) async -> MetricSample {
         let start = Date()
         switch check.type {
         case "http":
             return await performHTTPCheck(check, start: start)
         case "dns":
             return await performDNSCheck(check, start: start)
         case "port":
             return await performPortCheck(check, start: start)
         default:
             return MetricSample(
                 checkUid: check.uid,
                 timestamp: start.timeIntervalSince1970,
                 latencyMs: 0,
                 status: "error",
                 statusCode: 0,
                 errorMessage: "Unknown check type: \(check.type)"
             )
         }
     }
     private func performHTTPCheck(_ check: CheckTarget, start: Date) async -> MetricSample {
         guard !GuardUrl.isPrivateURL(check.target) else {
             return MetricSample(
                 checkUid: check.uid,
                 timestamp: start.timeIntervalSince1970,
                 latencyMs: 0,
                 status: "error",
                 statusCode: 0,
                 errorMessage: "Private target rejected: \(check.target)"
             )
         }
         do {
             let request = try HTTPClient.Request(url: check.target, method: .GET)
             let response = try await httpClient.execute(request, deadline: .now() + .seconds(15)).get()
             let latency = Date().timeIntervalSince(start) * 1000
             let status = response.status.code < 500 ? "up" : "down"
             return MetricSample(
                 checkUid: check.uid,
                 timestamp: start.timeIntervalSince1970,
                 latencyMs: latency,
                 status: status,
                 statusCode: Int(response.status.code),
                 errorMessage: nil
             )
         } catch {
             let latency = Date().timeIntervalSince(start) * 1000
             return MetricSample(
                 checkUid: check.uid,
                 timestamp: start.timeIntervalSince1970,
                 latencyMs: latency,
                 status: "down",
                 statusCode: 0,
                 errorMessage: error.localizedDescription
             )
         }
     }
     private func performDNSCheck(_ check: CheckTarget, start: Date) async -> MetricSample {
         let host = URL(string: check.target)?.host ?? check.target
         let result = DnsResolver.resolve(host)
         let latency = Date().timeIntervalSince(start) * 1000
         if let error = result.error {
             return MetricSample(
                 checkUid: check.uid,
                 timestamp: start.timeIntervalSince1970,
                 latencyMs: latency,
                 status: "down",
                 statusCode: 0,
                 errorMessage: error
             )
         }
         return MetricSample(
             checkUid: check.uid,
             timestamp: start.timeIntervalSince1970,
             latencyMs: latency,
             status: result.addresses.isEmpty ? "down" : "up",
             statusCode: result.addresses.count,
             errorMessage: result.addresses.isEmpty ? "No addresses resolved" : nil
         )
     }
     private func performPortCheck(_ check: CheckTarget, start: Date) async -> MetricSample {
         let host = URL(string: check.target)?.host ?? check.target
         let port = check.port ?? 80
         let useTLS = check.protocolType == "tls"
         guard !GuardUrl.isPrivateHost(host) else {
             return MetricSample(
                 checkUid: check.uid,
                 timestamp: start.timeIntervalSince1970,
                 latencyMs: 0,
                 status: "error",
                 statusCode: 0,
                 errorMessage: "Private host rejected: \(host)"
             )
         }
         let result = await PortScanner.scan(host: host, port: port, useTLS: useTLS)
         return MetricSample(
             checkUid: check.uid,
             timestamp: start.timeIntervalSince1970,
             latencyMs: result.latencyMs,
             status: result.open ? "up" : "down",
             statusCode: result.open ? 1 : 0,
             errorMessage: result.error
         )
     }
 }