Merge branch 'adding-section-header-view-for-select-location-sections-ios-549'

10 files changed, 310 insertions, 116 deletions

diff --git a/ios/MullvadVPN.xcodeproj/project.pbxproj b/ios/MullvadVPN.xcodeproj/project.pbxproj
index 815b5a268f..e11619a445 100644
--- a/ios/MullvadVPN.xcodeproj/project.pbxproj
+++ b/ios/MullvadVPN.xcodeproj/project.pbxproj
@@ -116,7 +116,6 @@
 		584023292A407F5F007B27AC /* libtunnel_obfuscator_proxy.a in Frameworks */ = {isa = PBXBuildFile; fileRef = 584023282A407F5F007B27AC /* libtunnel_obfuscator_proxy.a */; };
 		58421030282D8A3C00F24E46 /* UpdateAccountDataOperation.swift in Sources */ = {isa = PBXBuildFile; fileRef = 5842102F282D8A3C00F24E46 /* UpdateAccountDataOperation.swift */; };
 		58421032282E42B000F24E46 /* UpdateDeviceDataOperation.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58421031282E42B000F24E46 /* UpdateDeviceDataOperation.swift */; };
-		58435AC229CB2A350099C71B /* LocationCellFactory.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58435AC129CB2A350099C71B /* LocationCellFactory.swift */; };
 		584592612639B4A200EF967F /* TermsOfServiceContentView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 584592602639B4A200EF967F /* TermsOfServiceContentView.swift */; };
 		5846227126E229F20035F7C2 /* StoreSubscription.swift in Sources */ = {isa = PBXBuildFile; fileRef = 5846227026E229F20035F7C2 /* StoreSubscription.swift */; };
 		5846227326E22A160035F7C2 /* StorePaymentObserver.swift in Sources */ = {isa = PBXBuildFile; fileRef = 5846227226E22A160035F7C2 /* StorePaymentObserver.swift */; };
@@ -848,7 +847,9 @@
 		F09D04BD2AEBB7C5003D4F89 /* OutgoingConnectionService.swift in Sources */ = {isa = PBXBuildFile; fileRef = F09D04BC2AEBB7C5003D4F89 /* OutgoingConnectionService.swift */; };
 		F09D04C02AF39D63003D4F89 /* OutgoingConnectionServiceTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = F09D04BF2AF39D63003D4F89 /* OutgoingConnectionServiceTests.swift */; };
 		F09D04C12AF39EA2003D4F89 /* OutgoingConnectionService.swift in Sources */ = {isa = PBXBuildFile; fileRef = F09D04BC2AEBB7C5003D4F89 /* OutgoingConnectionService.swift */; };
+		F0A92B3C2B8E44F900DC7B37 /* InMemoryCustomListRepository.swift in Sources */ = {isa = PBXBuildFile; fileRef = F0A92B3B2B8E44F900DC7B37 /* InMemoryCustomListRepository.swift */; };
 		F0B0E6972AFE6E7E001DC66B /* XCTest+Async.swift in Sources */ = {isa = PBXBuildFile; fileRef = F0B0E6962AFE6E7E001DC66B /* XCTest+Async.swift */; };
+		F0BE65372B9F136A005CC385 /* LocationSectionHeaderView.swift in Sources */ = {isa = PBXBuildFile; fileRef = F0BE65362B9F136A005CC385 /* LocationSectionHeaderView.swift */; };
 		F0C2AEFD2A0BB5CC00986207 /* NotificationProviderIdentifier.swift in Sources */ = {isa = PBXBuildFile; fileRef = F0C2AEFC2A0BB5CC00986207 /* NotificationProviderIdentifier.swift */; };
 		F0C3333C2B31A29C00D1A478 /* MullvadSettings.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 58B2FDD32AA71D2A003EB5C6 /* MullvadSettings.framework */; };
 		F0C6A8432AB08E54000777A8 /* RedeemVoucherViewConfiguration.swift in Sources */ = {isa = PBXBuildFile; fileRef = F0C6A8422AB08E54000777A8 /* RedeemVoucherViewConfiguration.swift */; };
