From 1dea38cbedc17c41c89f751636b0c805060a38da Mon Sep 17 00:00:00 2001
From: Morgan Chen <morganchen12@gmail.com>
Date: Fri, 19 Oct 2018 12:07:53 -0700
Subject: [PATCH] Lock reads and writes to session map

---
 .../Network/GULNetworkURLSession.m            | 45 +++++++++++++------
 1 file changed, 31 insertions(+), 14 deletions(-)

diff --git a/GoogleUtilities/Network/GULNetworkURLSession.m b/GoogleUtilities/Network/GULNetworkURLSession.m
index 7fe134cdc65..26da579a3bb 100644
--- a/GoogleUtilities/Network/GULNetworkURLSession.m
+++ b/GoogleUtilities/Network/GULNetworkURLSession.m
@@ -158,7 +158,7 @@ - (NSString *)sessionIDFromAsyncPOSTRequest:(NSURLRequest *)request
   }
 
   // Save the session into memory.
-  [self setSessionInFetcherMap:self forSessionID:_sessionID];
+  [[self class] setSessionInFetcherMap:self forSessionID:_sessionID];
 
   _request = [request copy];
 
@@ -200,7 +200,7 @@ - (NSString *)sessionIDFromAsyncGETRequest:(NSURLRequest *)request
   }
 
   // Save the session into memory.
-  [self setSessionInFetcherMap:self forSessionID:_sessionID];
+  [[self class] setSessionInFetcherMap:self forSessionID:_sessionID];
 
   _request = [request copy];
 
@@ -291,7 +291,7 @@ - (void)URLSession:(NSURLSession *)session
     // Explicitly remove the session so it won't be reused. The weak map table should
     // remove the session on deallocation, but dealloc may not happen immediately after
     // calling `finishTasksAndInvalidate`.
-    [self setSessionInFetcherMap:nil forSessionID:sessionID];
+    [[self class] setSessionInFetcherMap:nil forSessionID:sessionID];
   }
 }
 
@@ -540,18 +540,19 @@ - (void)removeTempItemAtURL:(NSURL *)fileURL {
 
 /// Gets the fetcher with the session ID.
 + (instancetype)fetcherWithSessionIdentifier:(NSString *)sessionIdentifier {
-  NSMapTable<NSString *, GULNetworkURLSession *> *sessionIdentifierToFetcherMap =
-      [self sessionIDToFetcherMap];
-  GULNetworkURLSession *session = [sessionIdentifierToFetcherMap objectForKey:sessionIdentifier];
+  GULNetworkURLSession *session = [self sessionFromFetcherMapForSessionID:sessionIdentifier];
   if (!session && [sessionIdentifier hasPrefix:kGULNetworkBackgroundSessionConfigIDPrefix]) {
     session = [[GULNetworkURLSession alloc] initWithNetworkLoggerDelegate:nil];
     [session setSessionID:sessionIdentifier];
-    [sessionIdentifierToFetcherMap setObject:session forKey:sessionIdentifier];
+    [self setSessionInFetcherMap:session forSessionID:sessionIdentifier];
   }
   return session;
 }
 
 /// Returns a map of the fetcher by session ID. Creates a map if it is not created.
+/// When reading and writing from/to the session map, don't use this method directly.
+/// To avoid thread safety issues, use one of the helper methods at the bottom of the
+/// file: setSessionInFetcherMap:forSessionID:, sessionFromFetcherMapForSessionID:
 + (NSMapTable<NSString *, GULNetworkURLSession *> *)sessionIDToFetcherMap {
   static NSMapTable *sessionIDToFetcherMap;
 
@@ -562,6 +563,16 @@ + (instancetype)fetcherWithSessionIdentifier:(NSString *)sessionIdentifier {
   return sessionIDToFetcherMap;
 }
 
++ (NSLock *)sessionIDToFetcherMapReadWriteLock {
+  static NSLock *lock;
+
+  static dispatch_once_t onceToken;
+  dispatch_once(&onceToken, ^{
+    lock = [[NSLock alloc] init];
+  });
+  return lock;
+}
+
 /// Returns a map of system provided completion handler by session ID. Creates a map if it is not
 /// created.
 + (GULMutableDictionary *)sessionIDToSystemCompletionHandlerDictionary {
@@ -661,26 +672,32 @@ - (void)URLSession:(NSURLSession *)session
 
 #pragma mark - Helper Methods
 
-- (void)setSessionInFetcherMap:(GULNetworkURLSession *)session forSessionID:(NSString *)sessionID {
-  GULNetworkURLSession *existingSession = [self sessionFromFetcherMapForSessionID:sessionID];
++ (void)setSessionInFetcherMap:(GULNetworkURLSession *)session forSessionID:(NSString *)sessionID {
+  [[self sessionIDToFetcherMapReadWriteLock] lock];
+  GULNetworkURLSession *existingSession =
+      [[[self class] sessionIDToFetcherMap] objectForKey:sessionID];
   if (existingSession) {
     // Invalidating doesn't seem like the right thing to do here since it may cancel an active
     // background transfer if the background session is handling multiple requests. The old
     // session will be dropped from the map table, but still complete its request.
     NSString *message = [NSString stringWithFormat:@"Discarding session: %@", existingSession];
-    [_loggerDelegate GULNetwork_logWithLevel:kGULNetworkLogLevelInfo
-                                 messageCode:kGULNetworkMessageCodeURLSession019
-                                     message:message];
+    [existingSession->_loggerDelegate GULNetwork_logWithLevel:kGULNetworkLogLevelInfo
+                                                  messageCode:kGULNetworkMessageCodeURLSession019
+                                                      message:message];
   }
   if (session) {
     [[[self class] sessionIDToFetcherMap] setObject:session forKey:sessionID];
   } else {
     [[[self class] sessionIDToFetcherMap] removeObjectForKey:sessionID];
   }
+  [[self sessionIDToFetcherMapReadWriteLock] unlock];
 }
 
-- (nullable GULNetworkURLSession *)sessionFromFetcherMapForSessionID:(NSString *)sessionID {
-  return [[[self class] sessionIDToFetcherMap] objectForKey:sessionID];
++ (nullable GULNetworkURLSession *)sessionFromFetcherMapForSessionID:(NSString *)sessionID {
+  [[self sessionIDToFetcherMapReadWriteLock] lock];
+  GULNetworkURLSession *session = [[[self class] sessionIDToFetcherMap] objectForKey:sessionID];
+  [[self sessionIDToFetcherMapReadWriteLock] unlock];
+  return session;
 }
 
 - (void)callCompletionHandler:(GULNetworkURLSessionCompletionHandler)handler