Merge branch 'user-shouldnt-be-able-to-select-the-same-entry-location-or-ios-748'

6 files changed, 81 insertions, 50 deletions

diff --git a/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsDataSource.swift b/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsDataSource.swift
index a180b206aa..d2325e3550 100644
--- a/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsDataSource.swift
+++ b/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsDataSource.swift
@@ -161,10 +161,6 @@ extension AddLocationsDataSource {
     func nodeShouldBeSelected(_ node: LocationNode) -> Bool {
         customListLocationNode.children.contains(node)
     }
-
-    func excludedRelayTitle(_ node: LocationNode) -> String? {
-        nil // N/A
-    }
 }
 
 // MARK: - Toggle selection in table view
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationCell.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationCell.swift
index 1bb37ba097..69b6fcc4fd 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationCell.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationCell.swift
@@ -209,7 +209,6 @@ class LocationCell: UITableViewCell {
 
     private func updateDisabled(_ isDisabled: Bool) {
         locationLabel.alpha = isDisabled ? 0.2 : 1
-        collapseButton.alpha = isDisabled ? 0.2 : 1
 
         if isDisabled {
             accessibilityTraits.insert(.notEnabled)
@@ -339,14 +338,14 @@ extension LocationCell {
             }
         }
 
-        setExcludedRelayTitle(item.excludedRelayTitle)
         setBehavior(behavior)
     }
 
-    private func setExcludedRelayTitle(_ title: String?) {
-        if let title {
-            locationLabel.text! += " (\(title))"
-            updateDisabled(true)
+    func setExcluded(relayTitle: String? = nil) {
+        updateDisabled(true)
+
+        if let relayTitle {
+            locationLabel.text! += " (\(relayTitle))"
         }
     }
 
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationCellViewModel.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationCellViewModel.swift
index ca24c60552..355e7a6fb1 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationCellViewModel.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationCellViewModel.swift
@@ -31,11 +31,7 @@ struct LocationCellViewModel: Hashable {
 }
 
 extension [LocationCellViewModel] {
-    mutating func addSubNodes(
-        from item: LocationCellViewModel,
-        at indexPath: IndexPath,
-        excludedRelayTitleCallback: ((LocationNode) -> String?)?
-    ) {
+    mutating func addSubNodes(from item: LocationCellViewModel, at indexPath: IndexPath) {
         let section = LocationSection.allCases[indexPath.section]
         let row = indexPath.row + 1
 
@@ -44,8 +40,7 @@ extension [LocationCellViewModel] {
                 section: section,
                 node: $0,
                 indentationLevel: item.indentationLevel + 1,
-                isSelected: false,
-                excludedRelayTitle: excludedRelayTitleCallback?($0)
+                isSelected: false
             )
         }
 
@@ -65,3 +60,53 @@ extension [LocationCellViewModel] {
         }
     }
 }
+
+extension LocationCellViewModel {
+    /* Exclusion of other locations in the same node tree as the currently excluded location
+     happens when there are no more hosts in that tree that can be selected.
+     We check this by doing the following, in order:
+
+     1. Count host names in the tree. More than one means that there are other locations than
+     the excluded one for the relay selector to choose from. No exlusion.
+
+     2. Count host names in the excluded node. More than one means that there are multiple
+     locations for the relay selector to choose from. No exlusion.
+
+     3. Check existance of a location in the tree that match the currently excluded location.
+     No match means no exclusion.
+     */
+    func shouldExcludeLocation(_ excludedLocation: LocationCellViewModel?) -> Bool {
+        guard let excludedLocation else {
+            return false
+        }
+
+        var proxyNode = RootLocationNode(children: [node])
+        let allLocations = Set(proxyNode.flattened.flatMap { $0.locations })
+        let hostCount = allLocations.filter { location in
+            if case .hostname = location { true } else { false }
+        }.count
+
+        // If the there are more than one selectable relay in the current node we don't need
+        // show this in the location tree and can return early.
+        guard hostCount == 1 else { return false }
+
+        let proxyExcludedNode = RootLocationNode(children: [excludedLocation.node])
+        let allExcludedLocations = Set(proxyExcludedNode.flattened.flatMap { $0.locations })
+        let excludedHostCount = allExcludedLocations.filter { location in
+            if case .hostname = location { true } else { false }
+        }.count
+
+        // If the there are more than one selectable relay in the excluded node we don't need
+        // show this in the location tree and can return early.
+        guard excludedHostCount == 1 else { return false }
+
+        var containsExcludedLocation = false
+        if allLocations.contains(where: { allExcludedLocations.contains($0) }) {
+            containsExcludedLocation = true
+        }
+
+        // If the tree doesn't contain the excluded node we do nothing, otherwise the
+        // required conditions have now all been met.
+        return containsExcludedLocation
+    }
+}
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift
index 3dd3f42795..4d68eccbb5 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift
@@ -142,10 +142,8 @@ final class LocationDataSource:
     func setSelectedRelays(_ selectedRelays: RelaySelection) {
         selectedLocation = mapSelection(from: selectedRelays.selected)
 
-        if selectedRelays.hasExcludedRelay {
-            excludedLocation = mapSelection(from: selectedRelays.excluded)
-            excludedLocation?.excludedRelayTitle = selectedRelays.excludedTitle
-        }
+        excludedLocation = mapSelection(from: selectedRelays.excluded)
+        excludedLocation?.excludedRelayTitle = selectedRelays.excludedTitle
 
         tableView.reloadData()
     }
@@ -246,9 +244,24 @@ final class LocationDataSource:
     }
 
     override func tableView(_ tableView: UITableView, cellForRowAt indexPath: IndexPath) -> UITableViewCell {
-        // swiftlint:disable:next force_cast
-        let cell = super.tableView(tableView, cellForRowAt: indexPath) as! LocationCell
+        let cell = super.tableView(tableView, cellForRowAt: indexPath)
+        guard let cell = cell as? LocationCell, let item = itemIdentifier(for: indexPath) else {
+            return cell
+        }
+
         cell.delegate = self
+
+        if item.shouldExcludeLocation(excludedLocation) {
+            // Only host locations should have an excluded title. Since custom list nodes contain
+            // all locations of all child nodes, its first location could possibly be a host.
+            // Therefore we need to check for that as well.
+            if case .hostname = item.node.locations.first, !(item.node is CustomListLocationNode) {
+                cell.setExcluded(relayTitle: excludedLocation?.excludedRelayTitle)
+            } else {
+                cell.setExcluded()
+            }
+        }
+
         return cell
     }
 }
@@ -263,14 +276,6 @@ extension LocationDataSource {
     func nodeShouldBeSelected(_ node: LocationNode) -> Bool {
         false // N/A
     }
-
-    func excludedRelayTitle(_ node: LocationNode) -> String? {
-        if nodeMatchesExcludedLocation(node) {
-            excludedLocation?.excludedRelayTitle
-        } else {
-            nil
-        }
-    }
 }
 
 extension LocationDataSource: UITableViewDelegate {
@@ -307,7 +312,7 @@ extension LocationDataSource: UITableViewDelegate {
 
     func tableView(_ tableView: UITableView, shouldHighlightRowAt indexPath: IndexPath) -> Bool {
         guard let item = itemIdentifier(for: indexPath) else { return false }
-        return !nodeMatchesExcludedLocation(item.node) && item.node.isActive
+        return !item.shouldExcludeLocation(excludedLocation) && item.node.isActive
     }
 
     func tableView(_ tableView: UITableView, indentationLevelForRowAt indexPath: IndexPath) -> Int {
@@ -360,9 +365,7 @@ extension LocationDataSource: LocationCellDelegate {
         guard let indexPath = tableView.indexPath(for: cell),
               let item = itemIdentifier(for: indexPath) else { return }
 
-        let items = toggledItems(for: cell, excludedRelayTitleCallback: { node in
-            self.excludedRelayTitle(node)
-        })
+        let items = toggledItems(for: cell)
 
         updateDataSnapshot(with: items, reloadExisting: true, completion: {
             self.scroll(to: item, animated: true)
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationDiffableDataSourceProtocol.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationDiffableDataSourceProtocol.swift
index 59f7e12d21..0450be0a81 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationDiffableDataSourceProtocol.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationDiffableDataSourceProtocol.swift
@@ -14,7 +14,6 @@ protocol LocationDiffableDataSourceProtocol: UITableViewDiffableDataSource<Locat
     var sections: [LocationSection] { get }
     func nodeShowsChildren(_ node: LocationNode) -> Bool
     func nodeShouldBeSelected(_ node: LocationNode) -> Bool
-    func excludedRelayTitle(_ node: LocationNode) -> String?
 }
 
 extension LocationDiffableDataSourceProtocol {
@@ -40,10 +39,7 @@ extension LocationDiffableDataSourceProtocol {
         }
     }
 
-    func toggledItems(
-        for cell: LocationCell,
-        excludedRelayTitleCallback: ((LocationNode) -> String?)? = nil
-    ) -> [[LocationCellViewModel]] {
+    func toggledItems(for cell: LocationCell) -> [[LocationCellViewModel]] {
         guard let indexPath = tableView.indexPath(for: cell),
               let item = itemIdentifier(for: indexPath) else { return [[]] }
 
@@ -54,7 +50,7 @@ extension LocationDiffableDataSourceProtocol {
         item.node.showsChildren = !isExpanded
 
         if !isExpanded {
-            locationList.addSubNodes(from: item, at: indexPath, excludedRelayTitleCallback: excludedRelayTitleCallback)
+            locationList.addSubNodes(from: item, at: indexPath)
         } else {
             locationList.removeSubNodes(from: item.node)
         }
@@ -103,8 +99,7 @@ extension LocationDiffableDataSourceProtocol {
                     section: section,
                     node: childNode,
                     indentationLevel: indentationLevel,
-                    isSelected: nodeShouldBeSelected(childNode),
-                    excludedRelayTitle: excludedRelayTitle(childNode)
+                    isSelected: nodeShouldBeSelected(childNode)
                 )
             )
 
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/RelaySelection.swift b/ios/MullvadVPN/View controllers/SelectLocation/RelaySelection.swift
index 5ae7354fe8..85b185a574 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/RelaySelection.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/RelaySelection.swift
@@ -12,11 +12,4 @@ struct RelaySelection {
     var selected: UserSelectedRelays?
     var excluded: UserSelectedRelays?
     var excludedTitle: String?
-
-    var hasExcludedRelay: Bool {
-        if excluded?.locations.count == 1, case .hostname = excluded?.locations.first {
-            return true
-        }
-        return false
-    }
 }