@@ -1419,7 +1420,6 @@
 		5842102D282D3FC200F24E46 /* ResultBlockOperation.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ResultBlockOperation.swift; sourceTree = "<group>"; };
 		5842102F282D8A3C00F24E46 /* UpdateAccountDataOperation.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = UpdateAccountDataOperation.swift; sourceTree = "<group>"; };
 		58421031282E42B000F24E46 /* UpdateDeviceDataOperation.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = UpdateDeviceDataOperation.swift; sourceTree = "<group>"; };
-		58435AC129CB2A350099C71B /* LocationCellFactory.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = LocationCellFactory.swift; sourceTree = "<group>"; };
 		584592602639B4A200EF967F /* TermsOfServiceContentView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TermsOfServiceContentView.swift; sourceTree = "<group>"; };
 		5846226426E0D9630035F7C2 /* ProductsRequestOperation.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ProductsRequestOperation.swift; sourceTree = "<group>"; };
 		5846227026E229F20035F7C2 /* StoreSubscription.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = StoreSubscription.swift; sourceTree = "<group>"; };
@@ -1981,7 +1981,9 @@
 		F09D04BA2AE95396003D4F89 /* URLSessionStub.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = URLSessionStub.swift; sourceTree = "<group>"; };
 		F09D04BC2AEBB7C5003D4F89 /* OutgoingConnectionService.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = OutgoingConnectionService.swift; sourceTree = "<group>"; };
 		F09D04BF2AF39D63003D4F89 /* OutgoingConnectionServiceTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = OutgoingConnectionServiceTests.swift; sourceTree = "<group>"; };
+		F0A92B3B2B8E44F900DC7B37 /* InMemoryCustomListRepository.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = InMemoryCustomListRepository.swift; sourceTree = "<group>"; };
 		F0B0E6962AFE6E7E001DC66B /* XCTest+Async.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = "XCTest+Async.swift"; sourceTree = "<group>"; };
+		F0BE65362B9F136A005CC385 /* LocationSectionHeaderView.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = LocationSectionHeaderView.swift; sourceTree = "<group>"; };
 		F0C2AEFC2A0BB5CC00986207 /* NotificationProviderIdentifier.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = NotificationProviderIdentifier.swift; sourceTree = "<group>"; };
 		F0C6A8422AB08E54000777A8 /* RedeemVoucherViewConfiguration.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = RedeemVoucherViewConfiguration.swift; sourceTree = "<group>"; };
 		F0C6FA842A6A733700F521F0 /* InAppPurchaseInteractor.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = InAppPurchaseInteractor.swift; sourceTree = "<group>"; };
@@ -2403,13 +2405,14 @@
 			children = (
 				F050AE5F2B73A41E003F4EDB /* AllLocationDataSource.swift */,
 				F050AE612B74DBAC003F4EDB /* CustomListsDataSource.swift */,
+				F0A92B3B2B8E44F900DC7B37 /* InMemoryCustomListRepository.swift */,
 				5888AD82227B11080051EB06 /* LocationCell.swift */,
-				58435AC129CB2A350099C71B /* LocationCellFactory.swift */,
 				F050AE4D2B70D7F8003F4EDB /* LocationCellViewModel.swift */,
 				583DA21325FA4B5C00318683 /* LocationDataSource.swift */,
 				F050AE5D2B739A73003F4EDB /* LocationDataSourceProtocol.swift */,
 				7A6389F72B864CDF008E77E1 /* LocationNode.swift */,
 				F050AE512B70DFC0003F4EDB /* LocationSection.swift */,
+				F0BE65362B9F136A005CC385 /* LocationSectionHeaderView.swift */,
 				5888AD86227B17950051EB06 /* LocationViewController.swift */,
 			);
 			path = SelectLocation;
@@ -5115,7 +5118,6 @@
 				58B26E22294351EA00D5980C /* InAppNotificationProvider.swift in Sources */,
 				5893716A28817A45004EE76C /* DeviceManagementViewController.swift in Sources */,
 				7A9CCCB82A96302800DD6A34 /* SetupAccountCompletedCoordinator.swift in Sources */,
-				58435AC229CB2A350099C71B /* LocationCellFactory.swift in Sources */,
 				58BFA5C622A7C97F00A6173D /* RelayCacheTracker.swift in Sources */,
 				7A0B311E2B303A0D004B12E0 /* AccessbilityIdentifier.swift in Sources */,
 				E158B360285381C60002F069 /* String+AccountFormatting.swift in Sources */,
@@ -5129,6 +5131,8 @@
 				7ABE318D2A1CDD4500DF4963 /* UIFont+Weight.swift in Sources */,
 				58C774BE29A7A249003A1A56 /* CustomNavigationController.swift in Sources */,
 				E1FD0DF528AA7CE400299DB4 /* StatusActivityView.swift in Sources */,
