Implement getLimitedUseTokenWithCompletion: in DeviceCheck provider

google · Aug 14, 2023 · 1e21c01 · 1e21c01
1 parent 88edb4a
commit 1e21c01
Show file tree

Hide file tree

Showing 6 changed files with 60 additions and 41 deletions.
diff --git a/AppCheckCore/Sources/DeviceCheckProvider/API/GACDeviceCheckAPIService.h b/AppCheckCore/Sources/DeviceCheckProvider/API/GACDeviceCheckAPIService.h
@@ -24,7 +24,8 @@ NS_ASSUME_NONNULL_BEGIN
 
 @protocol GACDeviceCheckAPIServiceProtocol <NSObject>
 
-- (FBLPromise<GACAppCheckToken *> *)appCheckTokenWithDeviceToken:(NSData *)deviceToken;
+- (FBLPromise<GACAppCheckToken *> *)appCheckTokenWithDeviceToken:(NSData *)deviceToken
+                                                      limitedUse:(BOOL)limitedUse;
 
 @end
 

diff --git a/AppCheckCore/Sources/DeviceCheckProvider/API/GACDeviceCheckAPIService.m b/AppCheckCore/Sources/DeviceCheckProvider/API/GACDeviceCheckAPIService.m
@@ -35,6 +35,7 @@
 static NSString *const kContentTypeKey = @"Content-Type";
 static NSString *const kJSONContentType = @"application/json";
 static NSString *const kDeviceTokenField = @"device_token";
+static NSString *const kLimitedUseField = @"limited_use";
 
 @interface GACDeviceCheckAPIService ()
 
@@ -58,12 +59,13 @@ - (instancetype)initWithAPIService:(id<GACAppCheckAPIServiceProtocol>)APIService
 
 #pragma mark - Public API
 
