Merge branch 'intercept-back-button-when-leaving-an-unsaved-custom-list-ios-584'

16 files changed, 263 insertions, 86 deletions

diff --git a/ios/MullvadVPN.xcodeproj/project.pbxproj b/ios/MullvadVPN.xcodeproj/project.pbxproj
index 7879654647..15d379e53f 100644
--- a/ios/MullvadVPN.xcodeproj/project.pbxproj
+++ b/ios/MullvadVPN.xcodeproj/project.pbxproj
@@ -540,6 +540,7 @@
 		7A6F2FAB2AFD3097006D0856 /* CustomDNSCellFactory.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7A6F2FAA2AFD3097006D0856 /* CustomDNSCellFactory.swift */; };
 		7A6F2FAD2AFD3DA7006D0856 /* CustomDNSViewController.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7A6F2FAC2AFD3DA7006D0856 /* CustomDNSViewController.swift */; };
 		7A6F2FAF2AFE36E7006D0856 /* VPNSettingsInfoButtonItem.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7A6F2FAE2AFE36E7006D0856 /* VPNSettingsInfoButtonItem.swift */; };
+		7A7907332BC0280A00B61F81 /* InterceptibleNavigationController.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7A7907322BC0280A00B61F81 /* InterceptibleNavigationController.swift */; };
 		7A7AD28D29DC677800480EF1 /* FirstTimeLaunch.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7A7AD28C29DC677800480EF1 /* FirstTimeLaunch.swift */; };
 		7A818F1F29F0305800C7F0F4 /* RootConfiguration.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7A818F1E29F0305800C7F0F4 /* RootConfiguration.swift */; };
 		7A83A0C62B29A750008B5CE7 /* APIAccessMethodsTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7A83A0C52B29A750008B5CE7 /* APIAccessMethodsTests.swift */; };
@@ -1805,6 +1806,7 @@
 		7A6F2FAA2AFD3097006D0856 /* CustomDNSCellFactory.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = CustomDNSCellFactory.swift; sourceTree = "<group>"; };
 		7A6F2FAC2AFD3DA7006D0856 /* CustomDNSViewController.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = CustomDNSViewController.swift; sourceTree = "<group>"; };
 		7A6F2FAE2AFE36E7006D0856 /* VPNSettingsInfoButtonItem.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = VPNSettingsInfoButtonItem.swift; sourceTree = "<group>"; };
+		7A7907322BC0280A00B61F81 /* InterceptibleNavigationController.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = InterceptibleNavigationController.swift; sourceTree = "<group>"; };
 		7A7AD28C29DC677800480EF1 /* FirstTimeLaunch.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = FirstTimeLaunch.swift; sourceTree = "<group>"; };
 		7A818F1E29F0305800C7F0F4 /* RootConfiguration.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = RootConfiguration.swift; sourceTree = "<group>"; };
 		7A83A0C52B29A750008B5CE7 /* APIAccessMethodsTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = APIAccessMethodsTests.swift; sourceTree = "<group>"; };
@@ -2694,6 +2696,7 @@
 				58138E60294871C600684F0C /* DeviceDataThrottling.swift */,
 				7A7AD28C29DC677800480EF1 /* FirstTimeLaunch.swift */,
 				582AE30F2440A6CA00E6733A /* InputTextFormatter.swift */,
+				7A7907322BC0280A00B61F81 /* InterceptibleNavigationController.swift */,
 				58DFF7D12B0256A300F864E0 /* MarkdownStylingOptions.swift */,
 				58CC40EE24A601900019D96E /* ObserverList.swift */,
 			);
@@ -5463,6 +5466,7 @@
 				7A42DEC92A05164100B209BE /* SettingsInputCell.swift in Sources */,
 				5803B4B22940A48700C23744 /* TunnelStore.swift in Sources */,
 				586A950F29012BEE007BAF2B /* AddressCacheTracker.swift in Sources */,
+				7A7907332BC0280A00B61F81 /* InterceptibleNavigationController.swift in Sources */,
 				F02F41A02B9723AF00625A4F /* AddLocationsViewController.swift in Sources */,
 				587B753D2666468F00DEF7E9 /* NotificationController.swift in Sources */,
 			);
