Improve SPKI cache thread safety and handle protected data availability

datatheorem · Oct 21, 2024 · eb973a7 · eb973a7
1 parent d009bab
commit eb973a7
Show file tree

Hide file tree

Showing 4 changed files with 67 additions and 9 deletions.
diff --git a/TrustKit/Pinning/TSKSPKIHashCache.m b/TrustKit/Pinning/TSKSPKIHashCache.m
@@ -245,13 +245,27 @@ - (NSData *)hashSubjectPublicKeyInfoFromCertificate:(SecCertificateRef)certifica
     });
 
     // Update the cache on the filesystem
-    if (self.spkiCacheFilename.length > 0 && [[UIApplication sharedApplication] isProtectedDataAvailable])
-    {
-        NSData *serializedSpkiCache = [NSKeyedArchiver archivedDataWithRootObject:_spkiCache requiringSecureCoding:YES error:nil];
-        if ([serializedSpkiCache writeToURL:[self SPKICachePath] atomically:YES] == NO)
-        {
-            NSAssert(false, @"Failed to write cache");
-            TSKLog(@"Could not persist SPKI cache to the filesystem");
+    if (self.spkiCacheFilename.length > 0) {
+
+        __weak typeof(self) weakSelf = self;
+        void (^updateCacheBlock)(void) = ^{
+            if ([[UIApplication sharedApplication] isProtectedDataAvailable]) {
+                NSData *serializedSpkiCache = [NSKeyedArchiver archivedDataWithRootObject:weakSelf.spkiCache requiringSecureCoding:YES error:nil];
+                if ([serializedSpkiCache writeToURL:[weakSelf SPKICachePath] atomically:YES] == NO) {
+                    NSAssert(false, @"Failed to write cache");
+                    TSKLog(@"Could not persist SPKI cache to the filesystem");
+                }
+            }
+            else {
+                TSKLog(@"Protected data not available, skipping SPKI cache persistence");
+            }
+        };
+
+        if ([NSThread isMainThread]) {
+            updateCacheBlock();
+        }
+        else {
+            dispatch_async(dispatch_get_main_queue(), updateCacheBlock);
         }
     }
 

diff --git a/TrustKitTests/TSKEndToEndNSURLSessionTests.m b/TrustKitTests/TSKEndToEndNSURLSessionTests.m
@@ -197,7 +197,7 @@ - (void)testPinningValidationSucceeded
           @{
               @"www.datatheorem.com" : @{
                       kTSKEnforcePinning : @YES,
-                      kTSKPublicKeyHashes : @[@"cXjPgKdVe6iojP8s0YQJ3rtmDFHTnYZxcYvmYGFiYME=", // CA key for Google Trust Services (cert valid until 27 Jan 2028)
+                      kTSKPublicKeyHashes : @[@"hxqRlPTu1bMS/0DITB1SSu0vd4u/8l8TjPgfaAp63Gc=", // CA key for Google Trust Services (cert valid until 27 Jan 2028)
                                               @"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=" // Fake key
                                               ]}}};
 

diff --git a/TrustKitTests/TSKEndToEndSwizzlingTests.m b/TrustKitTests/TSKEndToEndSwizzlingTests.m
@@ -167,7 +167,7 @@ - (void)test
               // Valid pinning configuration
               @"www.datatheorem.com" : @{
                       kTSKEnforcePinning : @YES,
-                      kTSKPublicKeyHashes : @[@"F6jTih9VkkYZS8yuYqeU/4DUGehJ+niBGkkQ1yg8H3U=", // CA key
+                      kTSKPublicKeyHashes : @[@"a4FoHyEnsFhauIx0w/gB7ywslD6tWGk83J2F6Pv1phA=", // CA key
                                               @"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=" // Fake key
                                               ]},
               // Invalid pinning configuration

diff --git a/TrustKitTests/TSKPublicKeyAlgorithmTests.m b/TrustKitTests/TSKPublicKeyAlgorithmTests.m
@@ -18,6 +18,7 @@
 #import "../TrustKit/Reporting/reporting_utils.h"
 
 #import "TSKCertificateUtils.h"
+#import <OCMock/OCMock.h>
 
 
 @interface TSKSPKIHashCache (TestSupport)
@@ -112,4 +113,47 @@ - (void)testExtractEcDsaSecp384r1
 }
 
 
+- (void)testSPKICacheThreadSafetyAndProtectedData
+{
+    XCTestExpectation *expectation = [self expectationWithDescription:@"Cache operations completed"];
+
+    id mockApplication = OCMClassMock([UIApplication class]);
+    OCMStub([mockApplication sharedApplication]).andReturn(mockApplication);
+    // Simulate protected data being unavailable
+    OCMStub([mockApplication isProtectedDataAvailable]).andReturn(NO);
+
+
+    // Perform multiple cache operations in parallel on a background queue
+    SecCertificateRef certificate = [TSKCertificateUtils createCertificateFromDer:@"www.globalsign.com"];
+    dispatch_group_t group = dispatch_group_create();
+    for (int i = 0; i < 10; i++) {
+        dispatch_group_async(group, dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
+            [self->spkiCache hashSubjectPublicKeyInfoFromCertificate:certificate];
+        });
+    }
+
+    dispatch_group_notify(group, dispatch_get_main_queue(), ^{
+
+        NSDictionary *cache = [self->spkiCache getSubjectPublicKeyInfoHashesCache];
+        XCTAssertEqual(cache.count, 1, @"Cache should contain one entry");
+
+        // Simulate protected data becoming available
+        OCMStub([mockApplication isProtectedDataAvailable]).andReturn(YES);
+
+        // Perform one more cache operation to trigger filesystem write
+        [self->spkiCache hashSubjectPublicKeyInfoFromCertificate:certificate];
+
+        dispatch_after(dispatch_time(DISPATCH_TIME_NOW, (int64_t)(1 * NSEC_PER_SEC)), dispatch_get_main_queue(), ^{
+            NSDictionary *finalCache = [self->spkiCache getSubjectPublicKeyInfoHashesCache];
+            XCTAssertEqual(finalCache.count, 1, @"Cache should still contain one entry");
+
+            CFRelease(certificate);
+            [mockApplication stopMocking];
+            [expectation fulfill];
+        });
+    });
+
+    [self waitForExpectationsWithTimeout:5.0 handler:nil];
+}
+
 @end