-- (FBLPromise<GACAppCheckToken *> *)appCheckTokenWithDeviceToken:(NSData *)deviceToken {
+- (FBLPromise<GACAppCheckToken *> *)appCheckTokenWithDeviceToken:(NSData *)deviceToken
+                                                      limitedUse:(BOOL)limitedUse {
   NSString *URLString = [NSString stringWithFormat:@"%@/%@:exchangeDeviceCheckToken",
                                                    self.APIService.baseURL, self.resourceName];
   NSURL *URL = [NSURL URLWithString:URLString];
 
-  return [self HTTPBodyWithDeviceToken:deviceToken]
+  return [self HTTPBodyWithDeviceToken:deviceToken limitedUse:limitedUse]
       .then(^FBLPromise<GULURLSessionDataResponse *> *(NSData *HTTPBody) {
         return [self.APIService sendRequestWithURL:URL
                                         HTTPMethod:@"POST"
@@ -75,31 +77,34 @@ - (instancetype)initWithAPIService:(id<GACAppCheckAPIServiceProtocol>)APIService
       });
 }
 
-- (FBLPromise<NSData *> *)HTTPBodyWithDeviceToken:(NSData *)deviceToken {
+- (FBLPromise<NSData *> *)HTTPBodyWithDeviceToken:(NSData *)deviceToken
+                                       limitedUse:(BOOL)limitedUse {
   if (deviceToken.length <= 0) {
     FBLPromise *rejectedPromise = [FBLPromise pendingPromise];
     [rejectedPromise reject:[GACAppCheckErrorUtil
                                 errorWithFailureReason:@"DeviceCheck token must not be empty."]];
     return rejectedPromise;
   }
 
-  return [FBLPromise onQueue:[self backgroundQueue]
-                          do:^id _Nullable {
-                            NSString *base64EncodedToken =
-                                [deviceToken base64EncodedStringWithOptions:0];
-
-                            NSError *encodingError;
-                            NSData *payloadJSON = [NSJSONSerialization
-                                dataWithJSONObject:@{kDeviceTokenField : base64EncodedToken}
-                                           options:0
-                                             error:&encodingError];
-
-                            if (payloadJSON != nil) {
-                              return payloadJSON;
-                            } else {
-                              return [GACAppCheckErrorUtil JSONSerializationError:encodingError];
-                            }
-                          }];
+  return [FBLPromise
+      onQueue:[self backgroundQueue]
+           do:^id _Nullable {
+             NSString *base64EncodedToken = [deviceToken base64EncodedStringWithOptions:0];
+
+             NSError *encodingError;
+             NSData *payloadJSON = [NSJSONSerialization dataWithJSONObject:@{
+               kDeviceTokenField : base64EncodedToken,
+               kLimitedUseField : @(limitedUse)
+             }
+                                                                   options:0
+                                                                     error:&encodingError];
+
+             if (payloadJSON != nil) {
+               return payloadJSON;
+             } else {
+               return [GACAppCheckErrorUtil JSONSerializationError:encodingError];
+             }
+           }];
 }
 
 - (dispatch_queue_t)backgroundQueue {

diff --git a/AppCheckCore/Sources/DeviceCheckProvider/GACDeviceCheckProvider.m b/AppCheckCore/Sources/DeviceCheckProvider/GACDeviceCheckProvider.m
@@ -93,9 +93,22 @@ - (instancetype)initWithServiceName:(NSString *)serviceName
 
 - (void)getTokenWithCompletion:(void (^)(GACAppCheckToken *_Nullable token,
                                          NSError *_Nullable error))handler {
+  [self getTokenWithLimitedUse:NO completion:handler];
+}
+
+- (void)getLimitedUseTokenWithCompletion:(void (^)(GACAppCheckToken *_Nullable,
+                                                   NSError *_Nullable))handler {
+  [self getTokenWithLimitedUse:YES completion:handler];
+}
+
+#pragma mark - Internal
+
+- (void)getTokenWithLimitedUse:(BOOL)limitedUse
+                    completion:(void (^)(GACAppCheckToken *_Nullable token,
+                                         NSError *_Nullable error))handler {
   [self.backoffWrapper
       applyBackoffToOperation:^FBLPromise *_Nonnull {
-        return [self getTokenPromise];
+        return [self getTokenPromiseWithLimitedUse:limitedUse];
       }
                  errorHandler:[self.backoffWrapper defaultAppCheckProviderErrorHandler]]
       // Call the handler with either token or error.
@@ -108,20 +121,12 @@ - (void)getTokenWithCompletion:(void (^)(GACAppCheckToken *_Nullable token,
       });
 }
 
-- (void)getLimitedUseTokenWithCompletion:(void (^)(GACAppCheckToken *_Nullable,
-                                                   NSError *_Nullable))handler {
-  // TODO(andrewheard): Add support for generating limited-use tokens with a 5-minute TTL.
-  [self getTokenWithCompletion:handler];
-}
-
-#pragma mark - Internal
-
-- (FBLPromise<GACAppCheckToken *> *)getTokenPromise {
+- (FBLPromise<GACAppCheckToken *> *)getTokenPromiseWithLimitedUse:(BOOL)limitedUse {
   // Get DeviceCheck token
   return [self deviceToken]
       // Exchange DeviceCheck token for FAC token.
       .then(^FBLPromise<GACAppCheckToken *> *(NSData *deviceToken) {
-        return [self.APIService appCheckTokenWithDeviceToken:deviceToken];
+        return [self.APIService appCheckTokenWithDeviceToken:deviceToken limitedUse:limitedUse];
       });
 }
 

diff --git a/AppCheckCore/Tests/Integration/GACDeviceCheckAPIServiceE2ETests.m b/AppCheckCore/Tests/Integration/GACDeviceCheckAPIServiceE2ETests.m
@@ -67,7 +67,7 @@ - (void)tearDown {
 // TODO: Re-enable the test once secret with "GoogleService-Info.plist" is configured.
 - (void)temporaryDisabled_testAppCheckTokenSuccess {
   __auto_type appCheckPromise =
-      [self.deviceCheckAPIService appCheckTokenWithDeviceToken:[NSData data]];
+      [self.deviceCheckAPIService appCheckTokenWithDeviceToken:[NSData data] limitedUse:NO];
 
   XCTAssert(FBLWaitForPromisesWithTimeout(20));
 

diff --git a/AppCheckCore/Tests/Unit/DeviceCheckProvider/GACDeviceCheckAPIServiceTests.m b/AppCheckCore/Tests/Unit/DeviceCheckProvider/GACDeviceCheckAPIServiceTests.m
@@ -97,7 +97,8 @@ - (void)testAppCheckTokenSuccess {
       .andReturn([FBLPromise resolvedWith:expectedResult]);
 
   // 2. Send request.
-  __auto_type tokenPromise = [self.APIService appCheckTokenWithDeviceToken:deviceTokenData];
+  __auto_type tokenPromise = [self.APIService appCheckTokenWithDeviceToken:deviceTokenData
+                                                                limitedUse:NO];
 
   // 3. Verify.
   XCTAssert(FBLWaitForPromisesWithTimeout(1));
@@ -150,7 +151,8 @@ - (void)testAppCheckTokenResponseParsingError {
       .andReturn(rejectedPromise);
 
   // 2. Send request.
-  __auto_type tokenPromise = [self.APIService appCheckTokenWithDeviceToken:deviceTokenData];
+  __auto_type tokenPromise = [self.APIService appCheckTokenWithDeviceToken:deviceTokenData
+                                                                limitedUse:NO];
 
   // 3. Verify.
   XCTAssert(FBLWaitForPromisesWithTimeout(1));
@@ -180,7 +182,8 @@ - (void)testAppCheckTokenNetworkError {
       .andReturn(rejectedPromise);
 
   // 2. Send request.
-  __auto_type tokenPromise = [self.APIService appCheckTokenWithDeviceToken:deviceTokenData];
+  __auto_type tokenPromise = [self.APIService appCheckTokenWithDeviceToken:deviceTokenData
+                                                                limitedUse:NO];
 
   // 3. Verify.
   XCTAssert(FBLWaitForPromisesWithTimeout(1));
@@ -202,7 +205,8 @@ - (void)testAppCheckTokenEmptyDeviceToken {
                                   additionalHeaders:[OCMArg any]]);
 
   // 2. Send request.
-  __auto_type tokenPromise = [self.APIService appCheckTokenWithDeviceToken:deviceTokenData];
+  __auto_type tokenPromise = [self.APIService appCheckTokenWithDeviceToken:deviceTokenData
+                                                                limitedUse:NO];
 
   // 3. Verify.
   XCTAssert(FBLWaitForPromisesWithTimeout(1));

diff --git a/AppCheckCore/Tests/Unit/DeviceCheckProvider/GACDeviceCheckProviderTests.m b/AppCheckCore/Tests/Unit/DeviceCheckProvider/GACDeviceCheckProviderTests.m
@@ -81,8 +81,9 @@ - (void)testGetTokenSuccess {
   GACAppCheckToken *validToken = [[GACAppCheckToken alloc] initWithToken:@"valid_token"
                                                           expirationDate:[NSDate distantFuture]
                                                           receivedAtDate:[NSDate date]];
-  OCMExpect([self.fakeAPIService appCheckTokenWithDeviceToken:deviceToken])
+  OCMExpect([self.fakeAPIService appCheckTokenWithDeviceToken:deviceToken limitedUse:NO])
       .andReturn([FBLPromise resolvedWith:validToken]);
+  OCMReject([self.fakeAPIService appCheckTokenWithDeviceToken:OCMOCK_ANY limitedUse:YES]);
 
   // 3. Expect backoff wrapper to be used.
   self.fakeBackoffWrapper.backoffExpectation = [self expectationWithDescription:@"Backoff"];
@@ -136,7 +137,8 @@ - (void)testGetTokenWhenDeviceTokenFails {
   OCMExpect([self.fakeTokenGenerator generateTokenWithCompletionHandler:generateTokenArg]);
 
   // 2. Don't expect FAA token to be requested.
-  OCMReject([self.fakeAPIService appCheckTokenWithDeviceToken:[OCMArg any]]);
+  OCMReject([self.fakeAPIService appCheckTokenWithDeviceToken:OCMOCK_ANY limitedUse:NO])
+      .ignoringNonObjectArgs();
 
   // 3. Call getToken and validate the result.
   XCTestExpectation *completionExpectation =
@@ -186,8 +188,9 @@ - (void)testGetTokenWhenAPIServiceFails {
   // 2. Expect FAA token to be requested.
   FBLPromise *rejectedPromise = [FBLPromise pendingPromise];
   [rejectedPromise reject:APIServiceError];
-  OCMExpect([self.fakeAPIService appCheckTokenWithDeviceToken:deviceToken])
+  OCMExpect([self.fakeAPIService appCheckTokenWithDeviceToken:deviceToken limitedUse:NO])
       .andReturn(rejectedPromise);
+  OCMReject([self.fakeAPIService appCheckTokenWithDeviceToken:OCMOCK_ANY limitedUse:YES]);
 
   // 3. Call getToken and validate the result.
   XCTestExpectation *completionExpectation =
@@ -221,7 +224,8 @@ - (void)testGetTokenBackoff {
   self.fakeBackoffWrapper.backoffExpectation = [self expectationWithDescription:@"Backoff"];
 
   // 2. Don't expect any operations.
-  OCMReject([self.fakeAPIService appCheckTokenWithDeviceToken:[OCMArg any]]);
+  OCMReject([self.fakeAPIService appCheckTokenWithDeviceToken:OCMOCK_ANY limitedUse:NO])
+      .ignoringNonObjectArgs();
   OCMReject([self.fakeTokenGenerator generateTokenWithCompletionHandler:OCMOCK_ANY]);
 
   // 3. Call getToken and validate the result.