From 92aba98cb7689eeba33e025c0cb71325e0c82350 Mon Sep 17 00:00:00 2001
From: gmbnt <andriesseniels@gmail.com>
Date: Thu, 2 Apr 2020 16:05:56 +1100
Subject: [PATCH] Ensure target snode isn't included in onion request path

---
 .../Loki/API/Onion Requests/OnionRequestAPI.swift    | 12 ++++++------
 .../API/Onion Requests/OnionRequestAPITests.swift    | 12 ------------
 2 files changed, 6 insertions(+), 18 deletions(-)

diff --git a/SignalServiceKit/src/Loki/API/Onion Requests/OnionRequestAPI.swift b/SignalServiceKit/src/Loki/API/Onion Requests/OnionRequestAPI.swift
index 6cb8f4260..344cceb0c 100644
--- a/SignalServiceKit/src/Loki/API/Onion Requests/OnionRequestAPI.swift	
+++ b/SignalServiceKit/src/Loki/API/Onion Requests/OnionRequestAPI.swift	
@@ -123,16 +123,16 @@ internal enum OnionRequestAPI {
     }
 
     /// Returns a `Path` to be used for building an onion request. Builds new paths as needed.
-    ///
-    /// - Note: Exposed for testing purposes.
-    internal static func getPath() -> Promise<Path> {
+    private static func getPath(excluding snode: LokiAPITarget) -> Promise<Path> {
         guard pathSize >= 1 else { preconditionFailure("Cannot build path of size zero.") }
         // randomElement() uses the system's default random generator, which is cryptographically secure
         if paths.count >= pathCount {
-            return Promise<Path> { $0.fulfill(paths.randomElement()!) }
+            return Promise<Path> { seal in
+                seal.fulfill(paths.filter { !$0.contains(snode) }.randomElement()!)
+            }
         } else {
             return buildPaths().map(on: workQueue) { paths in
-                let path = paths.randomElement()!
+                let path = paths.filter { !$0.contains(snode) }.randomElement()!
                 OnionRequestAPI.paths = paths
                 return path
             }
@@ -148,7 +148,7 @@ internal enum OnionRequestAPI {
         var guardSnode: LokiAPITarget!
         var targetSnodeSymmetricKey: Data! // Needed by invoke(_:on:with:) to decrypt the response sent back by the target snode
         var encryptionResult: EncryptionResult!
-        return getPath().then(on: workQueue) { path -> Promise<EncryptionResult> in
+        return getPath(excluding: snode).then(on: workQueue) { path -> Promise<EncryptionResult> in
             guardSnode = path.first!
             // Encrypt in reverse order, i.e. the target snode first
             return encrypt(payload, forTargetSnode: snode).then(on: workQueue) { r -> Promise<EncryptionResult> in
diff --git a/SignalServiceKit/src/Loki/API/Onion Requests/OnionRequestAPITests.swift b/SignalServiceKit/src/Loki/API/Onion Requests/OnionRequestAPITests.swift
index 0ebef0106..3620ae028 100644
--- a/SignalServiceKit/src/Loki/API/Onion Requests/OnionRequestAPITests.swift	
+++ b/SignalServiceKit/src/Loki/API/Onion Requests/OnionRequestAPITests.swift	
@@ -4,18 +4,6 @@ import XCTest
 
 class OnionRequestAPITests : XCTestCase {
 
-    func testPathBuilding() {
-        let semaphore = DispatchSemaphore(value: 0)
-        var error: Error? = nil
-        let _ = OnionRequestAPI.getPath().done(on: OnionRequestAPI.workQueue) { _ in
-            semaphore.signal()
-        }.catch(on: OnionRequestAPI.workQueue) {
-            error = $0; semaphore.signal()
-        }
-        semaphore.wait()
-        XCTAssert(error == nil)
-    }
-
     /// Builds a path and then routes the same request through it several times. Logs the number of successes
     /// versus the number of failures.
     func testOnionRequestSending() {