Adds Auth Result to Sign In with Custom token (#489)

* Adds Auth Result to Sign In with Custom token

* Addresses comomments

7 files changed, 199 insertions, 15 deletions

diff --git a/Example/Auth/Sample/MainViewController.m b/Example/Auth/Sample/MainViewController.m
index 6980ec6..5326463 100644
--- a/Example/Auth/Sample/MainViewController.m
+++ b/Example/Auth/Sample/MainViewController.m
@@ -118,6 +118,12 @@ static NSString *const kSignInEmailPasswordAuthDataResultButtonText =
  */
 static NSString *const kSignInWithCustomTokenButtonText = @"Sign In (BYOAuth)";
 
+/** @var kSignInWithCustomAuthResultTokenButtonText
+    @brief The text of the "Sign In with Custom Token (Auth Result)" button.
+ */
+static NSString *const kSignInWithCustomAuthResultTokenButtonText = @"Sign In with Custom Token"
+    " (Auth Result)";
+
 /** @var kSignInAnonymouslyButtonText
     @brief The text of the "Sign In Anonymously" button.
  */
@@ -722,6 +728,8 @@ typedef enum {
                                            action:^{ [weakSelf signInEmailPasswordAuthDataResult]; }],
         [StaticContentTableViewCell cellWithTitle:kSignInWithCustomTokenButtonText
                                            action:^{ [weakSelf signInWithCustomToken]; }],
+        [StaticContentTableViewCell cellWithTitle:kSignInWithCustomAuthResultTokenButtonText
+                                           action:^{ [weakSelf signInWithCustomTokenAuthResult]; }],
         [StaticContentTableViewCell cellWithTitle:kSignInAnonymouslyButtonText
                                            action:^{ [weakSelf signInAnonymously]; }],
         [StaticContentTableViewCell cellWithTitle:kSignInAnonymouslyWithAuthResultButtonText
@@ -1754,6 +1762,24 @@ static NSDictionary<NSString *, NSString *> *parseURL(NSString *urlString) {
   [self presentViewController:dataEntryViewController animated:YES completion:nil];
 }
 
+/** @fn signInWithCustomTokenAuthResult
+    @brief Signs the user in using a manually-entered custom token.
+ */
+- (void)signInWithCustomTokenAuthResult {
+  CustomTokenDataEntryViewControllerCompletion action =
+      ^(BOOL cancelled, NSString *_Nullable userEnteredTokenText) {
+        if (cancelled) {
+          [self log:@"CANCELLED:sign-in with custom token cancelled."];
+          return;
+        }
+
+        [self doSignInAndRetrieveDataWithCustomToken:userEnteredTokenText];
+      };
+  CustomTokenDataEntryViewController *dataEntryViewController =
+      [[CustomTokenDataEntryViewController alloc] initWithCompletion:action];
+  [self presentViewController:dataEntryViewController animated:YES completion:nil];
+}
+
 /** @fn signOut
     @brief Signs the user out.
  */
@@ -3070,6 +3096,26 @@ static NSDictionary<NSString *, NSString *> *parseURL(NSString *urlString) {
   }];
 }
 
+- (void)doSignInAndRetrieveDataWithCustomToken:(NSString *_Nullable)userEnteredTokenText {
+  [[AppManager auth] signInAndRetrieveDataWithCustomToken:userEnteredTokenText
+                                               completion:^(FIRAuthDataResult *_Nullable result,
+                                                            NSError *_Nullable error) {
+    if (error) {
+      [self logFailure:@"sign-in with custom token failed" error:error];
+      [self showMessagePromptWithTitle:kSignInErrorAlertTitle
+                               message:error.localizedDescription
+                      showCancelButton:NO
+                            completion:nil];
+      return;
+    }
+    [self logSuccess:@"sign-in with custom token succeeded."];
+    [self showMessagePromptWithTitle:kSignedInAlertTitle
+                             message:result.user.displayName
+                    showCancelButton:NO
+                          completion:nil];
+  }];
+}
+
 - (void)updateUserInfo {
   [_userInfoTableViewCell updateContentsWithUser:[AppManager auth].currentUser];
   [_userInMemoryInfoTableViewCell updateContentsWithUser:_userInMemory];
diff --git a/Example/Auth/Tests/FIRAuthTests.m b/Example/Auth/Tests/FIRAuthTests.m
index 00a7c7a..80a9ae9 100644
--- a/Example/Auth/Tests/FIRAuthTests.m
+++ b/Example/Auth/Tests/FIRAuthTests.m
@@ -1248,6 +1248,61 @@ static const NSTimeInterval kWaitInterval = .5;
   OCMVerifyAll(_mockBackend);
 }
 