diff --git a/ios/MullvadVPN/Classes/InterceptibleNavigationController.swift b/ios/MullvadVPN/Classes/InterceptibleNavigationController.swift
new file mode 100644
index 0000000000..f93795c591
--- /dev/null
+++ b/ios/MullvadVPN/Classes/InterceptibleNavigationController.swift
@@ -0,0 +1,38 @@
+//
+//  InterceptibleNavigationController.swift
+//  MullvadVPN
+//
+//  Created by Jon Petersson on 2024-04-05.
+//  Copyright © 2024 Mullvad VPN AB. All rights reserved.
+//
+
+import UIKit
+
+class InterceptibleNavigationController: CustomNavigationController {
+    var shouldPopViewController: ((UIViewController) -> Bool)?
+    var shouldPopToViewController: ((UIViewController) -> Bool)?
+
+    // Called when popping the topmost view controller in the stack, eg. by pressing a navigation
+    // bar back button.
+    @discardableResult
+    override func popViewController(animated: Bool) -> UIViewController? {
+        guard let viewController = viewControllers.last else { return nil }
+
+        if shouldPopViewController?(viewController) ?? true {
+            return super.popViewController(animated: animated)
+        } else {
+            return nil
+        }
+    }
+
+    // Called when popping to a specific view controller, eg. by long pressing a navigation bar
+    // back button (revealing a navigation menu) and selecting a destination view controller.
+    @discardableResult
+    override func popToViewController(_ viewController: UIViewController, animated: Bool) -> [UIViewController]? {
+        if shouldPopToViewController?(viewController) ?? true {
+            return super.popToViewController(viewController, animated: animated)
+        } else {
+            return nil
+        }
+    }
+}
diff --git a/ios/MullvadVPN/Coordinators/CustomLists/AddCustomListCoordinator.swift b/ios/MullvadVPN/Coordinators/CustomLists/AddCustomListCoordinator.swift
index bbbf45ad54..6ef8d044fd 100644
--- a/ios/MullvadVPN/Coordinators/CustomLists/AddCustomListCoordinator.swift
+++ b/ios/MullvadVPN/Coordinators/CustomLists/AddCustomListCoordinator.swift
@@ -84,17 +84,10 @@ extension AddCustomListCoordinator: CustomListViewControllerDelegate {
         let coordinator = AddLocationsCoordinator(
             navigationController: navigationController,
             nodes: nodes,
-            customList: list
+            subject: subject
         )
 
-        coordinator.didFinish = { [weak self] locationsCoordinator, customList in
-            guard let self else { return }
-            subject.send(CustomListViewModel(
-                id: customList.id,
-                name: customList.name,
-                locations: customList.locations,
-                tableSections: subject.value.tableSections
-            ))
+        coordinator.didFinish = { locationsCoordinator in
             locationsCoordinator.removeFromParent()
         }
 
diff --git a/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsCoordinator.swift b/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsCoordinator.swift
index feb5bd415e..1634a24e80 100644
--- a/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsCoordinator.swift
+++ b/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsCoordinator.swift
@@ -6,6 +6,7 @@
 //  Copyright © 2024 Mullvad VPN AB. All rights reserved.
 //
 
+import Combine
 import MullvadSettings
 import MullvadTypes
 import Routing
@@ -14,9 +15,9 @@ import UIKit
 class AddLocationsCoordinator: Coordinator, Presentable, Presenting {
     private let navigationController: UINavigationController
     private let nodes: [LocationNode]
-    private var customList: CustomList
+    private var subject: CurrentValueSubject<CustomListViewModel, Never>
 
-    var didFinish: ((AddLocationsCoordinator, CustomList) -> Void)?
+    var didFinish: ((AddLocationsCoordinator) -> Void)?
 
     var presentedViewController: UIViewController {
         navigationController
@@ -25,17 +26,17 @@ class AddLocationsCoordinator: Coordinator, Presentable, Presenting {
     init(
         navigationController: UINavigationController,
         nodes: [LocationNode],
-        customList: CustomList
+        subject: CurrentValueSubject<CustomListViewModel, Never>
     ) {
         self.navigationController = navigationController
         self.nodes = nodes
-        self.customList = customList
+        self.subject = subject
     }
 
     func start() {
         let controller = AddLocationsViewController(
             allLocationsNodes: nodes,
-            customList: customList
+            subject: subject
         )
         controller.delegate = self
 
@@ -51,11 +52,7 @@ class AddLocationsCoordinator: Coordinator, Presentable, Presenting {
 }
 
 extension AddLocationsCoordinator: AddLocationsViewControllerDelegate {
-    func didUpdateSelectedLocations(locations: [RelayLocation]) {
-        customList.locations = locations
-    }
-
     func didBack() {
-        didFinish?(self, customList)
+        didFinish?(self)
     }
 }
diff --git a/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsDataSource.swift b/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsDataSource.swift
index 048d3e51fa..0af9f14d48 100644
--- a/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsDataSource.swift
+++ b/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsDataSource.swift
@@ -6,6 +6,7 @@
 //  Copyright © 2024 Mullvad VPN AB. All rights reserved.
 //
 
+import Combine
 import MullvadSettings
 import MullvadTypes
 import UIKit
@@ -15,20 +16,21 @@ class AddLocationsDataSource:
     LocationDiffableDataSourceProtocol {
     private var customListLocationNode: CustomListLocationNode
     private let nodes: [LocationNode]
-    var didUpdateCustomList: ((CustomListLocationNode) -> Void)?
+    private let subject: CurrentValueSubject<CustomListViewModel, Never>
     let tableView: UITableView
     let sections: [LocationSection]
 
     init(
         tableView: UITableView,
         allLocationNodes: [LocationNode],
-        customList: CustomList
+        subject: CurrentValueSubject<CustomListViewModel, Never>
     ) {
         self.tableView = tableView
         self.nodes = allLocationNodes
+        self.subject = subject
 
         self.customListLocationNode = CustomListLocationNodeBuilder(
-            customList: customList,
+            customList: subject.value.customList,
             allLocations: self.nodes
         ).customListLocationNode
 
@@ -51,10 +53,12 @@ class AddLocationsDataSource:
         reloadWithSelectedLocations()
     }
 
+    // Called from `LocationDiffableDataSourceProtocol`.
     func nodeShowsChildren(_ node: LocationNode) -> Bool {
         isLocationInCustomList(node: node)
     }
 
+    // Called from `LocationDiffableDataSourceProtocol`.
     func nodeShouldBeSelected(_ node: LocationNode) -> Bool {
         customListLocationNode.children.contains(node)
     }
@@ -149,7 +153,10 @@ extension AddLocationsDataSource: LocationCellDelegate {
             customListLocationNode.remove(selectedLocation: item.node, with: locationList)
         }
         updateDataSnapshot(with: [locationList], completion: {
-            self.didUpdateCustomList?(self.customListLocationNode)
+            let locations = self.customListLocationNode.children.reduce([]) { partialResult, locationNode in
+                partialResult + locationNode.locations
+            }
+            self.subject.value.locations = locations
         })
     }
 }
diff --git a/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsViewController.swift b/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsViewController.swift
index c728982fdb..1dbd7ac7ae 100644
--- a/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsViewController.swift
+++ b/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsViewController.swift
@@ -6,19 +6,19 @@
 //  Copyright © 2024 Mullvad VPN AB. All rights reserved.
 //
 
+import Combine
 import MullvadSettings
 import MullvadTypes
 import UIKit
 
 protocol AddLocationsViewControllerDelegate: AnyObject {
-    func didUpdateSelectedLocations(locations: [RelayLocation])
     func didBack()
 }
 
 class AddLocationsViewController: UIViewController {
     private var dataSource: AddLocationsDataSource?
     private let nodes: [LocationNode]
-    private let customList: CustomList
+    private let subject: CurrentValueSubject<CustomListViewModel, Never>
 
     weak var delegate: AddLocationsViewControllerDelegate?
     private let tableView: UITableView = {
@@ -33,10 +33,10 @@ class AddLocationsViewController: UIViewController {
 
     init(
         allLocationsNodes: [LocationNode],
-        customList: CustomList
+        subject: CurrentValueSubject<CustomListViewModel, Never>
     ) {
         self.nodes = allLocationsNodes
-        self.customList = customList
+        self.subject = subject
         super.init(nibName: nil, bundle: nil)
     }
 
@@ -70,17 +70,8 @@ class AddLocationsViewController: UIViewController {
         dataSource = AddLocationsDataSource(
             tableView: tableView,
             allLocationNodes: nodes.copy(),
-            customList: customList
+            subject: subject
         )
-
-        dataSource?.didUpdateCustomList = { [weak self] customListLocationNode in
-            guard let self else { return }
-            delegate?.didUpdateSelectedLocations(
-                locations: customListLocationNode.children.reduce([]) { partialResult, locationNode in
-                    partialResult + locationNode.locations
-                }
-            )
-        }
     }
 }
 
diff --git a/ios/MullvadVPN/Coordinators/CustomLists/CustomListViewController.swift b/ios/MullvadVPN/Coordinators/CustomLists/CustomListViewController.swift
index 4e5891658d..a3a7518af8 100644
--- a/ios/MullvadVPN/Coordinators/CustomLists/CustomListViewController.swift
+++ b/ios/MullvadVPN/Coordinators/CustomLists/CustomListViewController.swift
@@ -27,6 +27,14 @@ class CustomListViewController: UIViewController {
     private let alertPresenter: AlertPresenter
     private var validationErrors: Set<CustomListFieldValidationError> = []
 
+    private var persistedCustomList: CustomList? {
+        return interactor.fetchAll().first(where: { $0.id == subject.value.id })
+    }
+
+    var hasUnsavedChanges: Bool {
+        persistedCustomList != subject.value.customList
+    }
+
     private lazy var cellConfiguration: CustomListCellConfiguration = {
         CustomListCellConfiguration(tableView: tableView, subject: subject)
     }()
@@ -75,12 +83,12 @@ class CustomListViewController: UIViewController {
     override func viewDidLoad() {
         super.viewDidLoad()
 
+        navigationItem.rightBarButtonItem = saveBarButton
         view.directionalLayoutMargins = UIMetrics.contentLayoutMargins
         view.backgroundColor = .secondaryColor
         isModalInPresentation = true
 
         addSubviews()
-        configureNavigationItem()
         configureDataSource()
         configureTableView()
 
@@ -90,10 +98,6 @@ class CustomListViewController: UIViewController {
         }.store(in: &cancellables)
     }
 
-    private func configureNavigationItem() {
-        navigationItem.rightBarButtonItem = saveBarButton
-    }
-
     private func configureTableView() {
         tableView.delegate = dataSourceConfiguration
         tableView.backgroundColor = .secondaryColor
diff --git a/ios/MullvadVPN/Coordinators/CustomLists/CustomListViewModel.swift b/ios/MullvadVPN/Coordinators/CustomLists/CustomListViewModel.swift
index b41d52d2f5..10b9e13592 100644
--- a/ios/MullvadVPN/Coordinators/CustomLists/CustomListViewModel.swift
+++ b/ios/MullvadVPN/Coordinators/CustomLists/CustomListViewModel.swift
@@ -18,4 +18,9 @@ struct CustomListViewModel {
     var customList: CustomList {
         CustomList(id: id, name: name, locations: locations)
     }
+
+    mutating func update(with list: CustomList) {
+        name = list.name
+        locations = list.locations
+    }
 }
diff --git a/ios/MullvadVPN/Coordinators/CustomLists/EditCustomListCoordinator.swift b/ios/MullvadVPN/Coordinators/CustomLists/EditCustomListCoordinator.swift
index 5545f1bc95..2e3c8c9a0c 100644
--- a/ios/MullvadVPN/Coordinators/CustomLists/EditCustomListCoordinator.swift
+++ b/ios/MullvadVPN/Coordinators/CustomLists/EditCustomListCoordinator.swift
@@ -21,12 +21,16 @@ class EditCustomListCoordinator: Coordinator, Presentable, Presenting {
     let customList: CustomList
     let nodes: [LocationNode]
     let subject: CurrentValueSubject<CustomListViewModel, Never>
+    private lazy var alertPresenter: AlertPresenter = {
+        AlertPresenter(context: self)
+    }()
 
     var presentedViewController: UIViewController {
         navigationController
     }
 
     var didFinish: ((EditCustomListCoordinator, FinishAction, CustomList) -> Void)?
+    var didCancel: ((EditCustomListCoordinator) -> Void)?
 
     init(
         navigationController: UINavigationController,
@@ -50,7 +54,7 @@ class EditCustomListCoordinator: Coordinator, Presentable, Presenting {
         let controller = CustomListViewController(
             interactor: customListInteractor,
             subject: subject,
-            alertPresenter: AlertPresenter(context: self)
+            alertPresenter: alertPresenter
         )
         controller.delegate = self
 
@@ -61,7 +65,77 @@ class EditCustomListCoordinator: Coordinator, Presentable, Presenting {
             comment: ""
         )
 
+        navigationController.interactivePopGestureRecognizer?.delegate = self
         navigationController.pushViewController(controller, animated: true)
+
+        guard let interceptibleNavigationController = navigationController as? InterceptibleNavigationController else {
+            return
+        }
+
+        interceptibleNavigationController.shouldPopViewController = { [weak self] viewController in
+            guard
+                let self,
+                let customListViewController = viewController as? CustomListViewController,
+                customListViewController.hasUnsavedChanges
+            else { return true }
+
+            presentUnsavedChangesDialog()
+            return false
+        }
+
+        interceptibleNavigationController.shouldPopToViewController = { [weak self] viewController in
+            guard
+                let self,
+                let customListViewController = viewController as? CustomListViewController,
+                customListViewController.hasUnsavedChanges
+            else { return true }
+
+            presentUnsavedChangesDialog()
+            return false
+        }
+    }
+
+    private func presentUnsavedChangesDialog() {
+        let message = NSMutableAttributedString(
+            markdownString: NSLocalizedString(
+                "CUSTOM_LISTS_UNSAVED_CHANGES_PROMPT",
+                tableName: "CustomLists",
+                value: "You have unsaved changes.",
+                comment: ""
+            ),
+            options: MarkdownStylingOptions(font: .preferredFont(forTextStyle: .body))
+        )
+
+        let presentation = AlertPresentation(
+            id: "api-custom-lists-unsaved-changes-alert",
+            icon: .alert,
+            attributedMessage: message,
+            buttons: [
+                AlertAction(
+                    title: NSLocalizedString(
+                        "CUSTOM_LISTS_DISCARD_CHANGES_BUTTON",
+                        tableName: "CustomLists",
+                        value: "Discard changes",
+                        comment: ""
+                    ),
+                    style: .destructive,
+                    handler: {
+                        self.didCancel?(self)
+                    }
+                ),
+                AlertAction(
+                    title: NSLocalizedString(
+                        "CUSTOM_LISTS_BACK_TO_EDITING_BUTTON",
+                        tableName: "CustomLists",
+                        value: "Back to editing",
+                        comment: ""
+                    ),
+                    style: .default
+                ),
+            ]
+        )
+
+        alertPresenter.showAlert(presentation: presentation, animated: true)
     }
 }
 
@@ -78,17 +152,10 @@ extension EditCustomListCoordinator: CustomListViewControllerDelegate {
         let coordinator = EditLocationsCoordinator(
             navigationController: navigationController,
             nodes: nodes,
-            customList: list
+            subject: subject
         )
 
-        coordinator.didFinish = { [weak self] locationsCoordinator, customList in
-            guard let self else { return }
-            subject.send(CustomListViewModel(
-                id: customList.id,
-                name: customList.name,
-                locations: customList.locations,
-                tableSections: subject.value.tableSections
-            ))
+        coordinator.didFinish = { locationsCoordinator in
             locationsCoordinator.removeFromParent()
         }
 
@@ -97,3 +164,19 @@ extension EditCustomListCoordinator: CustomListViewControllerDelegate {
         addChild(coordinator)
     }
 }
+
+extension EditCustomListCoordinator: UIGestureRecognizerDelegate {
+    // For some reason, intercepting `popViewController(animated: Bool)` in `InterceptibleNavigationController`
+    // by SWIPING back leads to weird behaviour where subsequent navigation seem to happen systemwise but not
+    // UI-wise. This leads to the UI freezing up, and the only remedy is to restart the app.
+    //
+    // To get around this issue we can intercept the back swipe gesture and manually perform the transition
+    // instead, thereby bypassing the inner mechanisms that seem to go out of sync.
+    func gestureRecognizerShouldBegin(_ gestureRecognizer: UIGestureRecognizer) -> Bool {
+        guard gestureRecognizer == navigationController.interactivePopGestureRecognizer else {
+            return true
+        }
+        navigationController.popViewController(animated: true)
+        return false
+    }
+}
diff --git a/ios/MullvadVPN/Coordinators/CustomLists/EditLocationsCoordinator.swift b/ios/MullvadVPN/Coordinators/CustomLists/EditLocationsCoordinator.swift
index 9255a2bc29..9ca615ea69 100644
--- a/ios/MullvadVPN/Coordinators/CustomLists/EditLocationsCoordinator.swift
+++ b/ios/MullvadVPN/Coordinators/CustomLists/EditLocationsCoordinator.swift
@@ -6,6 +6,7 @@
 //  Copyright © 2024 Mullvad VPN AB. All rights reserved.
 //
 
+import Combine
 import MullvadSettings
 import MullvadTypes
 import Routing
@@ -14,9 +15,9 @@ import UIKit
 class EditLocationsCoordinator: Coordinator, Presentable, Presenting {
     private let navigationController: UINavigationController
     private let nodes: [LocationNode]
-    private var customList: CustomList
+    private var subject: CurrentValueSubject<CustomListViewModel, Never>
 
-    var didFinish: ((EditLocationsCoordinator, CustomList) -> Void)?
+    var didFinish: ((EditLocationsCoordinator) -> Void)?
 
     var presentedViewController: UIViewController {
         navigationController
@@ -25,17 +26,17 @@ class EditLocationsCoordinator: Coordinator, Presentable, Presenting {
     init(
         navigationController: UINavigationController,
         nodes: [LocationNode],
-        customList: CustomList
+        subject: CurrentValueSubject<CustomListViewModel, Never>
     ) {
         self.navigationController = navigationController
         self.nodes = nodes
-        self.customList = customList
+        self.subject = subject
     }
 
     func start() {
         let controller = AddLocationsViewController(
             allLocationsNodes: nodes,
-            customList: customList
+            subject: subject
         )
         controller.delegate = self
 
@@ -50,11 +51,7 @@ class EditLocationsCoordinator: Coordinator, Presentable, Presenting {
 }
 
 extension EditLocationsCoordinator: AddLocationsViewControllerDelegate {
-    func didUpdateSelectedLocations(locations: [RelayLocation]) {
-        customList.locations = locations
-    }
-
     func didBack() {
-        didFinish?(self, customList)
+        didFinish?(self)
     }
 }
diff --git a/ios/MullvadVPN/Coordinators/CustomLists/ListCustomListCoordinator.swift b/ios/MullvadVPN/Coordinators/CustomLists/ListCustomListCoordinator.swift
index 2b238dd1e5..713458e5b5 100644
--- a/ios/MullvadVPN/Coordinators/CustomLists/ListCustomListCoordinator.swift
+++ b/ios/MullvadVPN/Coordinators/CustomLists/ListCustomListCoordinator.swift
@@ -65,6 +65,12 @@ class ListCustomListCoordinator: Coordinator, Presentable, Presenting {
             self.updateRelayConstraints(for: action, in: list)
         }
 
+        coordinator.didCancel = { [weak self] editCustomListCoordinator in
+            guard let self else { return }
+            popToList()
+            editCustomListCoordinator.removeFromParent()
+        }
+
         coordinator.start()
         addChild(coordinator)
     }
@@ -84,6 +90,14 @@ class ListCustomListCoordinator: Coordinator, Presentable, Presenting {
                     customListSelection: UserSelectedRelays.CustomListSelection(listId: list.id, isList: true)
                 )
                 relayConstraints.locations = .only(selectedRelays)
+            } else {
+                let selectedConstraintIsRemovedFromList = list.locations.filter {
+                    relayConstraints.locations.value?.locations.contains($0) ?? false
+                }.isEmpty
+
+                if selectedConstraintIsRemovedFromList {
+                    relayConstraints.locations = .only(UserSelectedRelays(locations: []))
+                }
             }
         case .delete:
             relayConstraints.locations = .only(UserSelectedRelays(locations: []))
diff --git a/ios/MullvadVPN/Coordinators/LocationCoordinator.swift b/ios/MullvadVPN/Coordinators/LocationCoordinator.swift
index e146a7cd06..30fe23d8e3 100644
--- a/ios/MullvadVPN/Coordinators/LocationCoordinator.swift
+++ b/ios/MullvadVPN/Coordinators/LocationCoordinator.swift
@@ -147,7 +147,7 @@ class LocationCoordinator: Coordinator, Presentable, Presenting {
 
     private func showEditCustomLists(nodes: [LocationNode]) {
         let coordinator = ListCustomListCoordinator(
-            navigationController: CustomNavigationController(),
+            navigationController: InterceptibleNavigationController(),
             interactor: CustomListInteractor(repository: customListRepository),
             tunnelManager: tunnelManager,
             nodes: nodes
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationCellViewModel.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationCellViewModel.swift
index df0a3ba62c..14b7745efd 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationCellViewModel.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationCellViewModel.swift
@@ -15,14 +15,19 @@ struct LocationCellViewModel: Hashable {
     var isSelected = false
 
     func hash(into hasher: inout Hasher) {
-        hasher.combine(section)
         hasher.combine(node)
+        hasher.combine(node.children.count)
+        hasher.combine(section)
+        hasher.combine(isSelected)
+        hasher.combine(indentationLevel)
     }
 
     static func == (lhs: Self, rhs: Self) -> Bool {
         lhs.node == rhs.node &&
+            lhs.node.children.count == rhs.node.children.count &&
             lhs.section == rhs.section &&
-            lhs.isSelected == rhs.isSelected
+            lhs.isSelected == rhs.isSelected &&
+            lhs.indentationLevel == rhs.indentationLevel
     }
 }
 
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift
index 42957d24d6..3272f0e65b 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift
@@ -71,7 +71,7 @@ final class LocationDataSource:
         filterRelays(by: currentSearchString)
     }
 
-    func filterRelays(by searchString: String, scrollToSelected: Bool = true) {
+    func filterRelays(by searchString: String) {
         currentSearchString = searchString
 
         let list = sections.enumerated().map { index, section in
@@ -84,11 +84,6 @@ final class LocationDataSource:
         }
 
         updateDataSnapshot(with: list, reloadExisting: !searchString.isEmpty) {
-            guard scrollToSelected else {
-                self.setSelectedItem(self.selectedItem, animated: false)
-                return
-            }
-
             if searchString.isEmpty {
                 self.setSelectedItem(self.selectedItem, animated: false, completion: {
                     self.scrollToSelectedRelay()
@@ -99,23 +94,54 @@ final class LocationDataSource:
         }
     }
 
+    /// Refreshes the custom list section and keeps all modifications intact (selection and expanded states).
     func refreshCustomLists(selectedRelays: UserSelectedRelays?) {
-        let allLocationsDataSource =
-            dataSources.first(where: { $0 is AllLocationDataSource }) as? AllLocationDataSource
-
-        let customListsDataSource =
+        guard let allLocationsDataSource =
+            dataSources.first(where: { $0 is AllLocationDataSource }) as? AllLocationDataSource,
+            let customListsDataSource =
             dataSources.first(where: { $0 is CustomListsDataSource }) as? CustomListsDataSource
+        else {
+            return
+        }
 
-        customListsDataSource?.reload(allLocationNodes: allLocationsDataSource?.nodes ?? [], isFiltered: hasFilter)
+        // Take a "snapshot" of the currently expanded nodes.
+        let expandedNodes = customListsDataSource.nodes
+            .flatMap { [$0] + $0.flattened }
+            .filter { $0.showsChildren }
+
+        // Reload data source with (possibly) updated custom lists.
+        customListsDataSource.reload(allLocationNodes: allLocationsDataSource.nodes, isFiltered: hasFilter)
 
+        // Reapply current selection.
         mapSelectedItem(from: selectedRelays)
-        filterRelays(by: currentSearchString, scrollToSelected: false)
+
+        // Reapply current search filter.
+        let searchResultNodes = dataSources[0].search(by: currentSearchString)
+
+        // Reapply expanded status and override nodes being hidden by search filter.
+        RootLocationNode(children: searchResultNodes).forEachDescendant { node in
+            node.showsChildren = expandedNodes.contains(node)
+            node.isHiddenFromSearch = false
+        }
+
+        // Construct node tree.
+        let list = searchResultNodes.flatMap { node in
+            let rootNode = RootLocationNode(children: [node])
+            return recursivelyCreateCellViewModelTree(for: rootNode, in: .customLists, indentationLevel: 0)
+        }
+
+        updateDataSnapshot(with: [
+            list,
+            snapshot().itemIdentifiers(inSection: .allLocations),
+        ], reloadExisting: true)
     }
 
+    // Called from `LocationDiffableDataSourceProtocol`.
     func nodeShowsChildren(_ node: LocationNode) -> Bool {
         node.showsChildren
     }
 
+    // Called from `LocationDiffableDataSourceProtocol`.
     func nodeShouldBeSelected(_ node: LocationNode) -> Bool {
         false
     }
@@ -136,11 +162,19 @@ final class LocationDataSource:
             if let customListSelection = selectedRelays.customListSelection,
                let customList = customListsDataSource?.customList(by: customListSelection.listId),
                let selectedNode = customListsDataSource?.node(by: selectedRelays, for: customList) {
-                selectedItem = LocationCellViewModel(section: .customLists, node: selectedNode)
+                selectedItem = LocationCellViewModel(
+                    section: .customLists,
+                    node: selectedNode,
+                    indentationLevel: selectedNode.hierarchyLevel
+                )
                 // Look for a matching all locations node.
             } else if let location = selectedRelays.locations.first,
                       let selectedNode = allLocationsDataSource?.node(by: location) {
-                selectedItem = LocationCellViewModel(section: .allLocations, node: selectedNode)
+                selectedItem = LocationCellViewModel(
+                    section: .allLocations,
+                    node: selectedNode,
+                    indentationLevel: selectedNode.hierarchyLevel
+                )
             }
         }
     }
@@ -151,7 +185,7 @@ final class LocationDataSource:
 
         let rootNode = selectedItem.node.root
 
-        // Exit early if no changes to the node tree are necessary.
+        // Exit early if no changes to the node tree should be made.
         guard selectedItem.node != rootNode else {
             completion?()
             return
@@ -175,7 +209,7 @@ final class LocationDataSource:
             indentationLevel: 1
         )
 
-        // Insert the new node tree below the select item.
+        // Insert the new node tree below the selected item.
         var snapshotItems = snapshot().itemIdentifiers(inSection: selectedItem.section)
         snapshotItems.insert(contentsOf: nodesToAdd, at: indexPath.row + 1)
 
@@ -243,7 +277,7 @@ extension LocationDataSource: UITableViewDelegate {
 
     func tableView(_ tableView: UITableView, willDisplay cell: UITableViewCell, forRowAt indexPath: IndexPath) {
         if let item = itemIdentifier(for: indexPath), item == selectedItem {
-            tableView.selectRow(at: indexPath, animated: false, scrollPosition: .none)
+            cell.setSelected(true, animated: false)
         }
     }
 
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationNode.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationNode.swift
index fbf2fbf8fb..4534ac4d8e 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationNode.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationNode.swift
@@ -45,6 +45,12 @@ extension LocationNode {
         parent?.root ?? self
     }
 
+    var hierarchyLevel: Int {
+        var level = 0
+        forEachAncestor { _ in level += 1 }
+        return level
+    }
+
     func countryFor(code: String) -> LocationNode? {
         self.code == code ? self : children.first(where: { $0.code == code })
     }
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationSectionHeaderView.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationSectionHeaderView.swift
index 4a137d9cc1..220df6323d 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationSectionHeaderView.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationSectionHeaderView.swift
@@ -54,10 +54,9 @@ class LocationSectionHeaderView: UIView, UIContentView {
         addConstrainedSubviews([nameLabel, actionButton]) {
             nameLabel.pinEdgesToSuperviewMargins(.all().excluding(.trailing))
 
-            actionButton.pinEdgesToSuperviewMargins(PinnableEdges([.trailing(.zero)]))
-            actionButton.widthAnchor.constraint(equalToConstant: 24)
-            actionButton.heightAnchor.constraint(equalTo: actionButton.widthAnchor, multiplier: 1)
-            actionButton.centerYAnchor.constraint(equalTo: self.centerYAnchor)
+            actionButton.pinEdgesToSuperview(PinnableEdges([.trailing(8)]))
+            actionButton.heightAnchor.constraint(equalTo: heightAnchor)
+            actionButton.widthAnchor.constraint(equalTo: actionButton.heightAnchor)
 
             actionButton.leadingAnchor.constraint(equalTo: nameLabel.trailingAnchor, constant: 16)
         }