Replace setter override in category with initializer without category

Author: ameyapat

IdentityCore/IdentityCore.xcodeproj/project.pbxproj

@@ -752,9 +752,6 @@
 		728209D026FEA0F600B5F018 /* MSIDKeyOperationUtil.m in Sources */ = {isa = PBXBuildFile; fileRef = 728209CF26FEA0F600B5F018 /* MSIDKeyOperationUtil.m */; };
 		728209D126FEA0F600B5F018 /* MSIDKeyOperationUtil.m in Sources */ = {isa = PBXBuildFile; fileRef = 728209CF26FEA0F600B5F018 /* MSIDKeyOperationUtil.m */; };
 		728209D62702AF8900B5F018 /* MSIDBackgroundTaskManagerTests.m in Sources */ = {isa = PBXBuildFile; fileRef = 728209D32702AE9300B5F018 /* MSIDBackgroundTaskManagerTests.m */; };
-		728ABACC2E5A3B4E00FCE434 /* MSIDWPJKeyPairWithCert+TransportKey.h in Headers */ = {isa = PBXBuildFile; fileRef = 728ABACB2E5A3B2800FCE434 /* MSIDWPJKeyPairWithCert+TransportKey.h */; };
-		728ABACE2E5A41A800FCE434 /* MSIDWPJKeyPairWithCert+TransportKey.m in Sources */ = {isa = PBXBuildFile; fileRef = 728ABACD2E5A418F00FCE434 /* MSIDWPJKeyPairWithCert+TransportKey.m */; };
-		728ABACF2E5A41A800FCE434 /* MSIDWPJKeyPairWithCert+TransportKey.m in Sources */ = {isa = PBXBuildFile; fileRef = 728ABACD2E5A418F00FCE434 /* MSIDWPJKeyPairWithCert+TransportKey.m */; };
 		728D9E4628245DD7001D990F /* MSIDTestSecureEnclaveKeyPairGenerator.m in Sources */ = {isa = PBXBuildFile; fileRef = 728D9E4528245DD7001D990F /* MSIDTestSecureEnclaveKeyPairGenerator.m */; };
 		728D9E4728245DD7001D990F /* MSIDTestSecureEnclaveKeyPairGenerator.m in Sources */ = {isa = PBXBuildFile; fileRef = 728D9E4528245DD7001D990F /* MSIDTestSecureEnclaveKeyPairGenerator.m */; };
 		728D9E492824A323001D990F /* MSIDPkeyAuthHelperTests.m in Sources */ = {isa = PBXBuildFile; fileRef = 23CA0C5E220A68D400768729 /* MSIDPkeyAuthHelperTests.m */; };
@@ -2683,8 +2680,6 @@
 		728209CD26FEA0D800B5F018 /* MSIDKeyOperationUtil.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = MSIDKeyOperationUtil.h; sourceTree = "<group>"; };
 		728209CF26FEA0F600B5F018 /* MSIDKeyOperationUtil.m */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.objc; path = MSIDKeyOperationUtil.m; sourceTree = "<group>"; };
 		728209D32702AE9300B5F018 /* MSIDBackgroundTaskManagerTests.m */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.objc; path = MSIDBackgroundTaskManagerTests.m; sourceTree = "<group>"; };
