krz/domain-dig

  1import Foundation
  2
  3/// Opted out of the project's MainActor default isolation: this is a plain
  4/// value pair constructed inside `actor LookupRuntime`, and a MainActor-bound
  5/// memberwise init cannot be called from there under Swift 6.
  6nonisolated struct CachedLookupResult<Value> {
  7    let value: Value
  8    let source: LookupResultSource
  9}
 10
 11/// Opting out of MainActor isolation also opted out of the implicit
 12/// Sendable that globally-isolated types get, which is what lets this cross
 13/// from `actor LookupRuntime` back to its callers.
 14extension CachedLookupResult: Sendable where Value: Sendable {}
 15
 16actor LookupRuntime {
 17    static let shared = LookupRuntime()
 18
 19    private let ttl: TimeInterval = 300
 20
 21    private enum RequestKey: Hashable {
 22        case domain(String, LookupSectionKind)
 23        case subject(String, LookupSectionKind)
 24    }
 25
 26    private enum RateLimitBucket: Hashable {
 27        case crtsh
 28        case rdap
 29        case ipGeolocation
 30
 31        var minimumSpacing: TimeInterval {
 32            switch self {
 33            case .crtsh:
 34                return 1.0
 35            case .rdap:
 36                return 0.75
 37            case .ipGeolocation:
 38                return 0.75
 39            }
 40        }
 41    }
 42
 43    private enum CachedPayload {
 44        case dns(ServiceResult<[DNSSection]>)
 45        case availability(DomainAvailabilityResult)
 46        case ssl(ServiceResult<SSLCertificateInfo>)
 47        case hsts(Bool?)
 48        case http(ServiceResult<HTTPHeadersResult>)
 49        case reachability(ServiceResult<[PortReachability]>)
 50        case ownership(ServiceResult<DomainOwnership>)
 51        case redirect(ServiceResult<[RedirectHop]>)
 52        case subdomains(ServiceResult<[DiscoveredSubdomain]>)
 53        case portScan(ServiceResult<[PortScanResult]>)
 54        case email(ServiceResult<EmailSecurityResult>)
 55        case ptr(ServiceResult<String>)
 56        case ipGeolocation(ServiceResult<IPGeolocation>)
 57        case suggestions([DomainSuggestionResult])
 58    }
 59
 60    private struct CacheEntry {
 61        let payload: CachedPayload
 62        let expiresAt: Date
 63    }
 64
 65    private var cache: [RequestKey: CacheEntry] = [:]
 66    private var inFlight: [RequestKey: Task<CachedPayload, Never>] = [:]
 67    private var nextAllowedAt: [RateLimitBucket: Date] = [:]
 68
 69    func clearCache() {
 70        cache.removeAll()
 71        inFlight.values.forEach { $0.cancel() }
 72        inFlight.removeAll()
 73        nextAllowedAt.removeAll()
 74    }
 75
 76    func dns(domain: String) async -> CachedLookupResult<ServiceResult<[DNSSection]>> {
 77        await execute(
 78            key: .domain(domain, .dns),
 79            extract: { payload in
 80                guard case let .dns(result) = payload else { return nil }
 81                return result
 82            },
 83            operation: {
 84                .dns(await DNSLookupService.lookupAll(domain: domain))
 85            }
 86        )
 87    }
 88
 89    func availability(domain: String) async -> CachedLookupResult<DomainAvailabilityResult> {
 90        CachedLookupResult(
 91            value: await DomainAvailabilityService.check(domain: domain),
 92            source: .live
 93        )
 94    }
 95
 96    func ssl(domain: String) async -> CachedLookupResult<ServiceResult<SSLCertificateInfo>> {
 97        await execute(
 98            key: .domain(domain, .ssl),
 99            extract: { payload in
100                guard case let .ssl(result) = payload else { return nil }
101                return result
102            },
103            operation: {
104                .ssl(await SSLCheckService.check(domain: domain))
105            }
106        )
107    }
108
109    func hsts(domain: String) async -> CachedLookupResult<Bool?> {
110        await execute(
111            key: .domain(domain, .hsts),
112            extract: { payload in
113                guard case let .hsts(result) = payload else { return nil }
114                return result
115            },
116            operation: {
117                .hsts(await SSLCheckService.checkHSTSPreload(domain: domain))
118            }
119        )
120    }
121
122    func http(domain: String) async -> CachedLookupResult<ServiceResult<HTTPHeadersResult>> {
123        await execute(
124            key: .domain(domain, .httpHeaders),
125            extract: { payload in
126                guard case let .http(result) = payload else { return nil }
127                return result
128            },
129            operation: {
130                .http(await HTTPHeadersService.fetch(domain: domain))
131            }
132        )
133    }
134
135    func reachability(domain: String) async -> CachedLookupResult<ServiceResult<[PortReachability]>> {
136        await execute(
137            key: .domain(domain, .reachability),
138            extract: { payload in
139                guard case let .reachability(result) = payload else { return nil }
140                return result
141            },
142            operation: {
143                .reachability(await ReachabilityService.checkAll(domain: domain))
144            }
145        )
146    }
147
148    func ownership(domain: String) async -> CachedLookupResult<ServiceResult<DomainOwnership>> {
149        await execute(
150            key: .domain(domain, .ownership),
151            rateLimitBucket: .rdap,
152            extract: { payload in
153                guard case let .ownership(result) = payload else { return nil }
154                return result
155            },
156            operation: {
157                .ownership(await DomainOwnershipService.lookup(domain: domain))
158            }
159        )
160    }
161
162    func redirectChain(domain: String) async -> CachedLookupResult<ServiceResult<[RedirectHop]>> {
163        await execute(
164            key: .domain(domain, .redirectChain),
165            extract: { payload in
166                guard case let .redirect(result) = payload else { return nil }
167                return result
168            },
169            operation: {
170                .redirect(await RedirectChainService.trace(domain: domain))
171            }
172        )
173    }
174
175    func subdomains(domain: String) async -> CachedLookupResult<ServiceResult<[DiscoveredSubdomain]>> {
176        await execute(
177            key: .domain(domain, .subdomains),
178            rateLimitBucket: .crtsh,
179            extract: { payload in
180                guard case let .subdomains(result) = payload else { return nil }
181                return result
182            },
183            operation: {
184                .subdomains(await SubdomainDiscoveryService.discover(for: domain))
185            }
186        )
187    }
188
189    func portScan(domain: String) async -> CachedLookupResult<ServiceResult<[PortScanResult]>> {
190        await execute(
191            key: .domain(domain, .portScan),
192            extract: { payload in
193                guard case let .portScan(result) = payload else { return nil }
194                return result
195            },
196            operation: {
197                .portScan(await PortScanService.scanAll(domain: domain))
198            }
199        )
200    }
201
202    func email(domain: String, txtRecords: [DNSRecord]) async -> CachedLookupResult<ServiceResult<EmailSecurityResult>> {
203        await execute(
204            key: .domain(domain, .emailSecurity),
205            extract: { payload in
206                guard case let .email(result) = payload else { return nil }
207                return result
208            },
209            operation: {
210                .email(await EmailSecurityService.analyze(domain: domain, txtRecords: txtRecords))
211            }
212        )
213    }
214
215    func ptr(ip: String, resolverURLString: String) async -> CachedLookupResult<ServiceResult<String>> {
216        await execute(
217            key: .subject("\(resolverURLString)|\(ip)", .ptr),
218            extract: { payload in
219                guard case let .ptr(result) = payload else { return nil }
220                return result
221            },
222            operation: {
223                .ptr(await ReverseDNSService.lookup(ip: ip, resolverURLString: resolverURLString))
224            }
225        )
226    }
227
228    func ipGeolocation(ip: String) async -> CachedLookupResult<ServiceResult<IPGeolocation>> {
229        await execute(
230            key: .subject(ip, .ipGeolocation),
231            rateLimitBucket: .ipGeolocation,
232            extract: { payload in
233                guard case let .ipGeolocation(result) = payload else { return nil }
234                return result
235            },
236            operation: {
237                .ipGeolocation(await IPGeolocationService.lookup(ip: ip))
238            }
239        )
240    }
241
242    func suggestions(domain: String) async -> CachedLookupResult<[DomainSuggestionResult]> {
243        await execute(
244            key: .domain(domain, .suggestions),
245            extract: { payload in
246                guard case let .suggestions(result) = payload else { return nil }
247                return result
248            },
249            operation: {
250                .suggestions(await DomainAvailabilityService.suggestions(for: domain))
251            }
252        )
253    }
254
255    private func execute<T>(
256        key: RequestKey,
257        rateLimitBucket: RateLimitBucket? = nil,
258        extract: @escaping (CachedPayload) -> T?,
259        operation: @escaping @Sendable () async -> CachedPayload
260    ) async -> CachedLookupResult<T> {
261        if let cachedEntry = cache[key], cachedEntry.expiresAt > Date(), let value = extract(cachedEntry.payload) {
262            return CachedLookupResult(value: value, source: .cached)
263        }
264
265        if let task = inFlight[key], let value = extract(await task.value) {
266            return CachedLookupResult(value: value, source: .mixed)
267        }
268
269        let task = Task<CachedPayload, Never> {
270            if let rateLimitBucket {
271                await self.enforceRateLimit(for: rateLimitBucket)
272            }
273            return await operation()
274        }
275        inFlight[key] = task
276
277        let payload = await task.value
278        cache[key] = CacheEntry(payload: payload, expiresAt: Date().addingTimeInterval(ttl))
279        inFlight[key] = nil
280
281        guard let value = extract(payload) else {
282            fatalError("LookupRuntime payload extraction mismatch")
283        }
284
285        return CachedLookupResult(value: value, source: .live)
286    }
287
288    private func enforceRateLimit(for bucket: RateLimitBucket) async {
289        let now = Date()
290        if let nextAllowed = nextAllowedAt[bucket], nextAllowed > now {
291            let delay = nextAllowed.timeIntervalSince(now)
292            if delay > 0 {
293                try? await Task.sleep(nanoseconds: UInt64(delay * 1_000_000_000))
294            }
295        }
296        nextAllowedAt[bucket] = Date().addingTimeInterval(bucket.minimumSpacing)
297    }
298}