From 993ff0a1465d9eec35cbcf0b1520d19d096f9a7c Mon Sep 17 00:00:00 2001
From: mrFq1 <1xxbx0il0@mozmail.com>
Date: Sat, 21 Sep 2024 15:53:47 +0800
Subject: [PATCH] refactor: revertDns

---
 ProxyConfigHelper/MetaDNS.swift           | 46 ++++++++++-------------
 ProxyConfigHelper/ProxyConfigHelper.swift |  4 +-
 2 files changed, 21 insertions(+), 29 deletions(-)

diff --git a/ProxyConfigHelper/MetaDNS.swift b/ProxyConfigHelper/MetaDNS.swift
index 5fb51e30c..2b75dc4d6 100644
--- 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,28 +49,19 @@ 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]
         }
@@ -78,10 +69,11 @@ class MetaDNS: NSObject {
         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
index 0cc7f162a..4a3083954 100644
--- 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()
 		}