Fixed NetShield selection in setting window

ProtonVPN · Jan 11, 2021 · c8d66ed · c8d66ed
1 parent 558d50d
commit c8d66ed
Showing 1 changed file with 11 additions and 18 deletions.
diff --git a/ProtonVPN/GeneralViewModel.swift b/ProtonVPN/GeneralViewModel.swift
@@ -108,9 +108,17 @@ class GeneralViewModel {
             break
         }
 
-        let userTier = (try? vpnGateway.userTier()) ?? 0
+        guard netShieldPropertyProvider.isUserEligibleForNetShield else {
+            let upgradeAlert = NetShieldRequiresUpgradeAlert(continueHandler: {
+                SafariService.openLink(url: CoreAppConstants.ProtonVpnLinks.accountDashboard)
+            })
+
+            self.alertService.push(alert: upgradeAlert)
+            viewController?.reloadView()
+            return
+        }
 
-        if !isConnected && !netShieldType.isUserTierTooLow(userTier) {
+        guard isConnected else {
             netShieldPropertyProvider.netShieldType = netShieldType
             viewController?.reloadView()
             return
@@ -124,21 +132,6 @@ class GeneralViewModel {
             self.viewController?.reloadView()
         })
 
-        switch netShieldType {
-        case .off, .level1:
-            self.alertService.push(alert: reconnectAlert)
-        case .level2:
-            guard netShieldType.isUserTierTooLow(userTier) else {
-                self.alertService.push(alert: reconnectAlert)
-                return
-            }
-
-            let upgradeAlert = NetShieldRequiresUpgradeAlert(continueHandler: {
-                SafariService.openLink(url: CoreAppConstants.ProtonVpnLinks.accountDashboard)
-            })
-
-            self.alertService.push(alert: upgradeAlert)
-            viewController?.reloadView()
-        }
+        self.alertService.push(alert: reconnectAlert)
     }
 }