Add DCErrorInvalidKey handling test case

firebase · Oct 24, 2023 · 0f4f5b6 · 0f4f5b6
1 parent dbde712
commit 0f4f5b6
Show file tree

Hide file tree

Showing 2 changed files with 91 additions and 19 deletions.
diff --git a/FirebaseAppCheck/Sources/AppAttestProvider/FIRAppAttestProvider.m b/FirebaseAppCheck/Sources/AppAttestProvider/FIRAppAttestProvider.m
@@ -322,25 +322,27 @@ - (void)getTokenWithCompletion:(void (^)(FIRAppCheckToken *_Nullable, NSError *_
 
                 return [self attestKey:keyID challenge:challenge];
               })
-      .recoverOn(
-          self.queue,
-          ^id(NSError *error) {
-            // If Apple rejected the key then reset the attestation and throw a specific error.
-            if ([error.domain isEqualToString:DCErrorDomain] && error.code == DCErrorInvalidKey) {
-              FIRAppCheckDebugLog(
-                  kFIRLoggerAppCheckMessageCodeAttestationRejected,
-                  @"App Attest invalid key; the existing attestation will be reset.");
-
-              // Reset the attestation.
-              return [self resetAttestation].thenOn(self.queue, ^NSError *(id result) {
-                // Throw the rejection error.
-                return [[FIRAppAttestRejectionError alloc] init];
-              });
-            }
-
-            // Otherwise just re-throw the error.
-            return error;
-          })
+      .recoverOn(self.queue,
+                 ^id(NSError *error) {
+                   // If Apple rejected the key (DCErrorInvalidKey) then reset the attestation and
+                   // throw a specific error to signal retry (FIRAppAttestRejectionError).
+                   NSError *underlyingError = error.userInfo[NSUnderlyingErrorKey];
+                   if (underlyingError && [underlyingError.domain isEqualToString:DCErrorDomain] &&
+                       underlyingError.code == DCErrorInvalidKey) {
+                     FIRAppCheckDebugLog(
+                         kFIRLoggerAppCheckMessageCodeAttestationRejected,
+                         @"App Attest invalid key; the existing attestation will be reset.");
+
+                     // Reset the attestation.
+                     return [self resetAttestation].thenOn(self.queue, ^NSError *(id result) {
+                       // Throw the rejection error.
+                       return [[FIRAppAttestRejectionError alloc] init];
+                     });
+                   }
+
+                   // Otherwise just re-throw the error.
+                   return error;
+                 })
       .thenOn(self.queue,
               ^FBLPromise<NSArray *> *(FIRAppAttestKeyAttestationResult *result) {
                 // 3. Exchange the attestation to FAC token and pass the results to the next step.

diff --git a/FirebaseAppCheck/Tests/Unit/AppAttestProvider/FIRAppAttestProviderTests.m b/FirebaseAppCheck/Tests/Unit/AppAttestProvider/FIRAppAttestProviderTests.m
@@ -16,6 +16,7 @@
 
 #import <XCTest/XCTest.h>
 
+#import <DeviceCheck/DeviceCheck.h>
 #import <OCMock/OCMock.h>
 #import "FBLPromise+Testing.h"
 
@@ -602,6 +603,75 @@ - (void)testGetToken_WhenAttestationIsRejected_ThenAttestationIsResetAndRetriedO
   [self verifyAllMocks];
 }
 
+- (void)testGetToken_WhenExistingKeyIsRejectedByApple_ThenAttestationIsResetAndRetriedOnce_Success {
+  // 1. Expect FIRAppAttestService.isSupported.
+  [OCMExpect([self.mockAppAttestService isSupported]) andReturnValue:@(YES)];
+
+  // 2. Expect storage getAppAttestKeyID.
+  NSString *existingKeyID = @"existingKeyID";
+  OCMExpect([self.mockStorage getAppAttestKeyID])
+      .andReturn([FBLPromise resolvedWith:existingKeyID]);
+
+  // 3. Expect a stored artifact to be requested.
+  __auto_type rejectedPromise = [self rejectedPromiseWithError:[NSError errorWithDomain:self.name
+                                                                                   code:NSNotFound
+                                                                               userInfo:nil]];
+  OCMExpect([self.mockArtifactStorage getArtifactForKey:existingKeyID]).andReturn(rejectedPromise);
+
+  // 4. Expect random challenge to be requested.
+  OCMExpect([self.mockAPIService getRandomChallenge])
+      .andReturn([FBLPromise resolvedWith:self.randomChallenge]);
+
+  // 5. Expect the key to be attested with the challenge.
+  NSError *attestationError = [NSError errorWithDomain:DCErrorDomain
+                                                  code:DCErrorInvalidKey
+                                              userInfo:nil];
+  id attestCompletionArg = [OCMArg invokeBlockWithArgs:[NSNull null], attestationError, nil];
+  OCMExpect([self.mockAppAttestService attestKey:existingKeyID
+                                  clientDataHash:self.randomChallengeHash
+                               completionHandler:attestCompletionArg]);
+
+  // 6. Stored attestation to be reset.
+  [self expectAttestationReset];
+
+  // 7. Expect the App Attest key pair to be generated and attested.
+  NSString *newKeyID = @"newKeyID";
+  NSData *attestationData = [[NSUUID UUID].UUIDString dataUsingEncoding:NSUTF8StringEncoding];
+  [self expectAppAttestKeyGeneratedAndAttestedWithKeyID:newKeyID attestationData:attestationData];
+
+  // 8. Expect exchange request to be sent.
+  FIRAppCheckToken *FACToken = [[FIRAppCheckToken alloc] initWithToken:@"FAC token"
+                                                        expirationDate:[NSDate date]];
+  NSData *artifactData = [@"attestation artifact" dataUsingEncoding:NSUTF8StringEncoding];
+  __auto_type attestKeyResponse =
+      [[FIRAppAttestAttestationResponse alloc] initWithArtifact:artifactData token:FACToken];
+  OCMExpect([self.mockAPIService attestKeyWithAttestation:attestationData
+                                                    keyID:newKeyID
+                                                challenge:self.randomChallenge])
+      .andReturn([FBLPromise resolvedWith:attestKeyResponse]);
+
+  // 9. Expect the artifact received from Firebase backend to be saved.
+  OCMExpect([self.mockArtifactStorage setArtifact:artifactData forKey:newKeyID])
+      .andReturn([FBLPromise resolvedWith:artifactData]);
+
+  // 10. Call get token.
+  XCTestExpectation *completionExpectation =
+      [self expectationWithDescription:@"completionExpectation"];
+  [self.provider
+      getTokenWithCompletion:^(FIRAppCheckToken *_Nullable token, NSError *_Nullable error) {
+        [completionExpectation fulfill];
+
+        XCTAssertEqualObjects(token.token, FACToken.token);
+        XCTAssertEqualObjects(token.expirationDate, FACToken.expirationDate);
+        XCTAssertNil(error);
+      }];
+
+  [self waitForExpectations:@[ completionExpectation ] timeout:0.5 enforceOrder:YES];
+
+  // 11. Verify mocks.
+  [self verifyAllMocks];
+}
+
 #pragma mark - FAC token refresh (assertion)
 
 - (void)testGetToken_WhenKeyRegistered_Success {