+/** @fn testSignInAndRetrieveDataWithCustomTokenSuccess
+    @brief Tests the flow of a successful @c signInAndRetrieveDataWithCustomToken:completion: call.
+ */
+- (void)testSignInAndRetrieveDataWithCustomTokenSuccess {
+  OCMExpect([_mockBackend verifyCustomToken:[OCMArg any] callback:[OCMArg any]])
+      .andCallBlock2(^(FIRVerifyCustomTokenRequest *_Nullable request,
+                       FIRVerifyCustomTokenResponseCallback callback) {
+    XCTAssertEqualObjects(request.APIKey, kAPIKey);
+    XCTAssertEqualObjects(request.token, kCustomToken);
+    XCTAssertTrue(request.returnSecureToken);
+    dispatch_async(FIRAuthGlobalWorkQueue(), ^() {
+      id mockVeriyCustomTokenResponse = OCMClassMock([FIRVerifyCustomTokenResponse class]);
+      [self stubTokensWithMockResponse:mockVeriyCustomTokenResponse];
+      callback(mockVeriyCustomTokenResponse, nil);
+    });
+  });
+  [self expectGetAccountInfo];
+  XCTestExpectation *expectation = [self expectationWithDescription:@"callback"];
+  [[FIRAuth auth] signOut:NULL];
+  [[FIRAuth auth] signInAndRetrieveDataWithCustomToken:kCustomToken
+                                            completion:^(FIRAuthDataResult *_Nullable result,
+                                                         NSError *_Nullable error) {
+    XCTAssertTrue([NSThread isMainThread]);
+    [self assertUser:result.user];
+    XCTAssertFalse(result.additionalUserInfo.isNewUser);
+    XCTAssertNil(error);
+    [expectation fulfill];
+  }];
+  [self waitForExpectationsWithTimeout:kExpectationTimeout handler:nil];
+  [self assertUser:[FIRAuth auth].currentUser];
+  OCMVerifyAll(_mockBackend);
+}
+
+/** @fn testSignInAndRetrieveDataWithCustomTokenFailure
+    @brief Tests the flow of a failed @c signInAndRetrieveDataWithCustomToken:completion: call.
+ */
+- (void)testSignInAndRetrieveDataWithCustomTokenFailure {
+  OCMExpect([_mockBackend verifyCustomToken:[OCMArg any] callback:[OCMArg any]])
+      .andDispatchError2([FIRAuthErrorUtils invalidCustomTokenErrorWithMessage:nil]);
+  XCTestExpectation *expectation = [self expectationWithDescription:@"callback"];
+  [[FIRAuth auth] signOut:NULL];
+  [[FIRAuth auth] signInAndRetrieveDataWithCustomToken:kCustomToken
+                                            completion:^(FIRAuthDataResult *_Nullable result,
+                                                         NSError *_Nullable error) {
+    XCTAssertTrue([NSThread isMainThread]);
+    XCTAssertNil(result);
+    XCTAssertEqual(error.code, FIRAuthErrorCodeInvalidCustomToken);
+    XCTAssertNotNil(error.userInfo[NSLocalizedDescriptionKey]);
+    [expectation fulfill];
+  }];
+  [self waitForExpectationsWithTimeout:kExpectationTimeout handler:nil];
+  XCTAssertNil([FIRAuth auth].currentUser);
+  OCMVerifyAll(_mockBackend);
+}
+
 /** @fn testCreateUserWithEmailPasswordSuccess
     @brief Tests the flow of a successful @c createUserWithEmail:password:completion: call.
  */
diff --git a/Example/Auth/Tests/FIRVerifyCustomTokenResponseTests.m b/Example/Auth/Tests/FIRVerifyCustomTokenResponseTests.m
index b1d9d97..9fe95b1 100644
--- a/Example/Auth/Tests/FIRVerifyCustomTokenResponseTests.m
+++ b/Example/Auth/Tests/FIRVerifyCustomTokenResponseTests.m
@@ -47,6 +47,11 @@ static NSString *const kExpiresInKey = @"expiresIn";
  */
 static NSString *const kRefreshTokenKey = @"refreshToken";
 
