From 887eaf3ada219d859649b115fe079298c651ae71 Mon Sep 17 00:00:00 2001
From: nielsandriesse <andriesseniels@gmail.com>
Date: Tue, 2 Jun 2020 16:10:23 +1000
Subject: [PATCH] Don't block while countries are loading

---
 Signal/src/Loki/Utilities/IP2Country.swift        | 15 +++++++++------
 Signal/src/Loki/View Controllers/PathVC.swift     |  4 +++-
 .../src/Loki/Utilities/Notification+Loki.swift    |  2 ++
 3 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/Signal/src/Loki/Utilities/IP2Country.swift b/Signal/src/Loki/Utilities/IP2Country.swift
index 155245398..f715cee02 100644
--- a/Signal/src/Loki/Utilities/IP2Country.swift
+++ b/Signal/src/Loki/Utilities/IP2Country.swift
@@ -4,14 +4,14 @@ final class IP2Country {
 
     private let ipv4Table = try! CSV(name: "GeoLite2-Country-Blocks-IPv4", extension: "csv", bundle: .main, delimiter: ",", encoding: .utf8, loadColumns: true)!
     private let countryNamesTable = try! CSV(name: "GeoLite2-Country-Locations-English", extension: "csv", bundle: .main, delimiter: ",", encoding: .utf8, loadColumns: true)!
-    private var countryNamesCache: [String:String] = [:]
+    var countryNamesCache: [String:String] = [:]
 
     // MARK: Lifecycle
     static let shared = IP2Country()
 
     private init() {
-        preloadCountriesIfNeeded()
-        NotificationCenter.default.addObserver(self, selector: #selector(preloadCountriesIfNeeded), name: .pathsBuilt, object: nil)
+        populateCacheIfNeeded()
+        NotificationCenter.default.addObserver(self, selector: #selector(populateCacheIfNeeded), name: .pathsBuilt, object: nil)
     }
 
     deinit {
@@ -19,7 +19,7 @@ final class IP2Country {
     }
 
     // MARK: Implementation
-    func getCountry(_ ip: String) -> String {
+    private func cacheCountry(for ip: String) -> String {
         var truncatedIP = ip
         func getCountryInternal() -> String {
             if let country = countryNamesCache[ip] { return country }
@@ -42,7 +42,7 @@ final class IP2Country {
         return getCountryInternal()
     }
 
-    @objc private func preloadCountriesIfNeeded() {
+    @objc private func populateCacheIfNeeded() {
         DispatchQueue.global(qos: .userInitiated).async {
             if OnionRequestAPI.paths.count < OnionRequestAPI.pathCount {
                 let storage = OWSPrimaryStorage.shared()
@@ -53,7 +53,10 @@ final class IP2Country {
             guard OnionRequestAPI.paths.count >= OnionRequestAPI.pathCount else { return }
             let pathToDisplay = OnionRequestAPI.paths.first!
             pathToDisplay.forEach { snode in
-                let _ = self.getCountry(snode.ip) // Preload if needed
+                let _ = self.cacheCountry(for: snode.ip) // Preload if needed
+            }
+            DispatchQueue.main.async {
+                NotificationCenter.default.post(name: .onionRequestPathCountriesLoaded, object: nil)
             }
             print("[Loki] Finished preloading onion request path countries.")
         }
diff --git a/Signal/src/Loki/View Controllers/PathVC.swift b/Signal/src/Loki/View Controllers/PathVC.swift
index ae5e20a96..f4c984f1b 100644
--- a/Signal/src/Loki/View Controllers/PathVC.swift	
+++ b/Signal/src/Loki/View Controllers/PathVC.swift	
@@ -92,6 +92,7 @@ final class PathVC : BaseVC {
         let notificationCenter = NotificationCenter.default
         notificationCenter.addObserver(self, selector: #selector(handleBuildingPathsNotification), name: .buildingPaths, object: nil)
         notificationCenter.addObserver(self, selector: #selector(handlePathsBuiltNotification), name: .pathsBuilt, object: nil)
+        notificationCenter.addObserver(self, selector: #selector(handleOnionRequestPathCountriesLoadedNotification), name: .onionRequestPathCountriesLoaded, object: nil)
     }
 
     deinit {
@@ -101,6 +102,7 @@ final class PathVC : BaseVC {
     // MARK: Updating
     @objc private func handleBuildingPathsNotification() { update() }
     @objc private func handlePathsBuiltNotification() { update() }
+    @objc private func handleOnionRequestPathCountriesLoadedNotification() { update() }
 
     private func update() {
         pathStackView.arrangedSubviews.forEach { $0.removeFromSuperview() }
@@ -155,7 +157,7 @@ final class PathVC : BaseVC {
     }
 
     private func getPathRow(snode: LokiAPITarget, location: LineView.Location, dotAnimationStartDelay: Double, dotAnimationRepeatInterval: Double, isGuardSnode: Bool) -> UIStackView {
-        let country = IP2Country.shared.getCountry(snode.ip)
+        let country = IP2Country.shared.countryNamesCache[snode.ip] ?? "Resolving..."
         let title = isGuardSnode ? NSLocalizedString("Guard Node", comment: "") : NSLocalizedString("Service Node", comment: "")
         return getPathRow(title: title, subtitle: country, location: location, dotAnimationStartDelay: dotAnimationStartDelay, dotAnimationRepeatInterval: dotAnimationRepeatInterval)
     }
diff --git a/SignalServiceKit/src/Loki/Utilities/Notification+Loki.swift b/SignalServiceKit/src/Loki/Utilities/Notification+Loki.swift
index 27de5db22..1ccfc488c 100644
--- a/SignalServiceKit/src/Loki/Utilities/Notification+Loki.swift
+++ b/SignalServiceKit/src/Loki/Utilities/Notification+Loki.swift
@@ -23,6 +23,7 @@ public extension Notification.Name {
     // Onion requests
     public static let buildingPaths = Notification.Name("buildingPaths")
     public static let pathsBuilt = Notification.Name("pathsBuilt")
+    public static let onionRequestPathCountriesLoaded = Notification.Name("onionRequestPathCountriesLoaded")
 }
 
 @objc public extension NSNotification {
@@ -49,4 +50,5 @@ public extension Notification.Name {
     // Onion requests
     @objc public static let buildingPaths = Notification.Name.buildingPaths.rawValue as NSString
     @objc public static let pathsBuilt = Notification.Name.pathsBuilt.rawValue as NSString
+    @objc public static let onionRequestPathCountriesLoaded = Notification.Name.onionRequestPathCountriesLoaded.rawValue as NSString
 }