-		728ABACB2E5A3B2800FCE434 /* MSIDWPJKeyPairWithCert+TransportKey.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = "MSIDWPJKeyPairWithCert+TransportKey.h"; sourceTree = "<group>"; };
-		728ABACD2E5A418F00FCE434 /* MSIDWPJKeyPairWithCert+TransportKey.m */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.objc; path = "MSIDWPJKeyPairWithCert+TransportKey.m"; sourceTree = "<group>"; };
 		728D9E4528245DD7001D990F /* MSIDTestSecureEnclaveKeyPairGenerator.m */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.objc; path = MSIDTestSecureEnclaveKeyPairGenerator.m; sourceTree = "<group>"; };
 		728D9E4828247D4C001D990F /* MSIDTestSecureEnclaveKeyPairGenerator.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = MSIDTestSecureEnclaveKeyPairGenerator.h; sourceTree = "<group>"; };
 		729357E72DD810C70001D03C /* MSIDNonceTokenRequest.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = MSIDNonceTokenRequest.h; sourceTree = "<group>"; };
@@ -5213,8 +5208,6 @@
 		B2C0747E246B70DC0008D701 /* crypto */ = {
 			isa = PBXGroup;
 			children = (
-				728ABACD2E5A418F00FCE434 /* MSIDWPJKeyPairWithCert+TransportKey.m */,
-				728ABACB2E5A3B2800FCE434 /* MSIDWPJKeyPairWithCert+TransportKey.h */,
 				B27893792470CAF200627C28 /* mac */,
 				B2C0748E246B71470008D701 /* MSIDAssymetricKeyGenerating.h */,
 				B2C07490246B735B0008D701 /* MSIDAssymetricKeyKeychainGenerator.h */,
@@ -6182,7 +6175,6 @@
 				A07EB427259D0C6B00783943 /* MSIDThrottlingService.h in Headers */,
 				9658103120C7E1180025F4A4 /* MSIDWebviewResponse.h in Headers */,
 				1E707FDF2407335700716148 /* MSIDBrokerNativeAppOperationResponse.h in Headers */,
-				728ABACC2E5A3B4E00FCE434 /* MSIDWPJKeyPairWithCert+TransportKey.h in Headers */,
 				B28BDA7F217E964B003E5670 /* MSIDB2CTokenResponse.h in Headers */,
 				96B8D57D20946D2600E3F4A6 /* MSIDPkce.h in Headers */,
 				B286B9912389DC47007833AD /* MSIDIndividualClaimRequest.h in Headers */,
@@ -7438,7 +7430,6 @@
 				23C8981A2C892A3800071482 /* MSIDBrowserNativeMessageGetSupportedContractsResponse.m in Sources */,
 				B286B9992389DC9D007833AD /* MSIDSSOExtensionSilentTokenRequest.m in Sources */,
 				B2C7089921991D0000D917B8 /* MSIDAADV2BrokerResponse.m in Sources */,
-				728ABACE2E5A41A800FCE434 /* MSIDWPJKeyPairWithCert+TransportKey.m in Sources */,
 				B20E3CB61FC4FE400029C097 /* MSIDOAuth2Constants.m in Sources */,
 				B2BE924D21A2331A00F5AB8C /* MSIDTelemetryAuthorityValidationEvent.m in Sources */,
 				B2807FF9204CAFDF00944D89 /* MSIDHelpers.m in Sources */,
@@ -8151,7 +8142,6 @@
 				239E3BBF23E1004F00F7A50A /* MSIDClientSDKType.m in Sources */,
 				23B39ABD209BD47D000AA905 /* MSIDB2CAuthorityResolver.m in Sources */,
 				B2F671E92467A34400649855 /* MSIDAuthorizationCodeResult.m in Sources */,
-				728ABACF2E5A41A800FCE434 /* MSIDWPJKeyPairWithCert+TransportKey.m in Sources */,
 				23C10A9F2B40D9350063D97C /* MSIDBrowserNativeMessageSignOutResponse.m in Sources */,
 				23FB5C2B225517AA002BF1EB /* MSIDIndividualClaimRequestAdditionalInfo.m in Sources */,
 				1E707FDD2406FA9200716148 /* MSIDBrokerBrowserOperationResponse.m in Sources */,

IdentityCore/src/MSIDConstants.h

@@ -230,4 +230,6 @@ extern NSString * _Nonnull const MSID_FLIGHT_IGNORE_COOKIES_IN_DUNA_RESUME;
  */
 extern NSString * _Nonnull const MSID_FLIGHT_DISABLE_REMOVE_ACCOUNT_ARTIFACTS;
 
+extern NSString * _Nonnull const MSID_FLIGHT_DISABLE_QUERYING_STK;
+
 #define METHODANDLINE   [NSString stringWithFormat:@"%s [Line %d]", __PRETTY_FUNCTION__, __LINE__]

IdentityCore/src/MSIDConstants.m

@@ -94,5 +94,7 @@
 // Making the flight string short to avoid legacy broker url size limit
 NSString *const MSID_FLIGHT_DISABLE_REMOVE_ACCOUNT_ARTIFACTS = @"disable_rm_metadata";
 
+NSString *const MSID_FLIGHT_DISABLE_QUERYING_STK = @"disable_querying_stk";
+
 
 #define METHODANDLINE   [NSString stringWithFormat:@"%s [Line %d]", __PRETTY_FUNCTION__, __LINE__]

IdentityCore/src/cache/crypto/MSIDWPJKeyPairWithCert+TransportKey.h

@@ -63,6 +63,7 @@ typedef NS_ENUM(NSInteger, MSIDWPJKeychainAccessGroup)
                                 certificate:(SecCertificateRef)certificate
                           certificateIssuer:(nullable NSString *)issuer;
 
+- (nullable instancetype)initializePrivateTransportKeyRef:(SecKeyRef)privateTransportKeyRef;
 @end
 
 NS_ASSUME_NONNULL_END

IdentityCore/src/cache/crypto/MSIDWPJKeyPairWithCert+TransportKey.m

@@ -33,6 +33,7 @@ @interface MSIDWPJKeyPairWithCert()
 @property (nonatomic) NSString *certificateSubject;
 @property (nonatomic) NSString *certificateIssuer;
 @property (nonatomic) SecKeyRef privateKeyRef;
+@property (nonatomic) SecKeyRef privateTransportKeyRef;
 
 @end
 
@@ -88,6 +89,30 @@ - (nullable instancetype)initWithPrivateKey:(SecKeyRef)privateKey
     return self;
 }
 
+- (nullable instancetype)initializePrivateTransportKeyRef:(nonnull SecKeyRef)privateTransportKeyRef
+{
+    if (self && privateTransportKeyRef)
+    {
+        if (_privateTransportKeyRef != privateTransportKeyRef)
+        {
+            if (_privateTransportKeyRef)
+            {
+                CFReleaseNull(_privateTransportKeyRef);
+                _privateTransportKeyRef = NULL;
+            }
+            
+            _privateTransportKeyRef = privateTransportKeyRef;
+            
+            if (_privateTransportKeyRef)
+            {
+                CFRetain(_privateTransportKeyRef);
+            }
+        }
+    }
+    return self;
+}
+
+
 - (void)dealloc
 {
     if (_certificateRef)
@@ -104,7 +129,7 @@ - (void)dealloc
 
     if (_privateTransportKeyRef)
     {
-        CFRelease(_privateTransportKeyRef);
+        CFReleaseNull(_privateTransportKeyRef);
         _privateTransportKeyRef = NULL;
     }
 }

IdentityCore/src/cache/crypto/MSIDWPJKeyPairWithCert.h

@@ -27,7 +27,8 @@
 #import "MSIDWorkPlaceJoinConstants.h"
 #import "MSIDWPJKeyPairWithCert.h"
 #import "MSIDWPJMetadata.h"
-#import "MSIDWPJKeyPairWithCert+TransportKey.h"
+#import "MSIDFlightManager.h"
+#import "MSIDConstants.h"
 
 static NSString *kWPJPrivateKeyIdentifier = @"com.microsoft.workplacejoin.privatekey\0";
 static NSString *kECPrivateKeyTagSuffix = @"-EC";
@@ -381,13 +382,16 @@ + (MSIDWPJKeyPairWithCert *)getWPJKeysWithTenantId:(__unused NSString *)tenantId
     {
         defaultKeys.keyChainVersion = MSIDWPJKeychainAccessGroupV2;
         MSID_LOG_WITH_CTX(MSIDLogLevelInfo, context, @"Returning EC private device key from default registration.");
+#if TARGET_OS_IPHONE
+        bool isQueryingDisabledViaFlight = [MSIDFlightManager.sharedInstance boolForKey:MSID_FLIGHT_DISABLE_QUERYING_STK];
         // Query the session transport key only for iOS.
         // 1P apps use transport key to decrypt ECDH JWE responses when redeeming bound regular refresh tokens
         id keyType = privateKeyAttributes[(__bridge id)kSecAttrKeyType];
-        if (keyType && [keyType isEqual: (__bridge id)kSecAttrKeyTypeECSECPrimeRandom])
+        if (!isQueryingDisabledViaFlight && keyType && [keyType isEqual: (__bridge id)kSecAttrKeyTypeECSECPrimeRandom])
         {
-            defaultKeys.privateTransportKeyRef = [self getSessionTransportKeyRefFromSecureEnclaveForTenantId:tenantId context:context];
+            [defaultKeys initializePrivateTransportKeyRef:[self getSessionTransportKeyRefFromSecureEnclaveForTenantId:tenantId context:context]];
         }
+#endif
         return defaultKeys;
     }