Resolve EXC_BAD_ACCESS in GULNetworkURLSession via O(1) passive memory lifecycle cleanup

GoogleUtilities/Network/GULNetworkURLSession.m

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #import <Foundation/Foundation.h>
+#import <objc/runtime.h>
 
 #import "GoogleUtilities/Network/Public/GoogleUtilities/GULNetworkURLSession.h"
 
@@ -35,6 +36,41 @@ @interface GULNetworkURLSessionWeakHolder : NSObject
 @implementation GULNetworkURLSessionWeakHolder
 @end
 
+@interface GULNetworkURLSession (Private)
++ (NSMutableDictionary<NSString *, GULNetworkURLSessionWeakHolder *> *)sessionIDToFetcherMap;
++ (NSLock *)sessionIDToFetcherMapReadWriteLock;
+@end
+
+@interface GULSessionDeallocTracker : NSObject
+@property(nonatomic, copy) NSString *sessionID;
+@property(nonatomic, strong) GULNetworkURLSessionWeakHolder *holder;
+- (instancetype)initWithSessionID:(NSString *)sessionID
+                           holder:(GULNetworkURLSessionWeakHolder *)holder;
+@end
+
+@implementation GULSessionDeallocTracker
+- (instancetype)initWithSessionID:(NSString *)sessionID
+                           holder:(GULNetworkURLSessionWeakHolder *)holder {
+  self = [super init];
+  if (self) {
+    NSParameterAssert(sessionID);
+    _sessionID = [sessionID copy];
+    _holder = holder;
+  }
+  return self;
+}
+
+- (void)dealloc {
+  [[GULNetworkURLSession sessionIDToFetcherMapReadWriteLock] lock];
+  GULNetworkURLSessionWeakHolder *currentDictionaryHolder =
+      [[GULNetworkURLSession sessionIDToFetcherMap] objectForKey:_sessionID];
+  if (currentDictionaryHolder == _holder) {
+    [[GULNetworkURLSession sessionIDToFetcherMap] removeObjectForKey:_sessionID];
+  }
+  [[GULNetworkURLSession sessionIDToFetcherMapReadWriteLock] unlock];
+}
+@end
+
 @implementation GULNetworkURLSession {
   /// The handler to be called when the request completes or error has occurs.
   GULNetworkURLSessionCompletionHandler _completionHandler;
@@ -678,10 +714,15 @@ - (void)URLSession:(NSURLSession *)session
 
 #pragma mark - Helper Methods
 
+static const void *kGULSessionTrackerKey = &kGULSessionTrackerKey;
+
 + (void)setSessionInFetcherMap:(GULNetworkURLSession *)session forSessionID:(NSString *)sessionID {
   if (!sessionID) {
     return;
   }
+
+  GULSessionDeallocTracker *oldTrackerToReleaseOutsideLock = nil;
+
   [[self sessionIDToFetcherMapReadWriteLock] lock];
   GULNetworkURLSessionWeakHolder *holder =
       [[[self class] sessionIDToFetcherMap] objectForKey:sessionID];
@@ -693,27 +734,29 @@ + (void)setSessionInFetcherMap:(GULNetworkURLSession *)session forSessionID:(NSS
                                                     messageCode:kGULNetworkMessageCodeURLSession019
                                                         message:message];
     }
+    oldTrackerToReleaseOutsideLock =
+        objc_getAssociatedObject(existingSession, kGULSessionTrackerKey);
+    objc_setAssociatedObject(existingSession, kGULSessionTrackerKey, nil,
+                             OBJC_ASSOCIATION_RETAIN_NONATOMIC);
     [existingSession->_URLSession finishTasksAndInvalidate];
   }
   if (session) {
-    // Cleanup nil entries.
-    NSMutableArray *keysToRemove = [[NSMutableArray alloc] init];
-    [[[self class] sessionIDToFetcherMap]
-        enumerateKeysAndObjectsUsingBlock:^(NSString *key, GULNetworkURLSessionWeakHolder *holder,
-                                            BOOL *stop) {
-          if (holder.session == nil) {
-            [keysToRemove addObject:key];
-          }
-        }];
-    [[[self class] sessionIDToFetcherMap] removeObjectsForKeys:keysToRemove];
-
     GULNetworkURLSessionWeakHolder *newHolder = [[GULNetworkURLSessionWeakHolder alloc] init];
     newHolder.session = session;
     [[[self class] sessionIDToFetcherMap] setObject:newHolder forKey:sessionID];
+
+    GULSessionDeallocTracker *tracker =
+        [[GULSessionDeallocTracker alloc] initWithSessionID:sessionID holder:newHolder];
+    objc_setAssociatedObject(session, kGULSessionTrackerKey, tracker,
+                             OBJC_ASSOCIATION_RETAIN_NONATOMIC);
   } else {
     [[[self class] sessionIDToFetcherMap] removeObjectForKey:sessionID];
   }
   [[self sessionIDToFetcherMapReadWriteLock] unlock];
