Add retryable NSE auth fetches with fresh network sessions

Author: adrastaea

packages/mobile/ios/QuietNotificationServiceExtension/NSEAuthService.swift

@@ -3,21 +3,63 @@ import os.log
 
 private let authLog = OSLog(subsystem: "com.quietmobile.QuietNotificationServiceExtension", category: "NSEAuthService")
 
+struct NSEAuthTokenCacheKey: Hashable {
+    let qssUrl: URL
+    let teamId: String
+}
+
+final class NSEAuthTokenCache {
+    private let lock = NSLock()
+    private var tokens: [NSEAuthTokenCacheKey: (token: String, expiry: Date)] = [:]
+
+    func token(for qssUrl: URL, teamId: String, now: Date = Date()) -> (token: String, expiry: Date)? {
+        let key = NSEAuthTokenCacheKey(qssUrl: qssUrl, teamId: teamId)
+        lock.lock()
+        defer { lock.unlock() }
+
+        guard let cached = tokens[key] else {
+            return nil
+        }
+
+        guard cached.expiry > now else {
+            tokens.removeValue(forKey: key)
+            return nil
+        }
+
+        return cached
+    }
+
+    func store(token: String, expiresIn: Int, for qssUrl: URL, teamId: String, now: Date = Date()) {
+        let key = NSEAuthTokenCacheKey(qssUrl: qssUrl, teamId: teamId)
+        let expiry = now.addingTimeInterval(TimeInterval(expiresIn) - 30)
+        lock.lock()
+        tokens[key] = (token: token, expiry: expiry)
+        lock.unlock()
+    }
+
+    func removeToken(for qssUrl: URL, teamId: String) {
+        let key = NSEAuthTokenCacheKey(qssUrl: qssUrl, teamId: teamId)
+        lock.lock()
+        tokens.removeValue(forKey: key)
+        lock.unlock()
+    }
+}
+
 class NSEAuthService {
     private let client: NSENetworkClient
     private let crypto: DeviceCryptography
+    private let tokenCache: NSEAuthTokenCache
 
-    private var tokenCache: [String: (token: String, expiry: Date)] = [:]
-
-    init(client: NSENetworkClient, crypto: DeviceCryptography) {
+    init(client: NSENetworkClient, crypto: DeviceCryptography, tokenCache: NSEAuthTokenCache = NSEAuthTokenCache()) {
         self.client = client
         self.crypto = crypto
+        self.tokenCache = tokenCache
     }
 
     // MARK: - Full auth flow
 
     func authenticate(deviceId: String, teamId: String) async throws -> String {
-        if let cached = tokenCache[teamId], cached.expiry > Date() {
+        if let cached = tokenCache.token(for: client.baseURL, teamId: teamId) {
             os_log("authenticate: using cached token for teamId=%{public}@, expires=%{public}@",
                    log: authLog, type: .debug, teamId, "\(cached.expiry)")
             return cached.token
@@ -42,7 +84,7 @@ class NSEAuthService {
         )
         os_log("authenticate: token received, expiresIn=%{public}d", log: authLog, type: .info, tokenResp.expiresIn)
 
-        tokenCache[teamId] = (token: tokenResp.token, expiry: Date().addingTimeInterval(TimeInterval(tokenResp.expiresIn) - 30))
+        tokenCache.store(token: tokenResp.token, expiresIn: tokenResp.expiresIn, for: client.baseURL, teamId: teamId)
 
         return tokenResp.token
     }
@@ -63,12 +105,11 @@ class NSEAuthService {
         } catch NSEAuthError.logFetchFailed(let statusCode) where statusCode == 401 {
             os_log("fetchNewEntries: token rejected (401) for teamId=%{public}@, evicting cache and retrying",
                    log: authLog, type: .info, teamId)
-            tokenCache.removeValue(forKey: teamId)
+            tokenCache.removeToken(for: client.baseURL, teamId: teamId)
             let freshToken = try await authenticate(deviceId: deviceId, teamId: teamId)
             let resp = try await client.fetchLogEntries(teamId: teamId, afterSeq: afterSeq, token: freshToken)
             os_log("fetchNewEntries: retry succeeded, received %{public}d entries", log: authLog, type: .info, resp.entries.count)
             return resp
         }
     }
 }
-

packages/mobile/ios/QuietNotificationServiceExtension/NSEModels.swift

@@ -28,6 +28,57 @@ enum NSEAuthError: Error, LocalizedError {
     }
 }
 
+extension NSEAuthError {
+    private static let retryableURLCodes: Set<Int> = [
+        URLError.cannotConnectToHost.rawValue,
+        URLError.networkConnectionLost.rawValue,
+        URLError.timedOut.rawValue,
+        URLError.notConnectedToInternet.rawValue,
+        URLError.cannotFindHost.rawValue,
+        URLError.dnsLookupFailed.rawValue,
+        URLError.resourceUnavailable.rawValue,
+        URLError.callIsActive.rawValue,
+        URLError.dataNotAllowed.rawValue
+    ]
+
+    private static let retryablePOSIXCodes: Set<Int> = [53, 57, 60, 61, 64, 65]
+
+    var isRetryableNetworkFailure: Bool {
+        guard case .networkError(let error) = self else {
+            return false
+        }
+        return Self.isRetryableNetworkFailure(error)
+    }
+
+    private static func isRetryableNetworkFailure(_ error: Error) -> Bool {
+        let nsError = error as NSError
+
+        if nsError.domain == NSURLErrorDomain, retryableURLCodes.contains(nsError.code) {
+            return true
+        }
+
+        if nsError.domain == NSPOSIXErrorDomain, retryablePOSIXCodes.contains(nsError.code) {
+            return true
+        }
+
+        if let streamCode = nsError.userInfo["_kCFStreamErrorCodeKey"] as? Int,
+           retryablePOSIXCodes.contains(streamCode) {
+            return true
+        }
+
+        if let underlying = nsError.userInfo[NSUnderlyingErrorKey] as? NSError {
+            if underlying.domain == NSURLErrorDomain, retryableURLCodes.contains(underlying.code) {
+                return true
+            }
+            if underlying.domain == NSPOSIXErrorDomain, retryablePOSIXCodes.contains(underlying.code) {
+                return true
+            }
+        }
+
+        return false
+    }
+}
+
 // MARK: - Challenge
 
 struct ChallengePayload: Codable {

packages/mobile/ios/QuietNotificationServiceExtension/NSENetworkClient.swift

@@ -16,17 +16,18 @@ class NSENetworkClient {
 
     private static let decoder = JSONDecoder()
 
-    // Dedicated session with tight timeouts suitable for an NSE (30-second budget).
-    private static let defaultSession: URLSession = {
+    // Each client gets a fresh session so path transitions cannot poison pooled connections.
+    private static func makeDefaultSession() -> URLSession {
         let config = URLSessionConfiguration.ephemeral
         config.timeoutIntervalForRequest = 10
         config.timeoutIntervalForResource = 20
+        config.waitsForConnectivity = true
         return URLSession(configuration: config)
-    }()
+    }
 
     init(baseURL: URL, session: URLSession? = nil) {
         self.baseURL = baseURL
-        self.session = session ?? NSENetworkClient.defaultSession
+        self.session = session ?? NSENetworkClient.makeDefaultSession()
     }
 
     // MARK: - POST /nse-auth/challenge

packages/mobile/ios/QuietNotificationServiceExtension/NotificationService.swift

@@ -9,12 +9,19 @@ class NotificationService: UNNotificationServiceExtension {
         let message: NSEDecryptedNotificationMessage
     }
 
+    private static let retryDelaysNanoseconds: [UInt64] = [
+        250_000_000,
+        750_000_000
+    ]
+
+    private static let maxRetryWindow: TimeInterval = 5
+
     var contentHandler: ((UNNotificationContent) -> Void)?
     var bestAttemptContent: UNMutableNotificationContent?
     var fetchTask: Task<Void, Never>?
 
     private let crypto = NSECryptoService()
-    private var authCache: [URL: NSEAuthService] = [:]
+    private let tokenCache = NSEAuthTokenCache()
 
     private static func channelName(from channelId: String) -> String {
         guard let separatorIndex = channelId.firstIndex(of: "_") else {
@@ -76,23 +83,11 @@ class NotificationService: UNNotificationServiceExtension {
                log: nseLog, type: .info, teamId, qssUrlString)
 
         do {
-            let auth: NSEAuthService
-            if let cached = authCache[qssUrl] {
-                os_log("fetchAndUpdate: using cached NSEAuthService for %{public}@", log: nseLog, type: .debug, qssUrlString)
-                auth = cached
-            } else {
-                os_log("fetchAndUpdate: creating new NSEAuthService for %{public}@", log: nseLog, type: .debug, qssUrlString)
-                let client = NSENetworkClient(baseURL: qssUrl)
-                let newAuth = NSEAuthService(client: client, crypto: crypto)
-                authCache[qssUrl] = newAuth
-                auth = newAuth
-            }
-
             let afterSeq = SharedDefaults.getLastSyncSeq()
             os_log("fetchAndUpdate: fetching entries afterSeq=%{public}lld",
                    log: nseLog, type: .info, afterSeq)
 
-            let response = try await auth.fetchNewEntries(teamId: teamId, afterSeq: afterSeq)
+            let response = try await fetchEntriesWithRetry(qssUrl: qssUrl, teamId: teamId, afterSeq: afterSeq)
             let entries = response.entries
             let baselineSeq = afterSeq
             os_log("fetchAndUpdate: fetched %{public}d entries",
@@ -189,6 +184,61 @@ class NotificationService: UNNotificationServiceExtension {
         }
     }
 
+    private func fetchEntriesWithRetry(qssUrl: URL, teamId: String, afterSeq: Int64) async throws -> LogEntriesResponse {
+        let startedAt = Date()
+        var retryIndex = 0
+
+        while true {
+            guard !Task.isCancelled else {
+                throw CancellationError()
+            }
+
+            do {
+                let auth = makeAuthService(qssUrl: qssUrl)
+                return try await auth.fetchNewEntries(teamId: teamId, afterSeq: afterSeq)
+            } catch {
+                guard shouldRetryFetch(error: error, startedAt: startedAt, retryIndex: retryIndex) else {
+                    throw error
+                }
+
+                let delay = Self.retryDelaysNanoseconds[retryIndex]
+                retryIndex += 1
+                os_log(
+                    "fetchAndUpdate: retrying full auth fetch after transient network error (%{public}d/%{public}d): %{public}@",
+                    log: nseLog,
+                    type: .info,
+                    retryIndex,
+                    Self.retryDelaysNanoseconds.count,
+                    String(describing: error)
+                )
+                try await Task.sleep(nanoseconds: delay)
+            }
+        }
+    }
+
+    private func makeAuthService(qssUrl: URL) -> NSEAuthService {
+        os_log("fetchAndUpdate: creating fresh NSEAuthService for %{public}@",
+               log: nseLog, type: .debug, qssUrl.absoluteString)
+        let client = NSENetworkClient(baseURL: qssUrl)
+        return NSEAuthService(client: client, crypto: crypto, tokenCache: tokenCache)
+    }
+
+    private func shouldRetryFetch(error: Error, startedAt: Date, retryIndex: Int) -> Bool {
+        guard retryIndex < Self.retryDelaysNanoseconds.count else {
+            return false
+        }
+
+        guard Date().timeIntervalSince(startedAt) < Self.maxRetryWindow else {
+            return false
+        }
+
+        guard let authError = error as? NSEAuthError else {
+            return false
+        }
+
+        return authError.isRetryableNetworkFailure
+    }
+
     private func applyNotificationMessage(_ message: NSEDecryptedNotificationMessage, to content: UNMutableNotificationContent) {
         content.title = "#\(Self.channelName(from: message.channelId))"
         content.body = message.body