+				F0BE65372B9F136A005CC385 /* LocationSectionHeaderView.swift in Sources */,
+				F0A92B3C2B8E44F900DC7B37 /* InMemoryCustomListRepository.swift in Sources */,
 				7A2960FD2A964BB700389B82 /* AlertPresentation.swift in Sources */,
 				0697D6E728F01513007A9E99 /* TransportMonitor.swift in Sources */,
 				58968FAE28743E2000B799DC /* TunnelInteractor.swift in Sources */,
diff --git a/ios/MullvadVPN/Coordinators/ApplicationCoordinator.swift b/ios/MullvadVPN/Coordinators/ApplicationCoordinator.swift
index d0fc0c6bc3..b92a5492d8 100644
--- a/ios/MullvadVPN/Coordinators/ApplicationCoordinator.swift
+++ b/ios/MullvadVPN/Coordinators/ApplicationCoordinator.swift
@@ -43,6 +43,14 @@ final class ApplicationCoordinator: Coordinator, Presenting, RootContainerViewCo
      */
     private let secondaryNavigationContainer = RootContainerViewController()
 
+    private var customListRepository: CustomListRepositoryProtocol {
+        #if DEBUG
+        InMemoryCustomListRepository()
+        #else
+        CustomListRepository()
+        #endif
+    }
+
     /// Posts `preferredAccountNumber` notification when user inputs the account number instead of voucher code
     private let preferredAccountNumberSubject = PassthroughSubject<String, Never>()
 
