krz/domain-dig

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