+
+  // Retain the old tracker until the lock is released. If it deallocates inside the lock,
+  // its -dealloc method will attempt to re-acquire the lock, resulting in a deadlock.
+  (void)oldTrackerToReleaseOutsideLock;
 }
 
 + (nullable GULNetworkURLSession *)sessionFromFetcherMapForSessionID:(NSString *)sessionID {

GoogleUtilities/SwizzlerTestHelpers/GULRuntimeSnapshot.m

@@ -76,11 +76,13 @@ - (void)capture {
     for (int i = 0; i < numberOfClasses; i++) {
       Class aClass = classList[i];
       NSString *classString = NSStringFromClass(aClass);
-      GULRuntimeClassSnapshot *classSnapshot =
-          [[GULRuntimeClassSnapshot alloc] initWithClass:aClass];
-      _classSnapshots[classString] = classSnapshot;
-      [classSnapshot capture];
-      _runningHash ^= [classSnapshot hash];
+      if (classString) {
+        GULRuntimeClassSnapshot *classSnapshot =
+            [[GULRuntimeClassSnapshot alloc] initWithClass:aClass];
+        _classSnapshots[classString] = classSnapshot;
+        [classSnapshot capture];
+        _runningHash ^= [classSnapshot hash];
+      }
     }
   }
   free(classList);

GoogleUtilities/Tests/Unit/Environment/GULKeychainStorageTests.m

@@ -133,7 +133,7 @@ - (void)testGetExistingObjectClassMismatch {
                 OCMVerifyAll(self.mockCache);
                 [expectation fulfill];
               }];
-  [self waitForExpectations:@[ expectation ] timeout:1.0];
+  [self waitForExpectations:@[ expectation ] timeout:5.0];
 }
 
 - (void)testRemoveExistingObject {
@@ -174,7 +174,7 @@ - (void)assertSuccessWriteObject:(id<NSSecureCoding>)object forKey:(NSString *)k
           [expectation fulfill];
         }];
 
-  [self waitForExpectations:@[ expectation ] timeout:1.0];
+  [self waitForExpectations:@[ expectation ] timeout:5.0];
   OCMVerifyAll(self.mockCache);
 }
 
@@ -204,7 +204,7 @@ - (void)assertSuccessReadObject:(id<NSSecureCoding>)object
         XCTAssertEqualObjects([weakSelf.cache objectForKey:key], object, @"%@", weakSelf.name);
         [expectation fulfill];
       }];
-  [self waitForExpectations:@[ expectation ] timeout:1.0];
+  [self waitForExpectations:@[ expectation ] timeout:5.0];
   OCMVerifyAll(self.mockCache);
 }
 
@@ -224,7 +224,7 @@ - (void)assertNonExistingObjectForKey:(NSString *)key class:(Class)class {
                 XCTAssertNil(obj, @"%@", weakSelf.name);
                 [expectation fulfill];
               }];
-  [self waitForExpectations:@[ expectation ] timeout:1.0];
+  [self waitForExpectations:@[ expectation ] timeout:5.0];
   OCMVerifyAll(self.mockCache);
 }
 
@@ -238,7 +238,7 @@ - (void)assertRemoveObjectForKey:(NSString *)key {
                    XCTAssertNil(error);
                    [expectation fulfill];
                  }];
-  [self waitForExpectations:@[ expectation ] timeout:1.0];
+  [self waitForExpectations:@[ expectation ] timeout:5.0];
   OCMVerifyAll(self.mockCache);
 }
 

GoogleUtilities/Tests/Unit/Network/GULMutableDictionaryTest.m

