Add transport key test category

Author: ameyapat

IdentityCore/tests/MSIDWorkPlaceJoinUtilTests.m

@@ -808,5 +808,108 @@ - (void)testGetWPJKeysWithTenantId_whenEccRegistrationWithTransportKey_shouldRet
     XCTAssertTrue(result.privateTransportKeyRef != NULL);
 }
 
+- (void)testGetWPJKeysWithTenantId_whenEccRegistrationWithMissingTransportKey_shouldReturnOnlyDeviceKey
+{
+    // Insert device key and cert but skip transport key
+    SecCertificateRef certRef = [self dummyEccCertRef:kDummyTenant1CertIdentifier];
+    NSString *tag = [NSString stringWithFormat:@"%@#%@%@", kMSIDPrivateKeyIdentifier, self.tenantId, @"-EC"];
+    SecKeyRef keyRef = [self createAndGetdummyEccPrivateKey:NO privateKeyTag:tag];
+    NSString *keychainGroup = [self keychainGroup:NO];
+    
+    OSStatus status = [self insertDummyDRSIdentityIntoKeychain:certRef
+                                                 privateKeyRef:keyRef
+                                                 privateKeyTag:tag
+                                                   accessGroup:keychainGroup];
+    XCTAssertTrue(status == errSecSuccess || status == errSecDuplicateItem);
+    
+    // Don't insert transport key - simulate missing STK scenario
+    
+    MSIDWPJKeyPairWithCert *result = [MSIDWorkPlaceJoinUtil getWPJKeysWithTenantId:self.tenantId context:nil];
+    
+    XCTAssertNotNil(result);
+    XCTAssertEqual(result.keyChainVersion, MSIDWPJKeychainAccessGroupV2);
+    XCTAssertTrue(result.privateKeyRef != NULL);
+    XCTAssertTrue(result.privateTransportKeyRef == NULL, @"Expected privateTransportKeyRef to be nil when transport key is missing");
+}
+
+- (void)testGetWPJKeysWithTenantId_whenPrimaryEccRegistrationWithTransportKey_shouldReturnCorrectKeys
+{
+    [self addPrimaryEccDefaultRegistrationForTenantId:self.tenantId
+                                    sharedAccessGroup:[self keychainGroup:NO]
+                                       certIdentifier:kDummyTenant1CertIdentifier
+                                     useSecureEnclave:YES];
+    [self insertEccStkKeyForTenantIdentifier:self.tenantId];
+    
+    MSIDWPJKeyPairWithCert *result = [MSIDWorkPlaceJoinUtil getWPJKeysWithTenantId:nil context:nil];
+    
+    XCTAssertNotNil(result);
+    XCTAssertEqual(result.keyChainVersion, MSIDWPJKeychainAccessGroupV2);
+    XCTAssertTrue(result.privateKeyRef != NULL, @"Primary registration should have device key");
+    XCTAssertTrue(result.privateTransportKeyRef != NULL, @"Primary registration should have transport key");
+}
+
+- (void)testGetWPJKeysWithTenantId_whenLegacyRegistration_shouldHaveNoTransportKey
+{
+    [self insertDummyWPJInLegacyFormat:YES tenantIdentifier:self.tenantId writeTenantMetadata:YES certIdentifier:kDummyTenant1CertIdentifier];
+    
+    MSIDWPJKeyPairWithCert *result = [MSIDWorkPlaceJoinUtil getWPJKeysWithTenantId:self.tenantId context:nil];
+    
+    XCTAssertNotNil(result);
+    XCTAssertEqual(result.keyChainVersion, MSIDWPJKeychainAccessGroupV1);
+    XCTAssertTrue(result.privateKeyRef != NULL, @"Legacy registration should have device key");
+    XCTAssertTrue(result.privateTransportKeyRef == NULL, @"Legacy registration should not have transport key");
+}
+
+- (void)testGetWPJKeysWithTenantId_whenRSARegistrationInV2Format_shouldNotHaveTransportKey
+{
+    // For iOS, RSA keys in V2 format should not have transport keys
+    // This tests the case where we have RSA device key but no transport key expected
+    
+    [self insertDummyWPJInLegacyFormat:NO tenantIdentifier:@"tenantId" writeTenantMetadata:YES certIdentifier:kDummyTenant1CertIdentifier];
+    
+    MSIDWPJKeyPairWithCert *result = [MSIDWorkPlaceJoinUtil getWPJKeysWithTenantId:@"tenantId" context:nil];
+    
+    // For RSA registrations, transport key should be nil even in V2 format
+    XCTAssertNotNil(result);
+    XCTAssertEqual(result.keyChainVersion, MSIDWPJKeychainAccessGroupV2);
+    XCTAssertTrue(result.privateKeyRef != NULL, @"Expected privateKeyRef to be non-nil for RSA registration in V2 format");
+    XCTAssertTrue(result.privateTransportKeyRef == NULL, @"Expected privateTransportKeyRef to be nil for RSA registration in V2 format");
+}
+
+- (void)testGetWPJKeysWithTenantId_concurrentAccess_shouldBeThreadSafe
+{
+    [self insertDummyEccRegistrationForTenantIdentifier:self.tenantId certIdentifier:kDummyTenant1CertIdentifier useSecureEnclave:YES];
+    [self insertEccStkKeyForTenantIdentifier:self.tenantId];
+    dispatch_group_t group = dispatch_group_create();
+    __block NSMutableArray *results = [NSMutableArray array];
+    __block NSLock *lock = [[NSLock alloc] init];
+    
+    // Launch multiple concurrent requests
+    for (int i = 0; i < 5; i++) {
+        dispatch_group_async(group, dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
+            MSIDWPJKeyPairWithCert *result = [MSIDWorkPlaceJoinUtil getWPJKeysWithTenantId:self.tenantId context:nil];
+            
+            [lock lock];
+            if (result) {
+                [results addObject:result];
+            }
+            [lock unlock];
+        });
+    }
+    
+    // Wait for all requests to complete
+    dispatch_group_wait(group, dispatch_time(DISPATCH_TIME_NOW, 5 * NSEC_PER_SEC));
+
+    // All requests should succeed
+    XCTAssertTrue(results.count == 5, @"All concurrent requests should succeed");
+    
+    // Verify all results have transport keys
+    for (MSIDWPJKeyPairWithCert *result in results) {
+        XCTAssertTrue(result.privateKeyRef != NULL);
+        XCTAssertTrue(result.privateTransportKeyRef != NULL);
+    }
+
+}
+
 @end
 #endif