+/** @var kIsNewUserKey
+    @brief The name of the "isNewUser" property in the response.
+ */
+static NSString *const kIsNewUserKey = @"isNewUser";
+
 /** @var kTestIDToken
     @brief Testing ID token for verifying assertion.
  */
@@ -274,6 +279,7 @@ static const double kAllowedTimeDifference = 0.1;
     kIDTokenKey : kTestIDToken,
     kExpiresInKey : kTestExpiresIn,
     kRefreshTokenKey : kTestRefreshToken,
+    kIsNewUserKey : @YES,
   }];
   XCTAssert(callbackInvoked);
   XCTAssertNil(RPCError);
@@ -282,6 +288,7 @@ static const double kAllowedTimeDifference = 0.1;
   NSTimeInterval expiresIn = [RPCResponse.approximateExpirationDate timeIntervalSinceNow];
   XCTAssertEqualWithAccuracy(expiresIn, [kTestExpiresIn doubleValue], kAllowedTimeDifference);
   XCTAssertEqualObjects(RPCResponse.refreshToken, kTestRefreshToken);
+  XCTAssertTrue(RPCResponse.isNewUser);
 }
 
 @end
diff --git a/Firebase/Auth/Source/FIRAuth.m b/Firebase/Auth/Source/FIRAuth.m
index a104ad8..cadff23 100644
--- a/Firebase/Auth/Source/FIRAuth.m
+++ b/Firebase/Auth/Source/FIRAuth.m
@@ -790,25 +790,23 @@ static NSMutableDictionary *gKeychainServiceNameForAppName;
   dispatch_async(FIRAuthGlobalWorkQueue(), ^{
     FIRAuthResultCallback decoratedCallback =
         [self signInFlowAuthResultCallbackByDecoratingCallback:completion];
-    FIRVerifyCustomTokenRequest *request =
-        [[FIRVerifyCustomTokenRequest alloc] initWithToken:token
-                                      requestConfiguration:_requestConfiguration];
-    [FIRAuthBackend verifyCustomToken:request
-                             callback:^(FIRVerifyCustomTokenResponse *_Nullable response,
-                                        NSError *_Nullable error) {
-      if (error) {
-        decoratedCallback(nil, error);
-        return;
-      }
-      [self completeSignInWithAccessToken:response.IDToken
-                accessTokenExpirationDate:response.approximateExpirationDate
-                             refreshToken:response.refreshToken
-                                anonymous:NO
-                                 callback:decoratedCallback];
+    [self internalSignInAndRetrieveDataWithCustomToken:token
+                                            completion:^(FIRAuthDataResult *_Nullable authResult,
+                                                         NSError *_Nullable error) {
+      decoratedCallback(authResult.user, error);
     }];
   });
 }
 
+- (void)signInAndRetrieveDataWithCustomToken:(NSString *)token
+                                  completion:(nullable FIRAuthDataResultCallback)completion {
+  dispatch_async(FIRAuthGlobalWorkQueue(), ^{
+    FIRAuthDataResultCallback decoratedCallback =
+        [self signInFlowAuthDataResultCallbackByDecoratingCallback:completion];
+    [self internalSignInAndRetrieveDataWithCustomToken:token completion:decoratedCallback];
+  });
+}
+
 - (void)createUserWithEmail:(NSString *)email
                    password:(NSString *)password
                  completion:(nullable FIRAuthResultCallback)completion {
@@ -1196,6 +1194,52 @@ static NSMutableDictionary *gKeychainServiceNameForAppName;
 }
 #endif
 
+/** @fn internalSignInAndRetrieveDataWithCustomToken:completion:
+    @brief Signs in a Firebase user given a custom token.
+    @param token A self-signed custom auth token.
+    @param completion A block which is invoked when the custom token sign in request completes.
+ */
+- (void)internalSignInAndRetrieveDataWithCustomToken:(NSString *)token
+                                          completion:(nullable FIRAuthDataResultCallback)
+                                              completion {
+  FIRVerifyCustomTokenRequest *request =
+    [[FIRVerifyCustomTokenRequest alloc] initWithToken:token
+                                  requestConfiguration:_requestConfiguration];
+  [FIRAuthBackend verifyCustomToken:request
+                           callback:^(FIRVerifyCustomTokenResponse *_Nullable response,
+                                      NSError *_Nullable error) {
+    if (error) {
+      if (completion) {
+        completion(nil, error);
+        return;
+      }
+    }
+    [self completeSignInWithAccessToken:response.IDToken
+              accessTokenExpirationDate:response.approximateExpirationDate
+                           refreshToken:response.refreshToken
+                              anonymous:NO
+                               callback:^(FIRUser *_Nullable user,
+                                          NSError *_Nullable error) {
+      if (error) {
+        if (completion) {
+          completion(nil, error);
+        }
+        return;
+      }
+      FIRAdditionalUserInfo *additonalUserInfo =
+          [[FIRAdditionalUserInfo alloc] initWithProviderID:nil
+                                                   profile:nil
+                                                  username:nil
+                                                 isNewUser:response.isNewUser];
+      FIRAuthDataResult *result =
+          [[FIRAuthDataResult alloc] initWithUser:user additionalUserInfo:additonalUserInfo];
+      if (completion) {
+        completion(result, nil);
+      }
+    }];
+  }];
+}
+
 /** @fn internalCreateUserWithEmail:password:completion:
     @brief Makes a backend request attempting to create a new Firebase user given an email address
         and password.
diff --git a/Firebase/Auth/Source/Public/FIRAuth.h b/Firebase/Auth/Source/Public/FIRAuth.h
index 9a115ee..2935916 100644
--- a/Firebase/Auth/Source/Public/FIRAuth.h
+++ b/Firebase/Auth/Source/Public/FIRAuth.h
@@ -486,6 +486,32 @@ FIR_SWIFT_NAME(Auth)
 - (void)signInWithCustomToken:(NSString *)token
                    completion:(nullable FIRAuthResultCallback)completion;
 
+/** @fn signInAndRetrieveDataWithCustomToken:completion:
+    @brief Asynchronously signs in to Firebase with the given Auth token.
+
+    @param token A self-signed custom auth token.
+    @param completion Optionally; a block which is invoked when the sign in finishes, or is
+        canceled. Invoked asynchronously on the main thread in the future.
+
+    @remarks Possible error codes:
+    <ul>
+        <li>@c FIRAuthErrorCodeInvalidCustomToken - Indicates a validation error with
+            the custom token.
+        </li>
+        <li>@c FIRAuthErrorCodeCustomTokenMismatch - Indicates the service account and the API key
+            belong to different projects.
+        </li>
+    </ul>
+
+    @remarks See @c FIRAuthErrors for a list of error codes that are common to all API methods.
+    @remarks This method will only exist until the next major Firebase release following 4.x.x.
+        After the next major release the method @c createUserWithEmail:password:completion: will
+        support the @c FIRAuthDataResultCallback.
+ */
+- (void)signInAndRetrieveDataWithCustomToken:(NSString *)token
+                                  completion:(nullable FIRAuthDataResultCallback)completion;
+
+
 /** @fn createUserWithEmail:password:completion:
     @brief Creates and, on success, signs in a user with the given email address and password.
 
diff --git a/Firebase/Auth/Source/RPCs/FIRVerifyCustomTokenResponse.h b/Firebase/Auth/Source/RPCs/FIRVerifyCustomTokenResponse.h
index b8c215c..6957bf3 100644
--- a/Firebase/Auth/Source/RPCs/FIRVerifyCustomTokenResponse.h
+++ b/Firebase/Auth/Source/RPCs/FIRVerifyCustomTokenResponse.h
@@ -42,6 +42,11 @@ NS_ASSUME_NONNULL_BEGIN
  */
 @property(nonatomic, strong, readonly, nullable) NSString *refreshToken;
 
+/** @property isNewUser
+    @brief Flag indicating that the user signing in is a new user and not a returning user.
+ */
+@property(nonatomic, assign) BOOL isNewUser;
+
 @end
 
 NS_ASSUME_NONNULL_END
diff --git a/Firebase/Auth/Source/RPCs/FIRVerifyCustomTokenResponse.m b/Firebase/Auth/Source/RPCs/FIRVerifyCustomTokenResponse.m
index 12dc203..8a87141 100644
--- a/Firebase/Auth/Source/RPCs/FIRVerifyCustomTokenResponse.m
+++ b/Firebase/Auth/Source/RPCs/FIRVerifyCustomTokenResponse.m
@@ -26,6 +26,7 @@
   _approximateExpirationDate = [dictionary[@"expiresIn"] isKindOfClass:[NSString class]] ?
       [NSDate dateWithTimeIntervalSinceNow:[dictionary[@"expiresIn"] doubleValue]] : nil;
   _refreshToken = [dictionary[@"refreshToken"] copy];
+  _isNewUser = [dictionary[@"isNewUser"] boolValue];
   return YES;
 }