refactor: revertDns

mrFq1 · Sep 21, 2024 · 993ff0a · 993ff0a
1 parent ff60676
commit 993ff0a
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 29 deletions.
diff --git a/ProxyConfigHelper/MetaDNS.swift b/ProxyConfigHelper/MetaDNS.swift
@@ -12,8 +12,8 @@ import SystemConfiguration
 class MetaDNS: NSObject {
 	private var customDNS = "8.8.8.8"
 
-	static let savedDnsKey = "SavedSystemDNSs"
-    var savedDns = [String: [String]]()
+	static let savedDNSKey = "ProxyConfigHelper.SavedSystemDNSs"
+    var savedDNS = [String: [String]]()
     let authRef: AuthorizationRef
 
     override init() {
@@ -31,9 +31,9 @@ class MetaDNS: NSObject {
         }
 
 
-		if let data = UserDefaults.standard.data(forKey: MetaDNS.savedDnsKey),
+		if let data = UserDefaults.standard.data(forKey: MetaDNS.savedDNSKey),
 		   let saved = try? JSONDecoder().decode([String: [String]].self, from: data) {
-			self.savedDns = saved
+			self.savedDNS = saved
 		}
 
         authRef = auth!
@@ -49,39 +49,31 @@ class MetaDNS: NSObject {
 		customDNS = dns
 	}
 
-    @objc func updateDns() {
+    @objc func hijackDNS() {
 		let dns = getAllDns()
+		let hijacked = dns.allSatisfy {
+			$0.value.count == 1 && $0.value[0] == customDNS
+		}
+
+		guard !hijacked else { return }
 
-		if savedDns.isEmpty {
-			dns.forEach {
-				if $0.value.count == 1,
-				   $0.value[0] == customDNS {
-					if savedDns[$0.key] == nil {
-						savedDns[$0.key] = []
-					} else {
-						// ignore save
-					}
-				} else {
-					savedDns[$0.key] = $0.value
-				}
-			}
-			if !savedDns.isEmpty,
-				let data = try? JSONEncoder().encode(savedDns) {
-				UserDefaults.standard.set(data, forKey: MetaDNS.savedDnsKey)
-			}
+		savedDNS = dns
+		if let data = try? JSONEncoder().encode(savedDNS) {
+			UserDefaults.standard.set(data, forKey: MetaDNS.savedDNSKey)
 		}
-        
+
         let dnsDic = dns.reduce(into: [:]) {
             $0[$1.key] = [customDNS]
         }
 
         updateDNSConfigure(dnsDic)
     }
 
-    @objc func revertDns() {
-        updateDNSConfigure(savedDns)
-        savedDns.removeAll()
-		UserDefaults.standard.removeObject(forKey: MetaDNS.savedDnsKey)
+    @objc func revertDNS() {
+		guard savedDNS.count > 0 else { return }
+        updateDNSConfigure(savedDNS)
+        savedDNS.removeAll()
+		UserDefaults.standard.removeObject(forKey: MetaDNS.savedDNSKey)
     }
 
     func getAllDns() -> [String: [String]] {

diff --git a/ProxyConfigHelper/ProxyConfigHelper.swift b/ProxyConfigHelper/ProxyConfigHelper.swift
@@ -136,9 +136,9 @@ extension ProxyConfigHelper: ProxyConfigRemoteProcessProtocol {
 		DispatchQueue.main.async {
 			self.metaDNS.setCustomDNS(dns)
 			if state {
-				self.metaDNS.updateDns()
+				self.metaDNS.hijackDNS()
 			} else {
-				self.metaDNS.revertDns()
+				self.metaDNS.revertDNS()
 			}
 			self.metaDNS.flushDnsCache()
 		}