DomainDig/DNSLookupService.swift

ede15c586af71b8d6f2c190ffdad0324414cac06
domain-dig/DomainDig/DNSLookupService.swift history · blame · raw

90 lines · 3630 bytes

 1import Foundation
 2
 3struct DNSLookupService {
 4    private static let baseURL = "https://cloudflare-dns.com/dns-query"
 5
 6    static func lookup(domain: String, recordType: DNSRecordType) async throws -> [DNSRecord] {
 7        var components = URLComponents(string: baseURL)!
 8        components.queryItems = [
 9            URLQueryItem(name: "name", value: domain),
10            URLQueryItem(name: "type", value: String(recordType.queryType))
11        ]
12
13        var request = URLRequest(url: components.url!)
14        request.setValue("application/dns-json", forHTTPHeaderField: "Accept")
15
16        let (data, response) = try await URLSession.shared.data(for: request)
17
18        guard let httpResponse = response as? HTTPURLResponse,
19              httpResponse.statusCode == 200 else {
20            throw URLError(.badServerResponse)
21        }
22
23        let dnsResponse = try JSONDecoder().decode(CloudflareDNSResponse.self, from: data)
24
25        guard let answers = dnsResponse.Answer else {
26            return []
27        }
28
29        // Filter answers to only include the requested type
30        return answers
31            .filter { $0.type == recordType.queryType }
32            .map { answer in
33                let value = answer.data.trimmingCharacters(in: CharacterSet(charactersIn: "\""))
34                return DNSRecord(value: value, ttl: answer.TTL)
35            }
36    }
37
38    static func lookupAll(domain: String) async -> [DNSSection] {
39        // Each task returns (recordType, apex records, wildcard records).
40        typealias Result = (type: DNSRecordType, records: [DNSRecord], wildcard: [DNSRecord], error: String?)
41
42        // Record types that support wildcard queries
43        let wildcardTypes: Set<DNSRecordType> = [.A, .AAAA, .MX, .TXT]
44
45        return await withTaskGroup(of: Result.self, returning: [DNSSection].self) { group in
46            for recordType in DNSRecordType.allCases {
47                let shouldQueryWildcard = wildcardTypes.contains(recordType)
48                group.addTask {
49                    var apexRecords: [DNSRecord] = []
50                    var wildcardRecords: [DNSRecord] = []
51                    var lookupError: String?
52
53                    // Apex query
54                    do {
55                        apexRecords = try await lookup(domain: domain, recordType: recordType)
56                    } catch {
57                        lookupError = error.localizedDescription
58                    }
59
60                    // Wildcard query (only for applicable types, and only if apex didn't fail)
61                    if shouldQueryWildcard && lookupError == nil {
62                        do {
63                            wildcardRecords = try await lookup(domain: "*.\(domain)", recordType: recordType)
64                        } catch {
65                            // Wildcard failure is non-fatal; just leave empty
66                        }
67                    }
68
69                    return (recordType, apexRecords, wildcardRecords, lookupError)
70                }
71            }
72
73            var sections: [DNSSection] = []
74            for await result in group {
75                sections.append(DNSSection(
76                    recordType: result.type,
77                    records: result.records,
78                    wildcardRecords: result.wildcard,
79                    error: result.error
80                ))
81            }
82
83            // Sort to maintain consistent order
84            let order = DNSRecordType.allCases
85            return sections.sorted { a, b in
86                (order.firstIndex(of: a.recordType) ?? 0) < (order.firstIndex(of: b.recordType) ?? 0)
87            }
88        }
89    }
90}