From 0f84a5581a829fe42f193c0e45cab8fffd26c4fc Mon Sep 17 00:00:00 2001
From: nielsandriesse <andriesseniels@gmail.com>
Date: Fri, 8 May 2020 14:38:05 +1000
Subject: [PATCH] Fix race condition

---
 .../Protocol/Friend Requests/FriendRequestProtocol.swift     | 2 ++
 .../Session Management/SessionManagementProtocol.swift       | 5 +++++
 2 files changed, 7 insertions(+)

diff --git a/SignalServiceKit/src/Loki/Protocol/Friend Requests/FriendRequestProtocol.swift b/SignalServiceKit/src/Loki/Protocol/Friend Requests/FriendRequestProtocol.swift
index 3724a1d09..cec1f7dbf 100644
--- a/SignalServiceKit/src/Loki/Protocol/Friend Requests/FriendRequestProtocol.swift	
+++ b/SignalServiceKit/src/Loki/Protocol/Friend Requests/FriendRequestProtocol.swift	
@@ -264,6 +264,8 @@ public final class FriendRequestProtocol : NSObject {
         }
         // The envelope source is set during UD decryption
         let hexEncodedPublicKey = envelope.source!
+        // Don't handle a friend request if we're already handling a session request from someone
+        guard !SessionManagementProtocol.inFlightSessionRequests.contains(hexEncodedPublicKey) else { return }
         // The envelope type is set during UD decryption.
         guard envelope.type == .friendRequest else {
             print("[Loki] Ignoring friend request logic for non friend request type envelope.")
diff --git a/SignalServiceKit/src/Loki/Protocol/Session Management/SessionManagementProtocol.swift b/SignalServiceKit/src/Loki/Protocol/Session Management/SessionManagementProtocol.swift
index 82a4a0318..52dcc8aab 100644
--- a/SignalServiceKit/src/Loki/Protocol/Session Management/SessionManagementProtocol.swift	
+++ b/SignalServiceKit/src/Loki/Protocol/Session Management/SessionManagementProtocol.swift	
@@ -14,6 +14,9 @@ public final class SessionManagementProtocol : NSObject {
 
     internal static var storage: OWSPrimaryStorage { OWSPrimaryStorage.shared() }
 
+    /// The session requests that are currently being handled.
+    public static var inFlightSessionRequests: Set<String> = []
+
     // MARK: - General
     // BEHAVIOR NOTE: OWSMessageSender.throws_encryptedMessageForMessageSend:recipientId:plaintext:transaction: sets
     // isFriendRequest to true if the message in question is a friend request or a device linking request, but NOT if
@@ -192,6 +195,7 @@ public final class SessionManagementProtocol : NSObject {
     public static func handleSessionRequestMessage(_ dataMessage: SSKProtoDataMessage, wrappedIn envelope: SSKProtoEnvelope, using transaction: YapDatabaseReadWriteTransaction) {
         // The envelope source is set during UD decryption
         let hexEncodedPublicKey = envelope.source!
+        inFlightSessionRequests.insert(hexEncodedPublicKey)
         var closedGroupMembers: Set<String> = []
         TSGroupThread.enumerateCollectionObjects(with: transaction) { object, _ in
             guard let group = object as? TSGroupThread, group.groupModel.groupType == .closedGroup,
@@ -208,6 +212,7 @@ public final class SessionManagementProtocol : NSObject {
                 let ephemeralMessage = EphemeralMessage(in: thread)
                 let messageSenderJobQueue = SSKEnvironment.shared.messageSenderJobQueue
                 messageSenderJobQueue.add(message: ephemeralMessage, transaction: transaction)
+                inFlightSessionRequests.remove(hexEncodedPublicKey)
             }
         }
     }