adding section header view in select location

- each section should have its own distinct header view.
- custom list section has an action that should take user to add/edit custom list.it will be coming during upcoming changes.
- refactoring location cell
- removing LocationCellFactory

11 files changed, 210 insertions, 223 deletions

diff --git a/ios/MullvadVPN.xcodeproj/project.pbxproj b/ios/MullvadVPN.xcodeproj/project.pbxproj
index f709874aca..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 */; };
@@ -850,6 +849,7 @@
 		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 */; };
@@ -1420,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>"; };
@@ -1984,6 +1983,7 @@
 		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>"; };
@@ -2407,12 +2407,12 @@
 				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;
@@ -5118,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 */,
@@ -5132,6 +5131,7 @@
 				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 */,
diff --git a/ios/MullvadVPN/Coordinators/CustomLists/CustomListsDataSource.swift b/ios/MullvadVPN/Coordinators/CustomLists/CustomListsDataSource.swift
deleted file mode 100644
index e24a4c46dd..0000000000
--- a/ios/MullvadVPN/Coordinators/CustomLists/CustomListsDataSource.swift
+++ /dev/null
@@ -1,103 +0,0 @@
-//
-//  CustomListsDataSource.swift
-//  MullvadVPN
-//
-//  Created by Jon Petersson on 2024-02-22.
-//  Copyright © 2024 Mullvad VPN AB. All rights reserved.
-//
-
-import Foundation
-import MullvadREST
-import MullvadSettings
-import MullvadTypes
-import UIKit
-
-class CustomListsDataSource: LocationDataSourceProtocol {
-    private(set) var nodes = [LocationNode]()
-    var didTapEditCustomLists: (() -> Void)?
-
-    var viewForHeader: UIView? {
-        LocationSectionHeaderView(configuration: LocationSectionHeaderView.Configuration(
-            name: LocationSection.customLists.description,
-            primaryAction: UIAction(
-                handler: { [weak self] _ in
-                    self?.didTapEditCustomLists?()
-                }
-            )
-        ))
-    }
-
-    init(didTapEditCustomLists: (() -> Void)?) {
-        self.didTapEditCustomLists = didTapEditCustomLists
-    }
-
-    var searchableNodes: [LocationNode] {
-        nodes.flatMap { $0.children }
-    }
-
-    func reload(allLocationNodes: [LocationNode], customLists: [CustomList]) {
-        nodes = customLists.map { list in
-            let listNode = LocationListNode(
-                nodeName: list.name,
-                nodeCode: list.name.lowercased(),
-                locations: list.locations,
-                customList: list
-            )
-
-            listNode.children = list.locations.compactMap { location in
-                copy(location, from: allLocationNodes, withParent: listNode)
-            }
-
-            listNode.forEachDescendant { _, node in
-                node.nodeCode = "\(listNode.nodeCode)-\(node.nodeCode)"
-            }
-
-            return listNode
-        }
-    }
-
-    func node(by locations: [RelayLocation], for customList: CustomList) -> LocationNode? {
-        guard let customListNode = nodes.first(where: { $0.nodeName == customList.name })
-        else { return nil }
-
-        if locations.count > 1 {
-            return customListNode
-        } else {
-            return switch locations.first {
-            case let .country(countryCode):
-                customListNode.nodeFor(nodeCode: "\(customListNode.nodeCode)-\(countryCode)")
-            case let .city(_, cityCode):
-                customListNode.nodeFor(nodeCode: "\(customListNode.nodeCode)-\(cityCode)")
-            case let .hostname(_, _, hostCode):
-                customListNode.nodeFor(nodeCode: "\(customListNode.nodeCode)-\(hostCode)")
-            case .none:
-                nil
-            }
-        }
-    }
-
-    private func copy(
-        _ location: RelayLocation,
-        from allLocationNodes: [LocationNode],
-        withParent rootNode: LocationNode
-    ) -> LocationNode? {
-        let rootNode = RootNode(children: allLocationNodes)
-
-        return switch location {
-        case let .country(countryCode):
-            rootNode
-                .countryFor(countryCode: countryCode)?.copy(withParent: rootNode)
-
-        case let .city(countryCode, cityCode):
-            rootNode
-                .countryFor(countryCode: countryCode)?.copy(withParent: rootNode)
-                .cityFor(cityCode: cityCode)
-
-        case let .hostname(countryCode, cityCode, hostCode):
-            rootNode
-                .countryFor(countryCode: countryCode)?.copy(withParent: rootNode)
-                .cityFor(cityCode: cityCode)?
-                .hostFor(hostCode: hostCode)
-        }
-    }
-}
diff --git a/ios/MullvadVPN/Coordinators/LocationCoordinator.swift b/ios/MullvadVPN/Coordinators/LocationCoordinator.swift
index fefc4a6f3a..4d18672f5a 100644
--- a/ios/MullvadVPN/Coordinators/LocationCoordinator.swift
+++ b/ios/MullvadVPN/Coordinators/LocationCoordinator.swift
@@ -140,13 +140,6 @@ class LocationCoordinator: Coordinator, Presentable, Presenting, RelayCacheTrack
 
 extension LocationCoordinator: LocationViewControllerDelegate {
     func didRequestRouteToCustomLists(_ controller: LocationViewController) {
-        let coordinator = AddCustomListCoordinator(
-            navigationController: CustomNavigationController(),
-            customListInteractor: CustomListInteractor(
-                repository: customListRepository
-            )
-        )
-        coordinator.start()
-        presentChild(coordinator, animated: true)
+        // TODO: Show add/Edit bottom sheet.
     }
 }
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/AllLocationDataSource.swift b/ios/MullvadVPN/View controllers/SelectLocation/AllLocationDataSource.swift
index 7427b78cc9..a6e9e1bab0 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/AllLocationDataSource.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/AllLocationDataSource.swift
@@ -9,7 +9,6 @@
 import Foundation
 import MullvadREST
 import MullvadTypes
-import UIKit
 
 class AllLocationDataSource: LocationDataSourceProtocol {
     private(set) var nodes = [LocationNode]()
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/InMemoryCustomListRepository.swift b/ios/MullvadVPN/View controllers/SelectLocation/InMemoryCustomListRepository.swift
index fbc33ad071..999b4ad110 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/InMemoryCustomListRepository.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/InMemoryCustomListRepository.swift
@@ -17,12 +17,20 @@ class InMemoryCustomListRepository: CustomListRepositoryProtocol {
     }
 
     private var customRelayLists: [CustomList] = [
-        CustomList(id: UUID(), name: "Netflix", locations: [.city("al", "tia")]),
-        CustomList(id: UUID(), name: "Streaming", locations: [
-            .city("us", "dal"),
-            .country("se"),
-            .city("de", "ber"),
-        ]),
+        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>()
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 3c37884544..baa3cce181 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationViewController.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationViewController.swift
@@ -83,7 +83,7 @@ final class LocationViewController: UIViewController {
             })
         )
 
-        setUpDataSource()
+        setUpDataSources()
         setUpTableView()
         setUpTopContent()
 
@@ -119,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)
         }