@@ -710,7 +718,8 @@ final class ApplicationCoordinator: Coordinator, Presenting, RootContainerViewCo
         let locationCoordinator = LocationCoordinator(
             navigationController: navigationController,
             tunnelManager: tunnelManager,
-            relayCacheTracker: relayCacheTracker
+            relayCacheTracker: relayCacheTracker,
+            customListRepository: customListRepository
         )
 
         locationCoordinator.didFinish = { [weak self] _ in
diff --git a/ios/MullvadVPN/Coordinators/LocationCoordinator.swift b/ios/MullvadVPN/Coordinators/LocationCoordinator.swift
index 39bffdaab9..4d18672f5a 100644
--- a/ios/MullvadVPN/Coordinators/LocationCoordinator.swift
+++ b/ios/MullvadVPN/Coordinators/LocationCoordinator.swift
@@ -7,6 +7,7 @@
 //
 
 import MullvadREST
+import MullvadSettings
 import MullvadTypes
 import Routing
 import UIKit
@@ -15,6 +16,7 @@ class LocationCoordinator: Coordinator, Presentable, Presenting, RelayCacheTrack
     private let tunnelManager: TunnelManager
     private let relayCacheTracker: RelayCacheTracker
     private var cachedRelays: CachedRelays?
+    private var customListRepository: CustomListRepositoryProtocol
 
     let navigationController: UINavigationController
 
@@ -42,17 +44,21 @@ class LocationCoordinator: Coordinator, Presentable, Presenting, RelayCacheTrack
     init(
         navigationController: UINavigationController,
         tunnelManager: TunnelManager,
-        relayCacheTracker: RelayCacheTracker
+        relayCacheTracker: RelayCacheTracker,
+        customListRepository: CustomListRepositoryProtocol
     ) {
         self.navigationController = navigationController
         self.tunnelManager = tunnelManager
         self.relayCacheTracker = relayCacheTracker
+        self.customListRepository = customListRepository
     }
 
     func start() {
-        let selectLocationViewController = LocationViewController()
+        let locationViewController = LocationViewController(customListRepository: customListRepository)
+        locationViewController.delegate = self
+
+        locationViewController.didSelectRelays = { [weak self] locations in
 
-        selectLocationViewController.didSelectRelays = { [weak self] locations in
             guard let self else { return }
 
             var relayConstraints = tunnelManager.settings.relayConstraints
@@ -65,7 +71,7 @@ class LocationCoordinator: Coordinator, Presentable, Presenting, RelayCacheTrack
             didFinish?(self)
         }
 
-        selectLocationViewController.navigateToFilter = { [weak self] in
+        locationViewController.navigateToFilter = { [weak self] in
             guard let self else { return }
 
             let coordinator = makeRelayFilterCoordinator(forModalPresentation: true)
@@ -74,7 +80,7 @@ class LocationCoordinator: Coordinator, Presentable, Presenting, RelayCacheTrack
             presentChild(coordinator, animated: true)
         }
 
-        selectLocationViewController.didUpdateFilter = { [weak self] filter in
+        locationViewController.didUpdateFilter = { [weak self] filter in
             guard let self else { return }
 
             var relayConstraints = tunnelManager.settings.relayConstraints
@@ -83,7 +89,7 @@ class LocationCoordinator: Coordinator, Presentable, Presenting, RelayCacheTrack
             tunnelManager.updateSettings([.relayConstraints(relayConstraints)])
         }
 
-        selectLocationViewController.didFinish = { [weak self] in
+        locationViewController.didFinish = { [weak self] in
             guard let self else { return }
 
             didFinish?(self)
@@ -93,12 +99,12 @@ class LocationCoordinator: Coordinator, Presentable, Presenting, RelayCacheTrack
 
         if let cachedRelays = try? relayCacheTracker.getCachedRelays() {
             self.cachedRelays = cachedRelays
-            selectLocationViewController.setCachedRelays(cachedRelays, filter: relayFilter)
+            locationViewController.setCachedRelays(cachedRelays, filter: relayFilter)
         }
 
-        selectLocationViewController.relayLocations = tunnelManager.settings.relayConstraints.locations.value
+        locationViewController.relayLocations = tunnelManager.settings.relayConstraints.locations.value
 
-        navigationController.pushViewController(selectLocationViewController, animated: false)
+        navigationController.pushViewController(locationViewController, animated: false)
     }
 
     private func makeRelayFilterCoordinator(forModalPresentation isModalPresentation: Bool)
@@ -131,3 +137,9 @@ class LocationCoordinator: Coordinator, Presentable, Presenting, RelayCacheTrack
         selectLocationViewController?.setCachedRelays(cachedRelays, filter: relayFilter)
     }
 }
+
+extension LocationCoordinator: LocationViewControllerDelegate {
+    func didRequestRouteToCustomLists(_ controller: LocationViewController) {
+        // TODO: Show add/Edit bottom sheet.
+    }
+}
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/InMemoryCustomListRepository.swift b/ios/MullvadVPN/View controllers/SelectLocation/InMemoryCustomListRepository.swift
new file mode 100644
index 0000000000..999b4ad110
--- /dev/null
+++ b/ios/MullvadVPN/View controllers/SelectLocation/InMemoryCustomListRepository.swift
@@ -0,0 +1,63 @@
+//
+//  InMemoryCustomListRepository.swift
+//  MullvadVPN
+//
+//  Created by Mojgan on 2024-01-31.
+//  Copyright © 2024 Mullvad VPN AB. All rights reserved.
+//
+
+import Combine
+import Foundation
+import MullvadSettings
+import MullvadTypes
+
+class InMemoryCustomListRepository: CustomListRepositoryProtocol {
+    var publisher: AnyPublisher<[CustomList], Never> {
+        passthroughSubject.eraseToAnyPublisher()
+    }
+
+    private var customRelayLists: [CustomList] = [
+        CustomList(
+            id: UUID(uuidString: "F17948CB-18E2-4F84-82CD-5780F94216DB")!,
+            name: "Netflix",
+            locations: [.city("al", "tia")]
+        ),
+        CustomList(
+            id: UUID(uuidString: "4104C603-B35D-4A64-8865-96C0BF33D57F")!,
+            name: "Streaming",
+            locations: [
+                .city("us", "dal"),
+                .country("se"),
+                .city("de", "ber"),
+            ]
+        ),
+    ]
+
+    private let passthroughSubject = PassthroughSubject<[CustomList], Never>()
+
+    func update(_ list: CustomList) {
+        if let index = customRelayLists.firstIndex(where: { $0.id == list.id }) {
+            customRelayLists[index] = list
+        }
+    }
+
+    func delete(id: UUID) {
+        if let index = customRelayLists.firstIndex(where: { $0.id == id }) {
+            customRelayLists.remove(at: index)
+        }
+    }
+
+    func fetch(by id: UUID) -> CustomList? {
+        return customRelayLists.first(where: { $0.id == id })
+    }
+
+    func create(_ name: String, locations: [RelayLocation]) throws -> CustomList {
+        let item = CustomList(id: UUID(), name: name, locations: locations)
+        customRelayLists.append(item)
+        return item
+    }
+
+    func fetchAll() -> [CustomList] {
+        customRelayLists
+    }
+}
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationCell.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationCell.swift
index 8c1ed9334b..dfdd791de1 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationCell.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationCell.swift
@@ -8,25 +8,46 @@
 
 import UIKit
 
-private let kCollapseButtonWidth: CGFloat = 24
-private let kRelayIndicatorSize: CGFloat = 16
+protocol LocationCellDelegate: AnyObject {
+    func toggle(cell: LocationCell)
+}
 
 class LocationCell: UITableViewCell {
-    typealias CollapseHandler = (LocationCell) -> Void
+    weak var delegate: LocationCellDelegate?
+
+    private let locationLabel: UILabel = {
+        let label = UILabel()
+        label.font = UIFont.systemFont(ofSize: 16)
+        label.textColor = .white
+        label.lineBreakMode = .byWordWrapping
+        label.numberOfLines = 0
+        label.lineBreakStrategy = []
+        return label
+    }()
 
-    let locationLabel = UILabel()
-    let statusIndicator: UIView = {
+    private let statusIndicator: UIView = {
         let view = UIView()
-        view.layer.cornerRadius = kRelayIndicatorSize * 0.5
+        view.layer.cornerRadius = 8
         view.layer.cornerCurve = .circular
         return view
     }()
 
-    let tickImageView = UIImageView(image: UIImage(named: "IconTick"))
-    let collapseButton = UIButton(type: .custom)
+    private let tickImageView: UIImageView = {
+        let imageView = UIImageView(image: UIImage(resource: .iconTick))
+        imageView.tintColor = .white
+        return imageView
+    }()
+
+    private let collapseButton: UIButton = {
+        let button = UIButton(type: .custom)
+        button.accessibilityIdentifier = .collapseButton
+        button.isAccessibilityElement = false
+        button.tintColor = .white
+        return button
+    }()
 
-    private let chevronDown = UIImage(named: "IconChevronDown")
-    private let chevronUp = UIImage(named: "IconChevronUp")
+    private let chevronDown = UIImage(resource: .iconChevronDown)
+    private let chevronUp = UIImage(resource: .iconChevronUp)
 
     var isDisabled = false {
         didSet {
@@ -50,8 +71,6 @@ class LocationCell: UITableViewCell {
         }
     }
 
-    var didCollapseHandler: CollapseHandler?
-
     override var indentationLevel: Int {
         didSet {
             updateBackgroundColor()
@@ -103,17 +122,6 @@ class LocationCell: UITableViewCell {
         selectedBackgroundView = UIView()
         selectedBackgroundView?.backgroundColor = UIColor.Cell.Background.selected
 
-        locationLabel.font = UIFont.systemFont(ofSize: 17)
-        locationLabel.textColor = .white
-        locationLabel.lineBreakMode = .byWordWrapping
-        locationLabel.numberOfLines = 0
-        locationLabel.lineBreakStrategy = []
-
-        tickImageView.tintColor = .white
-
-        collapseButton.accessibilityIdentifier = .collapseButton
-        collapseButton.isAccessibilityElement = false
-        collapseButton.tintColor = .white
         collapseButton.addTarget(self, action: #selector(handleCollapseButton(_:)), for: .touchUpInside)
 
         [locationLabel, tickImageView, statusIndicator, collapseButton].forEach { subview in
@@ -131,7 +139,7 @@ class LocationCell: UITableViewCell {
             tickImageView.leadingAnchor.constraint(equalTo: contentView.layoutMarginsGuide.leadingAnchor),
             tickImageView.centerYAnchor.constraint(equalTo: contentView.centerYAnchor),
 
-            statusIndicator.widthAnchor.constraint(equalToConstant: kRelayIndicatorSize),
+            statusIndicator.widthAnchor.constraint(equalToConstant: 16),
             statusIndicator.heightAnchor.constraint(equalTo: statusIndicator.widthAnchor),
             statusIndicator.centerXAnchor.constraint(equalTo: tickImageView.centerXAnchor),
             statusIndicator.centerYAnchor.constraint(equalTo: tickImageView.centerYAnchor),
@@ -148,7 +156,7 @@ class LocationCell: UITableViewCell {
             collapseButton.widthAnchor
                 .constraint(
                     equalToConstant: UIMetrics.contentLayoutMargins.leading + UIMetrics
-                        .contentLayoutMargins.trailing + kCollapseButtonWidth
+                        .contentLayoutMargins.trailing + 24
                 ),
             collapseButton.topAnchor.constraint(equalTo: contentView.topAnchor),
             collapseButton.trailingAnchor.constraint(equalTo: contentView.trailingAnchor),
@@ -213,11 +221,11 @@ class LocationCell: UITableViewCell {
     }
 
     @objc private func handleCollapseButton(_ sender: UIControl) {
-        didCollapseHandler?(self)
+        delegate?.toggle(cell: self)
     }
 
     @objc private func toggleCollapseAccessibilityAction() -> Bool {
-        didCollapseHandler?(self)
+        delegate?.toggle(cell: self)
         return true
     }
 
@@ -255,3 +263,12 @@ class LocationCell: UITableViewCell {
         }
     }
 }
+
+extension LocationCell {
+    func configureCell(item: LocationCellViewModel) {
+        accessibilityIdentifier = item.node.code
+        locationLabel.text = item.node.name
+        showsCollapseControl = !item.node.children.isEmpty
+        isExpanded = item.node.showsChildren
+    }
+}
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationCellFactory.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationCellFactory.swift
deleted file mode 100644
index 1d0c1f9742..0000000000
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationCellFactory.swift
+++ /dev/null
@@ -1,51 +0,0 @@
-//
-//  LocationCellFactory.swift
-//  MullvadVPN
-//
-//  Created by Jon Petersson on 2023-03-17.
-//  Copyright © 2023 Mullvad VPN AB. All rights reserved.
-//
-
-import MullvadTypes
-import UIKit
-
-protocol LocationCellEventHandler {
-    func toggleCell(for item: LocationCellViewModel)
-}
-
-final class LocationCellFactory: CellFactoryProtocol {
-    var delegate: LocationCellEventHandler?
-    let tableView: UITableView
-    let reuseIdentifier: String
-
-    init(
-        tableView: UITableView,
-        reuseIdentifier: String
-    ) {
-        self.tableView = tableView
-        self.reuseIdentifier = reuseIdentifier
-    }
-
-    func makeCell(for item: LocationCellViewModel, indexPath: IndexPath) -> UITableViewCell {
-        let cell = tableView.dequeueReusableCell(
-            withIdentifier: reuseIdentifier,
-            for: indexPath
-        )
-
-        configureCell(cell, item: item, indexPath: indexPath)
-
-        return cell
-    }
-
-    func configureCell(_ cell: UITableViewCell, item: LocationCellViewModel, indexPath: IndexPath) {
-        guard let cell = cell as? LocationCell else { return }
-
-        cell.accessibilityIdentifier = item.node.code
-        cell.locationLabel.text = item.node.name
-        cell.showsCollapseControl = !item.node.children.isEmpty
-        cell.isExpanded = item.node.showsChildren
-        cell.didCollapseHandler = { [weak self] _ in
-            self?.delegate?.toggleCell(for: item)
-        }
-    }
-}
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift
index 34b78737a7..77da5c69b9 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift
@@ -15,11 +15,11 @@ import UIKit
 final class LocationDataSource: UITableViewDiffableDataSource<LocationSection, LocationCellViewModel> {
     private var currentSearchString = ""
     private let tableView: UITableView
-    private let locationCellFactory: LocationCellFactory
     private var dataSources: [LocationDataSourceProtocol] = []
     private var selectedItem: LocationCellViewModel?
 
     var didSelectRelayLocations: ((RelayLocations) -> Void)?
+    var didTapEditCustomLists: (() -> Void)?
 
     init(
         tableView: UITableView,
@@ -33,19 +33,18 @@ final class LocationDataSource: UITableViewDiffableDataSource<LocationSection, L
         #endif
         self.dataSources.append(allLocations)
 
-        let locationCellFactory = LocationCellFactory(
-            tableView: tableView,
-            reuseIdentifier: LocationSection.Cell.locationCell.reuseIdentifier
-        )
-        self.locationCellFactory = locationCellFactory
-
         super.init(tableView: tableView) { _, indexPath, itemIdentifier in
-            locationCellFactory.makeCell(for: itemIdentifier, indexPath: indexPath)
+            let reuseIdentifier = LocationSection.Cell.locationCell.reuseIdentifier
+            let cell = tableView.dequeueReusableCell(
+                withIdentifier: reuseIdentifier,
+                for: indexPath
+                // swiftlint:disable:next force_cast
+            ) as! LocationCell
+            cell.configureCell(item: itemIdentifier)
+            return cell
         }
 
         tableView.delegate = self
-        locationCellFactory.delegate = self
-
         defaultRowAnimation = .fade
         registerClasses()
     }
@@ -218,21 +217,44 @@ final class LocationDataSource: UITableViewDiffableDataSource<LocationSection, L
 
         return viewModels
     }
+
+    override func tableView(_ tableView: UITableView, cellForRowAt indexPath: IndexPath) -> UITableViewCell {
+        // swiftlint:disable:next force_cast
+        let cell = super.tableView(tableView, cellForRowAt: indexPath) as! LocationCell
+        cell.delegate = self
+        return cell
+    }
 }
 
 extension LocationDataSource: UITableViewDelegate {
+    func tableView(_ tableView: UITableView, viewForHeaderInSection section: Int) -> UIView? {
+        switch LocationSection.allCases[section] {
+        case .allLocations:
+            return LocationSectionHeaderView(
+                configuration: LocationSectionHeaderView.Configuration(name: LocationSection.allLocations.description)
+            )
+        case .customLists:
+            return LocationSectionHeaderView(configuration: LocationSectionHeaderView.Configuration(
+                name: LocationSection.customLists.description,
+                primaryAction: UIAction(
+                    handler: { [weak self] _ in
+                        self?.didTapEditCustomLists?()
+                    }
+                )
+            ))
+        }
+    }
+
     func tableView(_ tableView: UITableView, viewForFooterInSection section: Int) -> UIView? {
         nil
     }
 
     func tableView(_ tableView: UITableView, heightForFooterInSection section: Int) -> CGFloat {
-        let section = snapshot().sectionIdentifiers[section]
-
-        switch section {
+        switch LocationSection.allCases[section] {
+        case .allLocations:
+            return .zero
         case .customLists:
             return 24
-        case .allLocations:
-            return 0
         }
     }
 
@@ -263,9 +285,10 @@ extension LocationDataSource: UITableViewDelegate {
     }
 }
 
-extension LocationDataSource: LocationCellEventHandler {
-    func toggleCell(for item: LocationCellViewModel) {
-        guard let indexPath = indexPath(for: item) else { return }
+extension LocationDataSource: LocationCellDelegate {
+    func toggle(cell: LocationCell) {
+        guard let indexPath = tableView.indexPath(for: cell),
+              let item = itemIdentifier(for: indexPath) else { return }
 
         let sections = LocationSection.allCases
         let section = sections[indexPath.section]
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSourceProtocol.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSourceProtocol.swift
index 9fff54fff1..79f78ebc99 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSourceProtocol.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSourceProtocol.swift
@@ -9,7 +9,6 @@
 import Foundation
 import MullvadREST
 import MullvadTypes
-import UIKit
 
 protocol LocationDataSourceProtocol {
     var nodes: [LocationNode] { get }
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationSectionHeaderView.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationSectionHeaderView.swift
new file mode 100644
index 0000000000..49c9cbce20
--- /dev/null
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationSectionHeaderView.swift
@@ -0,0 +1,95 @@
+//
+//  LocationSectionHeaderView.swift
+//  MullvadVPN
+//
+//  Created by Mojgan on 2024-01-25.
+//  Copyright © 2024 Mullvad VPN AB. All rights reserved.
+//
+
+import Foundation
+import UIKit
+
+class LocationSectionHeaderView: UIView, UIContentView {
+    var configuration: UIContentConfiguration {
+        get {
+            actualConfiguration
+        } set {
+            guard let newConfiguration = newValue as? Configuration,
+                  actualConfiguration != newConfiguration else { return }
+            let previousConfiguration = actualConfiguration
+            actualConfiguration = newConfiguration
+            apply(configuration: previousConfiguration)
+        }
+    }
+
+    private var actualConfiguration: Configuration
+    private let nameLabel: UILabel = {
+        let label = UILabel()
+        label.numberOfLines = 1
+        label.textColor = .primaryTextColor
+        label.font = .systemFont(ofSize: 16, weight: .semibold)
+        return label
+    }()
+
+    private let actionButton: UIButton = {
+        let button = UIButton(type: .system)
+        button.setImage(UIImage(systemName: "ellipsis"), for: .normal)
+        button.tintColor = UIColor(white: 1, alpha: 0.6)
+        return button
+    }()
+
+    init(configuration: Configuration) {
+        self.actualConfiguration = configuration
+        super.init(frame: .zero)
+        applyAppearance()
+        addSubviews()
+        apply(configuration: configuration)
+    }
+
+    required init?(coder: NSCoder) {
+        fatalError("init(coder:) has not been implemented")
+    }
+
+    private func addSubviews() {
+        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.leadingAnchor.constraint(equalTo: nameLabel.trailingAnchor, constant: 16)
+        }
+    }
+
+    private func apply(configuration: Configuration) {
+        let isActionHidden = configuration.primaryAction == nil
+        nameLabel.text = configuration.name
+        actionButton.isHidden = isActionHidden
+        actualConfiguration.primaryAction.flatMap { [weak self] action in
+            self?.actionButton.addAction(action, for: .touchUpInside)
+        }
+    }
+
+    private func applyAppearance() {
+        backgroundColor = .primaryColor
+        directionalLayoutMargins = NSDirectionalEdgeInsets(top: 8, leading: 16, bottom: 8, trailing: 24)
+    }
+}
+
+extension LocationSectionHeaderView {
+    struct Configuration: UIContentConfiguration, Equatable {
+        let name: String
+
+        var primaryAction: UIAction?
+
+        func makeContentView() -> UIView & UIContentView {
+            LocationSectionHeaderView(configuration: self)
+        }
+
+        func updated(for state: UIConfigurationState) -> Configuration {
+            self
+        }
+    }
+}
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationViewController.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationViewController.swift
index 16f4797d80..baa3cce181 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationViewController.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationViewController.swift
@@ -12,15 +12,21 @@ import MullvadSettings
 import MullvadTypes
 import UIKit
 
+protocol LocationViewControllerDelegate: AnyObject {
+    func didRequestRouteToCustomLists(_ controller: LocationViewController)
+}
+
 final class LocationViewController: UIViewController {
     private let searchBar = UISearchBar()
-    private let tableView = UITableView()
+    private let tableView = UITableView(frame: .zero, style: .grouped)
     private let topContentView = UIStackView()
     private let filterView = RelayFilterView()
     private var dataSource: LocationDataSource?
     private var cachedRelays: CachedRelays?
     private var filter = RelayFilter()
     var relayLocations: RelayLocations?
+    weak var delegate: LocationViewControllerDelegate?
+    var customListRepository: CustomListRepositoryProtocol
 
     override var preferredStatusBarStyle: UIStatusBarStyle {
         .lightContent
@@ -35,6 +41,15 @@ final class LocationViewController: UIViewController {
     var didUpdateFilter: ((RelayFilter) -> Void)?
     var didFinish: (() -> Void)?
 
+    init(customListRepository: CustomListRepositoryProtocol) {
+        self.customListRepository = customListRepository
+        super.init(nibName: nil, bundle: nil)
+    }
+
+    required init?(coder: NSCoder) {
+        fatalError("init(coder:) has not been implemented")
+    }
+
     // MARK: - View lifecycle
 
     override func viewDidLoad() {
@@ -68,7 +83,7 @@ final class LocationViewController: UIViewController {
             })
         )
 
-        setUpDataSource()
+        setUpDataSources()
         setUpTableView()
         setUpTopContent()
 
@@ -104,17 +119,24 @@ final class LocationViewController: UIViewController {
 
     // MARK: - Private
 
-    private func setUpDataSource() {
+    private func setUpDataSources() {
+        let allLocationDataSource = AllLocationDataSource()
+        let customListsDataSource = CustomListsDataSource(repository: customListRepository)
         dataSource = LocationDataSource(
             tableView: tableView,
-            allLocations: AllLocationDataSource(),
-            customLists: CustomListsDataSource(repository: CustomListRepository())
+            allLocations: allLocationDataSource,
+            customLists: customListsDataSource
         )
 
         dataSource?.didSelectRelayLocations = { [weak self] locations in
             self?.didSelectRelays?(locations)
         }
 
+        dataSource?.didTapEditCustomLists = { [weak self] in
+            guard let self else { return }
+            delegate?.didRequestRouteToCustomLists(self)
+        }
+
         if let cachedRelays {
             dataSource?.setRelays(cachedRelays.relays, selectedLocations: relayLocations, filter: filter)
         }
@@ -128,6 +150,7 @@ final class LocationViewController: UIViewController {
         tableView.indicatorStyle = .white
         tableView.keyboardDismissMode = .onDrag
         tableView.accessibilityIdentifier = .selectLocationTableView
+        tableView.sectionHeaderHeight = 56.0
     }
 
     private func setUpTopContent() {