@@ -15,12 +15,6 @@
 #import <XCTest/XCTest.h>
 
 #import "GoogleUtilities/Network/Public/GoogleUtilities/GULMutableDictionary.h"
-#import "GoogleUtilities/Network/Public/GoogleUtilities/GULNetworkURLSession.h"
-
-@interface GULNetworkURLSession (Test)
-+ (void)setSessionInFetcherMap:(GULNetworkURLSession *)session forSessionID:(NSString *)sessionID;
-+ (nullable GULNetworkURLSession *)sessionFromFetcherMapForSessionID:(NSString *)sessionID;
-@end
 
 const static NSString *const kKey = @"testKey1";
 const static NSString *const kValue = @"testValue1";
@@ -90,46 +84,4 @@ - (void)testUnderlyingDictionary {
   XCTAssertEqual(dict[kKey2], kValue2);
 }
 
-- (void)testSessionMapStress {
-  int iterations = 1000;
-  int numSessionIDs = 10;
-
-  XCTestExpectation *expectation = [self expectationWithDescription:@"Stress test finished"];
-
-  dispatch_queue_t queue =
-      dispatch_queue_create("com.google.utilities.stress", DISPATCH_QUEUE_CONCURRENT);
-
-  dispatch_group_t group = dispatch_group_create();
-
-  // Reader threads
-  for (int t = 0; t < 5; t++) {
-    dispatch_group_async(group, queue, ^{
-      for (int i = 0; i < iterations; i++) {
-        int id_idx = i % numSessionIDs;
-        NSString *sessionID = [NSString stringWithFormat:@"session_%d", id_idx];
-        [GULNetworkURLSession sessionFromFetcherMapForSessionID:sessionID];
-      }
-    });
-  }
-
-  // Writer threads
-  for (int t = 0; t < 5; t++) {
-    dispatch_group_async(group, queue, ^{
-      for (int i = 0; i < iterations; i++) {
-        int id_idx = i % numSessionIDs;
-        NSString *sessionID = [NSString stringWithFormat:@"session_%d", id_idx];
-        GULNetworkURLSession *session =
-            [[GULNetworkURLSession alloc] initWithNetworkLoggerDelegate:nil];
-        [GULNetworkURLSession setSessionInFetcherMap:session forSessionID:sessionID];
-      }
-    });
-  }
-
-  dispatch_group_notify(group, dispatch_get_main_queue(), ^{
-    [expectation fulfill];
-  });
-
-  [self waitForExpectationsWithTimeout:30 handler:nil];
-}
-
 @end

GoogleUtilities/Tests/Unit/Network/GULNetworkTest.m

@@ -122,6 +122,36 @@ - (void)testReachability {
   reachabilityMock = nil;
 }
 
+#pragma mark - Test Passive Deallocation
+
+- (void)testSessionPassiveRemovalOnDeallocation {
+  NSString *testSessionID = @"test_passive_removal_id";
+
+  @autoreleasepool {
+    // 1. Create a session inside a limited scope.
+    GULNetworkURLSession *session =
+        [[GULNetworkURLSession alloc] initWithNetworkLoggerDelegate:nil];
+
+    // 2. Insert it into the fetcher map (which stores a WeakHolder and sets an Associated Object).
+    [GULNetworkURLSession setSessionInFetcherMap:session forSessionID:testSessionID];
+
+    // 3. Verify it is accessible from the map.
+    GULNetworkURLSession *retrievedSession =
+        [GULNetworkURLSession sessionFromFetcherMapForSessionID:testSessionID];
+    XCTAssertEqualObjects(session, retrievedSession, @"Session should be in the fetcher map.");
+  }
+  // 4. Exiting the autoreleasepool destroys the 'session' local variable.
+  //    ARC calls -dealloc on the session, which triggers the destruction of its
+  //    Associated Objects (including the GULSessionDeallocTracker), which in turn
+  //    cleans up the sessionID from the fetcher map.
+
+  // 5. Verify the session was passively removed.
+  GULNetworkURLSession *retrievedAfterDealloc =
+      [GULNetworkURLSession sessionFromFetcherMapForSessionID:testSessionID];
+  XCTAssertNil(retrievedAfterDealloc,
+               @"Session should be automatically removed from fetcher map upon deallocation.");
+}
+
 #pragma mark - Test POST Foreground
 
 - (void)testSessionNetwork_POST_foreground {