devplacepy/molouptime/Sources/CheckEngine.swift

// 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
        )
    }
}