From 5b50c9f7fb09e84b77c2124f71e050c366ca9c07 Mon Sep 17 00:00:00 2001 From: nielsandriesse Date: Mon, 1 Jun 2020 10:10:08 +1000 Subject: [PATCH] Enforce minimum snode pool size & clean --- .../src/Loki/API/LokiAPI+SwarmAPI.swift | 31 ++++++++++--------- .../src/Loki/API/LokiMessage.swift | 2 +- .../src/Loki/API/LokiMessageWrapper.swift | 2 +- .../src/Loki/API/LokiPoller.swift | 2 +- .../API/Onion Requests/OnionRequestAPI.swift | 16 +++++----- .../src/Loki/API/SignalMessage.swift | 8 ++--- 6 files changed, 31 insertions(+), 30 deletions(-) diff --git a/SignalServiceKit/src/Loki/API/LokiAPI+SwarmAPI.swift b/SignalServiceKit/src/Loki/API/LokiAPI+SwarmAPI.swift index 95c39304a..4bbb2749c 100644 --- a/SignalServiceKit/src/Loki/API/LokiAPI+SwarmAPI.swift +++ b/SignalServiceKit/src/Loki/API/LokiAPI+SwarmAPI.swift @@ -4,18 +4,19 @@ public extension LokiAPI { private static var snodeVersion: [LokiAPITarget:String] = [:] /// Only ever modified from `LokiAPI.errorHandlingQueue` to avoid race conditions. - internal static var failureCount: [LokiAPITarget:UInt] = [:] + internal static var snodeFailureCount: [LokiAPITarget:UInt] = [:] // MARK: Settings - private static let minimumSnodeCount = 2 - private static let targetSnodeCount = 3 + private static let minimumSnodePoolCount = 32 + private static let minimumSwarmSnodeCount = 2 + private static let targetSwarmSnodeCount = 3 - internal static let failureThreshold = 2 + internal static let snodeFailureThreshold = 2 // MARK: Caching internal static var swarmCache: [String:[LokiAPITarget]] = [:] - internal static func dropSnodeIfNeeded(_ target: LokiAPITarget, hexEncodedPublicKey: String) { + internal static func dropSnodeFromSwarmIfNeeded(_ target: LokiAPITarget, hexEncodedPublicKey: String) { let swarm = LokiAPI.swarmCache[hexEncodedPublicKey] if var swarm = swarm, let index = swarm.firstIndex(of: target) { swarm.remove(at: index) @@ -40,12 +41,12 @@ public extension LokiAPI { // MARK: Internal API internal static func getRandomSnode() -> Promise { - if snodePool.isEmpty { + if snodePool.count < minimumSnodePoolCount { storage.dbReadConnection.read { transaction in snodePool = storage.getSnodePool(in: transaction) } } - if snodePool.isEmpty { + if snodePool.count < minimumSnodePoolCount { let target = seedNodePool.randomElement()! let url = "\(target)/json_rpc" let parameters: JSON = [ @@ -95,7 +96,7 @@ public extension LokiAPI { } internal static func getSwarm(for hexEncodedPublicKey: String) -> Promise<[LokiAPITarget]> { - if let cachedSwarm = swarmCache[hexEncodedPublicKey], cachedSwarm.count >= minimumSnodeCount { + if let cachedSwarm = swarmCache[hexEncodedPublicKey], cachedSwarm.count >= minimumSwarmSnodeCount { return Promise<[LokiAPITarget]> { $0.fulfill(cachedSwarm) } } else { let parameters: [String:Any] = [ "pubKey" : hexEncodedPublicKey ] @@ -105,7 +106,7 @@ public extension LokiAPI { internal static func getTargetSnodes(for hexEncodedPublicKey: String) -> Promise<[LokiAPITarget]> { // shuffled() uses the system's default random generator, which is cryptographically secure - return getSwarm(for: hexEncodedPublicKey).map { Array($0.shuffled().prefix(targetSnodeCount)) } + return getSwarm(for: hexEncodedPublicKey).map { Array($0.shuffled().prefix(targetSwarmSnodeCount)) } } internal static func getFileServerProxy() -> Promise { @@ -169,13 +170,13 @@ internal extension Promise { switch error.statusCode { case 0, 400, 500, 503: // The snode is unreachable - let oldFailureCount = LokiAPI.failureCount[target] ?? 0 + let oldFailureCount = LokiAPI.snodeFailureCount[target] ?? 0 let newFailureCount = oldFailureCount + 1 - LokiAPI.failureCount[target] = newFailureCount + LokiAPI.snodeFailureCount[target] = newFailureCount print("[Loki] Couldn't reach snode at: \(target); setting failure count to \(newFailureCount).") - if newFailureCount >= LokiAPI.failureThreshold { + if newFailureCount >= LokiAPI.snodeFailureThreshold { print("[Loki] Failure threshold reached for: \(target); dropping it.") - LokiAPI.dropSnodeIfNeeded(target, hexEncodedPublicKey: hexEncodedPublicKey) // Remove it from the swarm cache associated with the given public key + LokiAPI.dropSnodeFromSwarmIfNeeded(target, hexEncodedPublicKey: hexEncodedPublicKey) // Remove it from the swarm cache associated with the given public key LokiAPI.snodePool.remove(target) // Remove it from the snode pool // Dispatch async on the main queue to avoid nested write transactions DispatchQueue.main.async { @@ -184,7 +185,7 @@ internal extension Promise { storage.dropSnode(target, in: transaction) } } - LokiAPI.failureCount[target] = 0 + LokiAPI.snodeFailureCount[target] = 0 } case 406: print("[Loki] The user's clock is out of sync with the service node network.") @@ -192,7 +193,7 @@ internal extension Promise { case 421: // The snode isn't associated with the given public key anymore print("[Loki] Invalidating swarm for: \(hexEncodedPublicKey).") - LokiAPI.dropSnodeIfNeeded(target, hexEncodedPublicKey: hexEncodedPublicKey) + LokiAPI.dropSnodeFromSwarmIfNeeded(target, hexEncodedPublicKey: hexEncodedPublicKey) case 432: // The PoW difficulty is too low if case LokiHTTPClient.HTTPError.networkError(_, let result, _) = error, let json = result as? JSON, let powDifficulty = json["difficulty"] as? Int { diff --git a/SignalServiceKit/src/Loki/API/LokiMessage.swift b/SignalServiceKit/src/Loki/API/LokiMessage.swift index 620bfd396..2e72ba112 100644 --- a/SignalServiceKit/src/Loki/API/LokiMessage.swift +++ b/SignalServiceKit/src/Loki/API/LokiMessage.swift @@ -33,7 +33,7 @@ public struct LokiMessage { do { let wrappedMessage = try LokiMessageWrapper.wrap(message: signalMessage) let data = wrappedMessage.base64EncodedString() - let destination = signalMessage.recipientID + let destination = signalMessage.recipientPublicKey var ttl = TTLUtilities.fallbackMessageTTL if let messageTTL = signalMessage.ttl, messageTTL > 0 { ttl = UInt64(messageTTL) } let isPing = signalMessage.isPing diff --git a/SignalServiceKit/src/Loki/API/LokiMessageWrapper.swift b/SignalServiceKit/src/Loki/API/LokiMessageWrapper.swift index 86c3cfce2..84b3a5669 100644 --- a/SignalServiceKit/src/Loki/API/LokiMessageWrapper.swift +++ b/SignalServiceKit/src/Loki/API/LokiMessageWrapper.swift @@ -31,7 +31,7 @@ public enum LokiMessageWrapper { private static func createEnvelope(around message: SignalMessage) throws -> SSKProtoEnvelope { do { let builder = SSKProtoEnvelope.builder(type: message.type, timestamp: message.timestamp) - builder.setSource(message.senderID) + builder.setSource(message.senderPublicKey) builder.setSourceDevice(message.senderDeviceID) if let content = try Data(base64Encoded: message.content) { builder.setContent(content) diff --git a/SignalServiceKit/src/Loki/API/LokiPoller.swift b/SignalServiceKit/src/Loki/API/LokiPoller.swift index e455fe794..9ebdf1c84 100644 --- a/SignalServiceKit/src/Loki/API/LokiPoller.swift +++ b/SignalServiceKit/src/Loki/API/LokiPoller.swift @@ -84,7 +84,7 @@ public final class LokiPoller : NSObject { self?.pollCount = 0 } else { print("[Loki] Polling \(nextSnode) failed; dropping it and switching to next snode.") - LokiAPI.dropSnodeIfNeeded(nextSnode, hexEncodedPublicKey: userHexEncodedPublicKey) + LokiAPI.dropSnodeFromSwarmIfNeeded(nextSnode, hexEncodedPublicKey: userHexEncodedPublicKey) } self?.pollNextSnode(seal: seal) } diff --git a/SignalServiceKit/src/Loki/API/Onion Requests/OnionRequestAPI.swift b/SignalServiceKit/src/Loki/API/Onion Requests/OnionRequestAPI.swift index 6ce07203d..e54b59385 100644 --- a/SignalServiceKit/src/Loki/API/Onion Requests/OnionRequestAPI.swift +++ b/SignalServiceKit/src/Loki/API/Onion Requests/OnionRequestAPI.swift @@ -6,10 +6,10 @@ public enum OnionRequestAPI { /// - Note: Must only be modified from `LokiAPI.workQueue`. public static var guardSnodes: Set = [] /// - Note: Must only be modified from `LokiAPI.workQueue`. - public static var paths: [Path] = [] + public static var paths: [Path] = [] // Not a set to ensure we consistently show the same path to the user private static var snodePool: Set { - let unreliableSnodes = Set(LokiAPI.failureCount.keys) + let unreliableSnodes = Set(LokiAPI.snodeFailureCount.keys) return LokiAPI.snodePool.subtracting(unreliableSnodes) } @@ -269,13 +269,13 @@ private extension Promise where T == JSON { switch statusCode { case 0, 400, 500, 503: // The snode is unreachable - let oldFailureCount = LokiAPI.failureCount[snode] ?? 0 + let oldFailureCount = LokiAPI.snodeFailureCount[snode] ?? 0 let newFailureCount = oldFailureCount + 1 - LokiAPI.failureCount[snode] = newFailureCount + LokiAPI.snodeFailureCount[snode] = newFailureCount print("[Loki] Couldn't reach snode at: \(snode); setting failure count to \(newFailureCount).") - if newFailureCount >= LokiAPI.failureThreshold { + if newFailureCount >= LokiAPI.snodeFailureThreshold { print("[Loki] Failure threshold reached for: \(snode); dropping it.") - LokiAPI.dropSnodeIfNeeded(snode, hexEncodedPublicKey: hexEncodedPublicKey) // Remove it from the swarm cache associated with the given public key + LokiAPI.dropSnodeFromSwarmIfNeeded(snode, hexEncodedPublicKey: hexEncodedPublicKey) // Remove it from the swarm cache associated with the given public key LokiAPI.snodePool.remove(snode) // Remove it from the snode pool // Dispatch async on the main queue to avoid nested write transactions DispatchQueue.main.async { @@ -284,7 +284,7 @@ private extension Promise where T == JSON { storage.dropSnode(snode, in: transaction) } } - LokiAPI.failureCount[snode] = 0 + LokiAPI.snodeFailureCount[snode] = 0 } case 406: print("[Loki] The user's clock is out of sync with the service node network.") @@ -292,7 +292,7 @@ private extension Promise where T == JSON { case 421: // The snode isn't associated with the given public key anymore print("[Loki] Invalidating swarm for: \(hexEncodedPublicKey).") - LokiAPI.dropSnodeIfNeeded(snode, hexEncodedPublicKey: hexEncodedPublicKey) + LokiAPI.dropSnodeFromSwarmIfNeeded(snode, hexEncodedPublicKey: hexEncodedPublicKey) case 432: // The proof of work difficulty is too low if let powDifficulty = json["difficulty"] as? Int { diff --git a/SignalServiceKit/src/Loki/API/SignalMessage.swift b/SignalServiceKit/src/Loki/API/SignalMessage.swift index f345fc01f..c4cb88249 100644 --- a/SignalServiceKit/src/Loki/API/SignalMessage.swift +++ b/SignalServiceKit/src/Loki/API/SignalMessage.swift @@ -3,10 +3,10 @@ public final class SignalMessage : NSObject { @objc public let type: SSKProtoEnvelope.SSKProtoEnvelopeType @objc public let timestamp: UInt64 - @objc public let senderID: String + @objc public let senderPublicKey: String @objc public let senderDeviceID: UInt32 @objc public let content: String - @objc public let recipientID: String + @objc public let recipientPublicKey: String @objc(ttl) public let objc_ttl: UInt64 @objc public let isPing: Bool @@ -17,10 +17,10 @@ public final class SignalMessage : NSObject { content: String, recipientID: String, ttl: UInt64, isPing: Bool) { self.type = type self.timestamp = timestamp - self.senderID = senderID + self.senderPublicKey = senderID self.senderDeviceID = senderDeviceID self.content = content - self.recipientID = recipientID + self.recipientPublicKey = recipientID self.objc_ttl = ttl self.isPing = isPing super.init()