Repurpose compound domains with match/search policy

Author: keeshux

Sources/PartoutCore/Modules/DNSModule.swift

@@ -12,6 +12,11 @@ public struct DNSModule: Module, BuildableType, Hashable, Codable {
         case tls(hostname: String)
     }
 
+    public enum DomainPolicy: Hashable, Codable, Sendable {
+        case search
+        case match
+    }
+
     public static let moduleType = ModuleType("DNS")
 
     public let id: UniqueID
@@ -24,6 +29,8 @@ public struct DNSModule: Module, BuildableType, Hashable, Codable {
 
     public let searchDomains: [Address]?
 
+    public let domainPolicy: DomainPolicy?
+
     public let routesThroughVPN: Bool?
 
     fileprivate init(
@@ -32,13 +39,15 @@ public struct DNSModule: Module, BuildableType, Hashable, Codable {
         servers: [Address],
         domainName: Address?,
         searchDomains: [Address]?,
+        domainPolicy: DomainPolicy?,
         routesThroughVPN: Bool?
     ) {
         self.id = id
         self.protocolType = protocolType
         self.servers = servers
         self.domainName = domainName
         self.searchDomains = searchDomains
+        self.domainPolicy = domainPolicy
         self.routesThroughVPN = routesThroughVPN
     }
 
@@ -50,17 +59,29 @@ public struct DNSModule: Module, BuildableType, Hashable, Codable {
         switch protocolType {
         case .cleartext:
             break
-
         case .https(let url):
             builder.protocolType = .https
             builder.dohURL = url.absoluteString
-
         case .tls(let hostname):
             builder.protocolType = .tls
             builder.dotHostname = hostname
         }
-        builder.domainName = domainName?.rawValue
-        builder.searchDomains = searchDomains?.map(\.rawValue)
+        if let domainName {
+            builder.isFirstDomainPrimary = true
+            if let searchDomains {
+                if domainName == searchDomains.first {
+                    builder.domains = searchDomains.map(\.rawValue)
+                } else {
+                    builder.domains = [domainName.rawValue] + searchDomains.map(\.rawValue)
+                }
+            } else {
+                builder.domains = [domainName.rawValue]
+            }
+        } else if let searchDomains {
+            builder.isFirstDomainPrimary = false
+            builder.domains = searchDomains.map(\.rawValue)
+        }
+        builder.domainPolicy = domainPolicy
         builder.routesThroughVPN = routesThroughVPN
         return builder
     }
@@ -78,9 +99,11 @@ extension DNSModule {
 
         public var dotHostname: String
 
-        public var domainName: String?
+        public var domains: [String]?
+
+        public var domainPolicy: DomainPolicy?
 
-        public var searchDomains: [String]?
+        public var isFirstDomainPrimary: Bool
 
         public var routesThroughVPN: Bool?
 
@@ -94,17 +117,19 @@ extension DNSModule {
             servers: [String] = [],
             dohURL: String = "",
             dotHostname: String = "",
-            domainName: String? = nil,
-            searchDomains: [String]? = nil,
+            domains: [String]? = nil,
+            domainPolicy: DomainPolicy? = nil,
+            isFirstDomainPrimary: Bool = true,
             routesThroughVPN: Bool? = nil
         ) {
             self.id = id
             self.protocolType = protocolType
             self.servers = servers
             self.dohURL = dohURL
             self.dotHostname = dotHostname
-            self.domainName = domainName
-            self.searchDomains = searchDomains
+            self.domains = domains
+            self.domainPolicy = domainPolicy
+            self.isFirstDomainPrimary = isFirstDomainPrimary
             self.routesThroughVPN = routesThroughVPN
         }
 
@@ -118,25 +143,15 @@ extension DNSModule {
                 }
                 return addr
             }
-            let validDomainName = try domainName.flatMap {
+            let validDomains = try domains?.compactMap {
                 guard !$0.isEmpty else {
                     return nil as Address?
                 }
                 guard let addr = Address(rawValue: $0), !addr.isIPAddress else {
-                    throw PartoutError.invalidFields(["domainName": $0])
+                    throw PartoutError.invalidFields(["domains": $0])
                 }
                 return addr
             }
-            let validSearchDomains = try searchDomains?.compactMap {
-                guard !$0.isEmpty else {
-                    return nil as Address?
-                }
-                guard let addr = Address(rawValue: $0), !addr.isIPAddress else {
-                    throw PartoutError.invalidFields(["searchDomains": $0])
-                }
-                return addr
-            }
-
             let validProtocolType: ProtocolType
             switch protocolType {
             case .cleartext:
@@ -158,8 +173,9 @@ extension DNSModule {
                 id: id,
                 protocolType: validProtocolType,
                 servers: validServers,
-                domainName: validDomainName,
-                searchDomains: validSearchDomains,
+                domainName: isFirstDomainPrimary ? validDomains?.first : nil,
+                searchDomains: validDomains,
+                domainPolicy: domainPolicy,
                 routesThroughVPN: routesThroughVPN
             )
         }

Tests/PartoutCoreTests/DNSModuleTests.swift

@@ -39,6 +39,24 @@ struct DNSModuleTests {
         #expect(sut == module.builder())
     }
 
+    @Test(arguments: [true, false])
+    func givenDomains_whenRebuild_thenIsRestored(isFirstDomainPrimary: Bool) throws {
+        let sut = DNSModule.Builder(
+            protocolType: .cleartext,
+            servers: ["1.2.3.4"],
+            domains: ["primary.example.com", "search.example.com"],
+            isFirstDomainPrimary: isFirstDomainPrimary
+        )
+        let module = try sut.build()
+        let rebuilt = module.builder()
+
+        #expect(module.domainName?.rawValue == (isFirstDomainPrimary ? "primary.example.com" : nil))
+        #expect(module.searchDomains?.map(\.rawValue) == ["primary.example.com", "search.example.com"])
+        #expect(rebuilt.domains == sut.domains)
+        #expect(rebuilt.isFirstDomainPrimary == sut.isFirstDomainPrimary)
+        #expect(rebuilt == sut)
+    }
+
     @Test
     func givenHTTPSWithoutURL_whenBuild_thenFails() {
         let sut = DNSModule.Builder(

scripts/openapi.yaml

@@ -17,6 +17,8 @@ components:
       properties:
         domainName:
           "$ref": "#/components/schemas/Address"
+        domainPolicy:
+          title: DomainPolicy
         id:
           "$ref": "#/components/schemas/UniqueID"
         protocolType: