Reset App Attest key state if attestKey fails (#41)

Author: andrewheard

AppCheckCore/Sources/AppAttestProvider/GACAppAttestProvider.m

@@ -361,6 +361,27 @@ - (void)getTokenWithLimitedUse:(BOOL)limitedUse
 
                 return [self attestKey:keyID challenge:challenge];
               })
+      .recoverOn(self.queue,
+                 ^id(NSError *error) {
+                   // If Apple rejected the key (DCErrorInvalidKey) then reset the attestation and
+                   // throw a specific error to signal retry (GACAppAttestRejectionError).
+                   NSError *underlyingError = error.userInfo[NSUnderlyingErrorKey];
+                   if (underlyingError && [underlyingError.domain isEqualToString:DCErrorDomain] &&
+                       underlyingError.code == DCErrorInvalidKey) {
+                     GACAppCheckLog(
+                         GACLoggerAppCheckMessageCodeAttestationRejected, GACAppCheckLogLevelDebug,
+                         @"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 [[GACAppAttestRejectionError alloc] init];
+                     });
+                   }
+
+                   // Otherwise just re-throw the error.
+                   return error;
+                 })
       .thenOn(self.queue,
               ^FBLPromise<NSArray *> *(GACAppAttestKeyAttestationResult *result) {
                 // 3. Exchange the attestation to FAC token and pass the results to the next step.

AppCheckCore/Tests/Unit/AppAttestProvider/GACAppAttestProviderTests.m

@@ -14,13 +14,14 @@
  * limitations under the License.
  */
 
+#import "AppCheckCore/Sources/Public/AppCheckCore/GACAppAttestProvider.h"
+
+#import <DeviceCheck/DeviceCheck.h>
 #import <XCTest/XCTest.h>
 
 #import <OCMock/OCMock.h>
 #import "FBLPromise+Testing.h"
 
-#import "AppCheckCore/Sources/Public/AppCheckCore/GACAppAttestProvider.h"
-
 #import "AppCheckCore/Sources/AppAttestProvider/API/GACAppAttestAPIService.h"
 #import "AppCheckCore/Sources/AppAttestProvider/API/GACAppAttestAttestationResponse.h"
 #import "AppCheckCore/Sources/AppAttestProvider/GACAppAttestService.h"
@@ -626,6 +627,76 @@ - (void)testGetToken_WhenAttestationIsRejected_ThenAttestationIsResetAndRetriedO
   [self verifyAllMocks];
 }
 
+- (void)testGetToken_WhenExistingKeyIsRejectedByApple_ThenAttestationIsResetAndRetriedOnce_Success {
+  // 1. Expect GACAppAttestService.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.
+  GACAppCheckToken *appCheckToken = [[GACAppCheckToken alloc] initWithToken:@"App Check Token"
+                                                             expirationDate:[NSDate date]];
+  NSData *artifactData = [@"attestation artifact" dataUsingEncoding:NSUTF8StringEncoding];
+  __auto_type attestKeyResponse =
+      [[GACAppAttestAttestationResponse alloc] initWithArtifact:artifactData token:appCheckToken];
+  OCMExpect([self.mockAPIService attestKeyWithAttestation:attestationData
+                                                    keyID:newKeyID
+                                                challenge:self.randomChallenge
+                                               limitedUse:NO])
+      .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:^(GACAppCheckToken *_Nullable token, NSError *_Nullable error) {
+        [completionExpectation fulfill];
+
+        XCTAssertEqualObjects(token.token, appCheckToken.token);
+        XCTAssertEqualObjects(token.expirationDate, appCheckToken.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 {