From e0c9f430260fc8f20cc37cd5873b003e654353cd Mon Sep 17 00:00:00 2001
From: Stefan Ceriu <stefanc@matrix.org>
Date: Thu, 15 Feb 2024 09:17:36 +0200
Subject: [PATCH] Fix localization bundle and attributed string builder crashes
 on unsynchronised access to their internal caches

---
 .../Sources/Other/Extensions/Bundle.swift     | 32 ++++++++++++-----
 .../HTMLParsing/AttributedStringBuilder.swift | 36 ++++++++++++++-----
 2 files changed, 52 insertions(+), 16 deletions(-)

diff --git a/ElementX/Sources/Other/Extensions/Bundle.swift b/ElementX/Sources/Other/Extensions/Bundle.swift
index f41fc3e50..465c43a7e 100644
--- a/ElementX/Sources/Other/Extensions/Bundle.swift
+++ b/ElementX/Sources/Other/Extensions/Bundle.swift
@@ -32,13 +32,17 @@ public extension Bundle {
     
     // MARK: - Localisation
     
-    private static var cachedLocalizationBundles = [String: Bundle]()
+    /// Overrides `Bundle.app.preferredLocalizations` for testing translations.
+    static var overrideLocalizations: [String]?
+    
+    private static let cacheDispatchQueue = DispatchQueue(label: "io.element.elementx.localization_bundle_cache")
+    private static var cachedBundles = [String: Bundle]()
     
     /// Get an lproj language bundle from the receiver bundle.
     /// - Parameter language: The language to try to load.
     /// - Returns: The lproj bundle if found otherwise nil.
     static func lprojBundle(for language: String) -> Bundle? {
-        if let bundle = cachedLocalizationBundles[language] {
+        if let bundle = cachedValue(forKey: language) {
             return bundle
         }
         
@@ -48,13 +52,25 @@ public extension Bundle {
         
         let bundle = Bundle(url: lprojURL)
         
-        DispatchQueue.main.async {
-            cachedLocalizationBundles[language] = bundle
-        }
+        cacheValue(bundle, forKey: language)
         
         return bundle
     }
-
-    /// Overrides `Bundle.app.preferredLocalizations` for testing translations.
-    static var overrideLocalizations: [String]?
+    
+    // MARK: - Private
+    
+    private static func cacheValue(_ value: Bundle?, forKey key: String) {
+        cacheDispatchQueue.sync {
+            cachedBundles[key] = value
+        }
+    }
+    
+    private static func cachedValue(forKey key: String) -> Bundle? {
+        var result: Bundle?
+        cacheDispatchQueue.sync {
+            result = cachedBundles[key]
+        }
+        
+        return result
+    }
 }
diff --git a/ElementX/Sources/Other/HTMLParsing/AttributedStringBuilder.swift b/ElementX/Sources/Other/HTMLParsing/AttributedStringBuilder.swift
index 681bdc930..bb417fb33 100644
--- a/ElementX/Sources/Other/HTMLParsing/AttributedStringBuilder.swift
+++ b/ElementX/Sources/Other/HTMLParsing/AttributedStringBuilder.swift
@@ -26,6 +26,8 @@ struct AttributedStringBuilder: AttributedStringBuilderProtocol {
     private let mentionBuilder: MentionBuilderProtocol
     
     private static let defaultKey = "default"
+    
+    private static let cacheDispatchQueue = DispatchQueue(label: "io.element.elementx.attributed_string_builder_cache")
     private static var caches: [String: LRUCache<String, AttributedString>] = [:]
 
     static func invalidateCaches() {
@@ -34,9 +36,6 @@ struct AttributedStringBuilder: AttributedStringBuilderProtocol {
     
     init(cacheKey: String = defaultKey, permalinkBaseURL: URL, mentionBuilder: MentionBuilderProtocol) {
         self.cacheKey = cacheKey
-        if Self.caches[cacheKey] == nil {
-            Self.caches[cacheKey] = LRUCache<String, AttributedString>(countLimit: 1000)
-        }
         self.permalinkBaseURL = permalinkBaseURL
         self.mentionBuilder = mentionBuilder
     }
@@ -46,7 +45,7 @@ struct AttributedStringBuilder: AttributedStringBuilderProtocol {
             return nil
         }
         
-        if let cached = Self.caches[cacheKey]?.value(forKey: string) {
+        if let cached = Self.cachedValue(forKey: string, cacheKey: cacheKey) {
             return cached
         }
 
@@ -56,7 +55,8 @@ struct AttributedStringBuilder: AttributedStringBuilderProtocol {
         removeLinkColors(mutableAttributedString)
         
         let result = try? AttributedString(mutableAttributedString, including: \.elementX)
-        Self.caches[cacheKey]?.setValue(result, forKey: string)
+        Self.cacheValue(result, forKey: string, cacheKey: cacheKey)
+        
         return result
     }
         
@@ -72,10 +72,10 @@ struct AttributedStringBuilder: AttributedStringBuilderProtocol {
             return nil
         }
         
-        if let cached = Self.caches[cacheKey]?.value(forKey: originalHTMLString) {
+        if let cached = Self.cachedValue(forKey: originalHTMLString, cacheKey: cacheKey) {
             return cached
         }
-        
+                
         let htmlString = originalHTMLString.replacingHtmlBreaksOccurrences()
         
         guard let data = htmlString.data(using: .utf8) else {
@@ -115,12 +115,32 @@ struct AttributedStringBuilder: AttributedStringBuilderProtocol {
         removeDTCoreTextArtifacts(mutableAttributedString)
         
         let result = try? AttributedString(mutableAttributedString, including: \.elementX)
-        Self.caches[cacheKey]?.setValue(result, forKey: htmlString)
+        Self.cacheValue(result, forKey: htmlString, cacheKey: cacheKey)
+        
         return result
     }
     
     // MARK: - Private
     
+    private static func cacheValue(_ value: AttributedString?, forKey key: String, cacheKey: String) {
+        cacheDispatchQueue.sync {
+            if caches[cacheKey] == nil {
+                caches[cacheKey] = LRUCache<String, AttributedString>(countLimit: 1000)
+            }
+            
+            caches[cacheKey]?.setValue(value, forKey: key)
+        }
+    }
+    
+    private static func cachedValue(forKey key: String, cacheKey: String) -> AttributedString? {
+        var result: AttributedString?
+        cacheDispatchQueue.sync {
+            result = caches[cacheKey]?.value(forKey: key)
+        }
+        
+        return result
+    }
+    
     private func replaceMarkedBlockquotes(_ attributedString: NSMutableAttributedString) {
         // According to blockquotes in the string, DTCoreText can apply 2 policies:
         //     - define a `DTTextBlocksAttribute` attribute on a <blockquote> block