Add GIDKeychainHandler implementation (#265)

GoogleSignIn/Sources/GIDKeychainHandler/API/GIDKeychainHandler.h

@@ -0,0 +1,40 @@
+/*
+ * Copyright 2022 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#import <Foundation/Foundation.h>
+
+@class OIDAuthState;
+
+NS_ASSUME_NONNULL_BEGIN
+
+@protocol GIDKeychainHandler <NSObject>
+
+/// Loads the OIDAuthState object from the keychain.
+///
+/// @return an OIDAuthState object or nil if the keychain is empty.
+- (nullable OIDAuthState *)loadAuthState;
+
+/// Saves the OIDAuthState object to the keychain.
+///
+/// @return A `BOOL` indicating if the save succeeded.
+- (BOOL)saveAuthState:(OIDAuthState *)authState;
+
+/// Removes the OIDAuthState object saved in the keychain.
+- (void)removeAllKeychainEntries;
+
+@end
+
+NS_ASSUME_NONNULL_END

GoogleSignIn/Sources/GIDKeychainHandler/Implementations/Fakes/GIDFakeKeychainHandler.h

@@ -0,0 +1,31 @@
+/*
+ * Copyright 2022 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#import <Foundation/Foundation.h>
+
+#import "GoogleSignIn/Sources/GIDKeychainHandler/API/GIDKeychainHandler.h"
+
+NS_ASSUME_NONNULL_BEGIN
+
+@interface GIDFakeKeychainHandler : NSObject<GIDKeychainHandler>
+
+/// If YES, the method `saveAuthState:` returns NO and cleans the saved authState in the keychain.
+/// The default value is NO.
+@property(nonatomic) BOOL failToSave;
+
+@end
+
+NS_ASSUME_NONNULL_END

GoogleSignIn/Sources/GIDKeychainHandler/Implementations/Fakes/GIDFakeKeychainHandler.m

@@ -0,0 +1,33 @@
+#import "GoogleSignIn/Sources/GIDKeychainHandler/Implementations/Fakes/GIDFakeKeychainHandler.h"
+
+NS_ASSUME_NONNULL_BEGIN
+
+@interface GIDFakeKeychainHandler ()
+
+@property(nonatomic, nullable) OIDAuthState *savedAuthState;
+
+@end
+
+@implementation GIDFakeKeychainHandler
+
+- (nullable OIDAuthState *)loadAuthState {
+  return self.savedAuthState;
+}
+
+- (BOOL)saveAuthState:(OIDAuthState *)authState {
+  if (self.failToSave) {
+    self.savedAuthState = nil;
+    return NO;
+  } else {
+    self.savedAuthState = authState;
+    return YES;
+  }
+}
+
+- (void)removeAllKeychainEntries {
+  self.savedAuthState = nil;
+}
+
+@end
+
+NS_ASSUME_NONNULL_END

GoogleSignIn/Sources/GIDKeychainHandler/Implementations/GIDKeychainHandler.h

@@ -0,0 +1,26 @@
+/*
+ * Copyright 2022 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#import <Foundation/Foundation.h>
+
+#import "GoogleSignIn/Sources/GIDKeychainHandler/API/GIDKeychainHandler.h"
+
+NS_ASSUME_NONNULL_BEGIN
+
+@interface GIDKeychainHandler : NSObject<GIDKeychainHandler>
+@end
+
+NS_ASSUME_NONNULL_END

GoogleSignIn/Sources/GIDKeychainHandler/Implementations/GIDKeychainHandler.m

@@ -0,0 +1,39 @@
+#import "GoogleSignIn/Sources/GIDKeychainHandler/Implementations/GIDKeychainHandler.h"
+
+#ifdef SWIFT_PACKAGE
+@import AppAuth;
+@import GTMAppAuth;
+#else
+#import <AppAuth/AppAuth.h>
+#import <GTMAppAuth/GTMAppAuth.h>
+#endif
+
+static NSString *const kGTMAppAuthKeychainName = @"auth";
+
+NS_ASSUME_NONNULL_BEGIN
+
+@implementation GIDKeychainHandler
+
+- (nullable OIDAuthState *)loadAuthState {
+  GTMAppAuthFetcherAuthorization *authorization =
+      [GTMAppAuthFetcherAuthorization authorizationFromKeychainForName:kGTMAppAuthKeychainName
+                                             useDataProtectionKeychain:YES];
+  return authorization.authState;
+}
+
+- (BOOL)saveAuthState:(OIDAuthState *)authState {
+  GTMAppAuthFetcherAuthorization *authorization =
+      [[GTMAppAuthFetcherAuthorization alloc] initWithAuthState:authState];
+  return [GTMAppAuthFetcherAuthorization saveAuthorization:authorization
+                                         toKeychainForName:kGTMAppAuthKeychainName
+                                 useDataProtectionKeychain:YES];
+}
+
+- (void)removeAllKeychainEntries {
+  [GTMAppAuthFetcherAuthorization removeAuthorizationFromKeychainForName:kGTMAppAuthKeychainName
+                                               useDataProtectionKeychain:YES];
+}
+
+@end
+
+NS_ASSUME_NONNULL_END

GoogleSignIn/Sources/GIDSignIn.m

@@ -22,6 +22,8 @@
 #import "GoogleSignIn/Sources/Public/GoogleSignIn/GIDUserAuth.h"
 
 #import "GoogleSignIn/Sources/GIDEMMSupport.h"
+#import "GoogleSignIn/Sources/GIDKeychainHandler/API/GIDKeychainHandler.h"
+#import "GoogleSignIn/Sources/GIDKeychainHandler/Implementations/GIDKeychainHandler.h"
 #import "GoogleSignIn/Sources/GIDSignInInternalOptions.h"
 #import "GoogleSignIn/Sources/GIDSignInPreferences.h"
 #import "GoogleSignIn/Sources/GIDCallbackQueue.h"
@@ -166,6 +168,8 @@ static NSString *const kConfigOpenIDRealmKey = @"GIDOpenIDRealm";
   id<OIDExternalUserAgentSession> _currentAuthorizationFlow;
   // Flag to indicate that the auth flow is restarting.
   BOOL _restarting;
+  
+  id<GIDKeychainHandler> _keychainHandler;
 }
 
 #pragma mark - Public methods
@@ -194,7 +198,7 @@ static NSString *const kConfigOpenIDRealmKey = @"GIDOpenIDRealm";
   if ([_currentUser.authState isAuthorized]) {
     return YES;
   }
-  OIDAuthState *authState = [self loadAuthState];
+  OIDAuthState *authState = [_keychainHandler loadAuthState];
   return [authState isAuthorized];
 }
 
@@ -216,7 +220,7 @@ static NSString *const kConfigOpenIDRealmKey = @"GIDOpenIDRealm";
   }
 
   // Try retrieving an authorization object from the keychain.
-  OIDAuthState *authState = [self loadAuthState];
+  OIDAuthState *authState = [_keychainHandler loadAuthState];
   if (!authState) {
     return NO;
   }
@@ -381,7 +385,7 @@ static NSString *const kConfigOpenIDRealmKey = @"GIDOpenIDRealm";
     self.currentUser = nil;
   }
   // Remove all state from the keychain.
-  [self removeAllKeychainEntries];
+  [_keychainHandler removeAllKeychainEntries];
 }
 
 - (void)disconnectWithCompletion:(nullable GIDDisconnectCompletion)completion {
@@ -389,7 +393,7 @@ static NSString *const kConfigOpenIDRealmKey = @"GIDOpenIDRealm";
   if (!authState) {
     // Even the user is not signed in right now, we still need to remove any token saved in the
     // keychain.
-    authState = [self loadAuthState];
+    authState = [_keychainHandler loadAuthState];
   }
   // Either access or refresh token would work, but we won't have access token if the auth is
   // retrieved from keychain.
@@ -447,6 +451,11 @@ static NSString *const kConfigOpenIDRealmKey = @"GIDOpenIDRealm";
 #pragma mark - Private methods
 
 - (id)initPrivate {
+  GIDKeychainHandler *keychainHandler = [[GIDKeychainHandler alloc] init];
+  return [self initWithKeychainHandler:keychainHandler];
+}
+
+- (instancetype)initWithKeychainHandler:(id<GIDKeychainHandler>)keychainHandler {
   self = [super init];
   if (self) {
     // Get the bundle of the current executable.
@@ -462,7 +471,7 @@ static NSString *const kConfigOpenIDRealmKey = @"GIDOpenIDRealm";
 
     // If this is a fresh install, ensure that any pre-existing keychain data is purged.
     if (isFreshInstall) {
-      [self removeAllKeychainEntries];
+      [_keychainHandler removeAllKeychainEntries];
     }
 
     NSString *authorizationEnpointURL = [NSString stringWithFormat:kAuthorizationURLTemplate,
@@ -480,6 +489,8 @@ static NSString *const kConfigOpenIDRealmKey = @"GIDOpenIDRealm";
                                           keychainName:kGTMAppAuthKeychainName
                                         isFreshInstall:isFreshInstall];
 #endif // TARGET_OS_IOS && !TARGET_OS_MACCATALYST
+    
+    _keychainHandler = keychainHandler;
   }
   return self;
 }
@@ -671,7 +682,7 @@ static NSString *const kConfigOpenIDRealmKey = @"GIDOpenIDRealm";
   }
 
   // Try retrieving an authorization object from the keychain.
-  OIDAuthState *authState = [self loadAuthState];
+  OIDAuthState *authState = [_keychainHandler loadAuthState];
 
   if (![authState isAuthorized]) {
     // No valid auth in keychain, per documentation/spec, notify callback of failure.
@@ -766,7 +777,7 @@ static NSString *const kConfigOpenIDRealmKey = @"GIDOpenIDRealm";
     GIDAuthFlow *handlerAuthFlow = weakAuthFlow;
     OIDAuthState *authState = handlerAuthFlow.authState;
     if (authState && !handlerAuthFlow.error) {
-      if (![self saveAuthState:authState]) {
+      if (![self->_keychainHandler saveAuthState:authState]) {
         handlerAuthFlow.error = [self errorWithString:kKeychainError
                                                  code:kGIDSignInErrorCodeKeychain];
         return;
@@ -971,26 +982,6 @@ static NSString *const kConfigOpenIDRealmKey = @"GIDOpenIDRealm";
   return YES;
 }
 
-- (void)removeAllKeychainEntries {
-  [GTMAppAuthFetcherAuthorization removeAuthorizationFromKeychainForName:kGTMAppAuthKeychainName
-                                               useDataProtectionKeychain:YES];
-}
-
-- (BOOL)saveAuthState:(OIDAuthState *)authState {
-  GTMAppAuthFetcherAuthorization *authorization =
-      [[GTMAppAuthFetcherAuthorization alloc] initWithAuthState:authState];
-  return [GTMAppAuthFetcherAuthorization saveAuthorization:authorization
-                                         toKeychainForName:kGTMAppAuthKeychainName
-                                 useDataProtectionKeychain:YES];
-}
-
-- (OIDAuthState *)loadAuthState {
-  GTMAppAuthFetcherAuthorization *authorization =
-      [GTMAppAuthFetcherAuthorization authorizationFromKeychainForName:kGTMAppAuthKeychainName
-                                             useDataProtectionKeychain:YES];
-  return authorization.authState;
-}
-
 // Generates user profile from OIDIDToken.
 - (GIDProfileData *)profileDataWithIDToken:(OIDIDToken *)idToken {
   if (!idToken ||

GoogleSignIn/Sources/GIDSignIn_Private.h

@@ -29,6 +29,8 @@ NS_ASSUME_NONNULL_BEGIN
 @class GIDGoogleUser;
 @class GIDSignInInternalOptions;
 
+@protocol GIDKeychainHandler;
+
 /// Represents a completion block that takes a `GIDUserAuth` on success or an error if the operation
 /// was unsuccessful.
 typedef void (^GIDUserAuthCompletion)(GIDUserAuth *_Nullable userAuth, NSError *_Nullable error);
@@ -45,6 +47,10 @@ typedef void (^GIDDisconnectCompletion)(NSError *_Nullable error);
 /// Private initializer for |GIDSignIn|.
 - (instancetype)initPrivate;
 
+/// The designated initializer.
+- (instancetype)initWithKeychainHandler:(id<GIDKeychainHandler>)keychainHandler
+    NS_DESIGNATED_INITIALIZER;
+
 /// Authenticates with extra options.
 - (void)signInWithOptions:(GIDSignInInternalOptions *)options;
 

GoogleSignIn/Tests/Unit/GIDKeychainHandlerTest.m

@@ -0,0 +1,81 @@
+// Copyright 2022 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#import "GoogleSignIn/Sources/GIDKeychainHandler/Implementations/GIDKeychainHandler.h"
+
+#import <XCTest/XCTest.h>
+
+#import "GoogleSignIn/Tests/Unit/OIDAuthState+Testing.h"
+#import "GoogleSignIn/Tests/Unit/OIDTokenResponse+Testing.h"
+
+#ifdef SWIFT_PACKAGE
+@import AppAuth;
+@import GTMAppAuth;
+@import OCMock;
+#else
+#import <AppAuth/AppAuth.h>
+#import <GTMAppAuth/GTMAppAuth.h>
+#import <OCMock/OCMock.h>
+#endif
+
+static NSTimeInterval const kIDTokenExpiresIn = 100;
+
+@interface GIDKeychainHandlerTest : XCTestCase {
+  GIDKeychainHandler *_keychainHandler;
+  id _authorization;
+}
+
+@end
+
+@implementation GIDKeychainHandlerTest
+
+- (void)setUp {
+  _keychainHandler = [[GIDKeychainHandler alloc] init];
+  _authorization = OCMStrictClassMock([GTMAppAuthFetcherAuthorization class]);
+}
+
+- (void)testLoadAuthState {
+  [_keychainHandler loadAuthState];
+  [[_authorization verify] authorizationFromKeychainForName:[OCMArg any]
+                                  useDataProtectionKeychain:YES];
+}
+
+- (void)testSaveAuthState {
+  NSString *idToken = [self idTokenWithExpiresIn:kIDTokenExpiresIn];
+  OIDAuthState *authState = [OIDAuthState testInstanceWithIDToken:idToken
+                                                      accessToken:kAccessToken
+                                             accessTokenExpiresIn:kAccessTokenExpiresIn
+                                                     refreshToken:kRefreshToken];
+
+  [_keychainHandler saveAuthState:authState];
+  [[_authorization verify] saveAuthorization:[OCMArg any]
+                           toKeychainForName:[OCMArg any]
+                   useDataProtectionKeychain:YES];
+}
+
+- (void)testRemoveAllKeychainEntries {
+  [_keychainHandler removeAllKeychainEntries];
+  [[_authorization verify] removeAuthorizationFromKeychainForName:[OCMArg any]
+                                        useDataProtectionKeychain:YES];
+}
+
+#pragma mark - Helpers
+
+- (NSString *)idTokenWithExpiresIn:(NSTimeInterval)expiresIn {
+  // The expireTime should be based on 1970.
+  NSTimeInterval expireTime = [[NSDate date] timeIntervalSince1970] + expiresIn;
+  return [OIDTokenResponse idTokenWithSub:kUserID exp:@(expireTime)];
+}
+
+@end

GoogleSignIn/Tests/Unit/GIDSignInTest.m

@@ -30,6 +30,7 @@
 #import "GoogleSignIn/Sources/GIDGoogleUser_Private.h"
 #import "GoogleSignIn/Sources/GIDSignIn_Private.h"
 #import "GoogleSignIn/Sources/GIDSignInPreferences.h"
+#import "GoogleSignIn/Sources/GIDKeychainHandler/Implementations/Fakes/GIDFakeKeychainHandler.h"
 
 #if TARGET_OS_IOS && !TARGET_OS_MACCATALYST
 #import "GoogleSignIn/Sources/GIDEMMErrorHandler.h"
@@ -192,6 +193,8 @@ static NSString *const kNewScope = @"newScope";
 
   // Mock |GTMAppAuthFetcherAuthorization|.
   id _authorization;
+  
+  GIDFakeKeychainHandler *_keychainHandler;
 
 #if TARGET_OS_IOS || TARGET_OS_MACCATALYST
   // Mock |UIViewController|.
@@ -219,12 +222,6 @@ static NSString *const kNewScope = @"newScope";
   // Fake [NSBundle mainBundle];
   GIDFakeMainBundle *_fakeMainBundle;
 
-  // Whether |saveParamsToKeychainForName:authentication:| has been called.
-  BOOL _keychainSaved;
-
-  // Whether |removeAuthFromKeychainForName:| has been called.
-  BOOL _keychainRemoved;
-
   // The |GIDSignIn| object being tested.
   GIDSignIn *_signIn;
 
@@ -256,9 +253,6 @@ static NSString *const kNewScope = @"newScope";
 
   // The saved token request callback.
   OIDTokenCallback _savedTokenCallback;
-
-  // Status returned by saveAuthorization:toKeychainForName:
-  BOOL _saveAuthorizationReturnValue;
 }
 @end
 
@@ -271,12 +265,9 @@ static NSString *const kNewScope = @"newScope";
 #if TARGET_OS_IOS && !TARGET_OS_MACCATALYST
   _isEligibleForEMM = [UIDevice currentDevice].systemVersion.integerValue >= 9;
 #endif // TARGET_OS_IOS && !TARGET_OS_MACCATALYST
-  _saveAuthorizationReturnValue = YES;
 
   // States
   _completionCalled = NO;
-  _keychainSaved = NO;
-  _keychainRemoved = NO;
 
   // Mocks
 #if TARGET_OS_IOS || TARGET_OS_MACCATALYST
@@ -290,22 +281,8 @@ static NSString *const kNewScope = @"newScope";
   _tokenResponse = OCMStrictClassMock([OIDTokenResponse class]);
   _tokenRequest = OCMStrictClassMock([OIDTokenRequest class]);
   _authorization = OCMStrictClassMock([GTMAppAuthFetcherAuthorization class]);
-  OCMStub([_authorization authorizationFromKeychainForName:OCMOCK_ANY
-                                 useDataProtectionKeychain:YES]).andReturn(_authorization);
   OCMStub([_authorization alloc]).andReturn(_authorization);
   OCMStub([_authorization initWithAuthState:OCMOCK_ANY]).andReturn(_authorization);
-  OCMStub([_authorization saveAuthorization:OCMOCK_ANY
-                          toKeychainForName:OCMOCK_ANY
-                  useDataProtectionKeychain:YES])
-      .andDo(^(NSInvocation *invocation) {
-        self->_keychainSaved = self->_saveAuthorizationReturnValue;
-        [invocation setReturnValue:&self->_saveAuthorizationReturnValue];
-      });
-  OCMStub([_authorization removeAuthorizationFromKeychainForName:OCMOCK_ANY
-                                       useDataProtectionKeychain:YES])
-      .andDo(^(NSInvocation *invocation) {
-        self->_keychainRemoved = YES;
-      });
   _user = OCMStrictClassMock([GIDGoogleUser class]);
   _oidAuthorizationService = OCMStrictClassMock([OIDAuthorizationService class]);
   OCMStub([_oidAuthorizationService
@@ -330,7 +307,8 @@ static NSString *const kNewScope = @"newScope";
   [[NSUserDefaults standardUserDefaults] setBool:YES
                                           forKey:kAppHasRunBeforeKey];
 
-  _signIn = [[GIDSignIn alloc] initPrivate];
+  _keychainHandler = [[GIDFakeKeychainHandler alloc] init];
+  _signIn = [[GIDSignIn alloc] initWithKeychainHandler:_keychainHandler];
   _hint = nil;
 
   __weak GIDSignInTest *weakSelf = self;
@@ -349,7 +327,6 @@ static NSString *const kNewScope = @"newScope";
   OCMVerifyAll(_authState);
   OCMVerifyAll(_tokenResponse);
   OCMVerifyAll(_tokenRequest);
-  OCMVerifyAll(_authorization);
   OCMVerifyAll(_user);
   OCMVerifyAll(_oidAuthorizationService);
 
@@ -359,7 +336,6 @@ static NSString *const kNewScope = @"newScope";
   OCMVerifyAll(_presentingWindow);
 #endif // TARGET_OS_IOS || TARGET_OS_MACCATALYST
 
-
   [_fakeMainBundle stopFaking];
   [super tearDown];
 }
@@ -418,6 +394,7 @@ static NSString *const kNewScope = @"newScope";
   OCMStub([_authState lastTokenResponse]).andReturn(_tokenResponse);
   OCMStub([_authState refreshToken]).andReturn(kRefreshToken);
   [[_authState expect] setStateChangeDelegate:OCMOCK_ANY];
+  [_keychainHandler saveAuthState:_authState];
 
   id idTokenDecoded = OCMClassMock([OIDIDToken class]);
   OCMStub([idTokenDecoded alloc]).andReturn(idTokenDecoded);
@@ -438,8 +415,8 @@ static NSString *const kNewScope = @"newScope";
   
   [_signIn restorePreviousSignInNoRefresh];
 
-  [_authorization verify];
   [_authState verify];
+  [_authorization verify];
   [_tokenResponse verify];
   [_tokenRequest verify];
   [idTokenDecoded verify];
@@ -449,37 +426,31 @@ static NSString *const kNewScope = @"newScope";
 }
 
 - (void)testRestoredPreviousSignInNoRefresh_hasNoPreviousUser {
-  [[[_authorization expect] andReturn:nil] authState];
-
+  XCTAssertNil([_keychainHandler loadAuthState]);
   [_signIn restorePreviousSignInNoRefresh];
 
-  [_authorization verify];
   XCTAssertNil(_signIn.currentUser);
 }
 
 - (void)testHasPreviousSignIn_HasBeenAuthenticated {
-  [[[_authorization expect] andReturn:_authState] authState];
+  [_keychainHandler saveAuthState:_authState];
   [[[_authState expect] andReturnValue:[NSNumber numberWithBool:YES]] isAuthorized];
   XCTAssertTrue([_signIn hasPreviousSignIn], @"should return |YES|");
-  [_authorization verify];
   [_authState verify];
-  XCTAssertFalse(_keychainRemoved, @"should not remove keychain");
   XCTAssertFalse(_completionCalled, @"should not call delegate");
   XCTAssertNil(_authError, @"should have no error");
 }
 
 - (void)testHasPreviousSignIn_HasNotBeenAuthenticated {
-  [[[_authorization expect] andReturn:_authState] authState];
+  [_keychainHandler saveAuthState:_authState];
   [[[_authState expect] andReturnValue:[NSNumber numberWithBool:NO]] isAuthorized];
   XCTAssertFalse([_signIn hasPreviousSignIn], @"should return |NO|");
-  [_authorization verify];
   [_authState verify];
-  XCTAssertFalse(_keychainRemoved, @"should not remove keychain");
   XCTAssertFalse(_completionCalled, @"should not call delegate");
 }
 
 - (void)testRestorePreviousSignInWhenSignedOut {
-  [[[_authorization expect] andReturn:_authState] authState];
+  [_keychainHandler saveAuthState:_authState];
   [[[_authState expect] andReturnValue:[NSNumber numberWithBool:NO]] isAuthorized];
   _completionCalled = NO;
   _authError = nil;
@@ -499,7 +470,6 @@ static NSString *const kNewScope = @"newScope";
   }];
 
   [self waitForExpectationsWithTimeout:1 handler:nil];
-  [_authorization verify];
   [_authState verify];
 }
 
@@ -730,13 +700,13 @@ static NSString *const kNewScope = @"newScope";
                      oldAccessToken:NO
                         modalCancel:NO];
   [self waitForExpectationsWithTimeout:1 handler:nil];
-  XCTAssertFalse(_keychainSaved, @"should save to keychain");
   XCTAssertTrue(_completionCalled, @"should call delegate");
   XCTAssertEqualObjects(_authError.domain, kGIDSignInErrorDomain);
   XCTAssertEqual(_authError.code, kGIDSignInErrorCodeKeychain);
 }
 
 - (void)testSignOut {
+  XCTAssert([_keychainHandler saveAuthState:_authState]);
   // Sign in a user so that we can then sign them out.
   [self OAuthLoginWithAddScopesFlow:NO
                           authError:nil
@@ -748,25 +718,21 @@ static NSString *const kNewScope = @"newScope";
                         modalCancel:NO];
 
   XCTAssertNotNil(_signIn.currentUser);
+  XCTAssertNotNil([_keychainHandler loadAuthState]);
 
   [_signIn signOut];
   XCTAssertNil(_signIn.currentUser, @"should not have a current user");
-  XCTAssertTrue(_keychainRemoved, @"should remove keychain");
-
-  OCMVerify([_authorization removeAuthorizationFromKeychainForName:kKeychainName
-                                         useDataProtectionKeychain:YES]);
+  XCTAssertNil([_keychainHandler loadAuthState]);
 }
 
 - (void)testNotHandleWrongScheme {
   XCTAssertFalse([_signIn handleURL:[NSURL URLWithString:kWrongSchemeURL]],
                  @"should not handle URL");
-  XCTAssertFalse(_keychainSaved, @"should not save to keychain");
   XCTAssertFalse(_completionCalled, @"should not call delegate");
 }
 
 - (void)testNotHandleWrongPath {
   XCTAssertFalse([_signIn handleURL:[NSURL URLWithString:kWrongPathURL]], @"should not handle URL");
-  XCTAssertFalse(_keychainSaved, @"should not save to keychain");
   XCTAssertFalse(_completionCalled, @"should not call delegate");
 }
 
@@ -774,7 +740,7 @@ static NSString *const kNewScope = @"newScope";
 
 // Verifies disconnect calls callback with no errors if access token is present.
 - (void)testDisconnect_accessToken {
-  [[[_authorization expect] andReturn:_authState] authState];
+  [_keychainHandler saveAuthState:_authState];
   [[[_authState expect] andReturn:_tokenResponse] lastTokenResponse];
   [[[_tokenResponse expect] andReturn:kAccessToken] accessToken];
   [[[_authorization expect] andReturn:_fetcherService] fetcherService];
@@ -789,11 +755,12 @@ static NSString *const kNewScope = @"newScope";
   [_authorization verify];
   [_authState verify];
   [_tokenResponse verify];
+  XCTAssertNil([_keychainHandler loadAuthState]);
 }
 
 // Verifies disconnect if access token is present.
 - (void)testDisconnectNoCallback_accessToken {
-  [[[_authorization expect] andReturn:_authState] authState];
+  [_keychainHandler saveAuthState:_authState];
   [[[_authState expect] andReturn:_tokenResponse] lastTokenResponse];
   [[[_tokenResponse expect] andReturn:kAccessToken] accessToken];
   [[[_authorization expect] andReturn:_fetcherService] fetcherService];
@@ -802,11 +769,12 @@ static NSString *const kNewScope = @"newScope";
   [_authorization verify];
   [_authState verify];
   [_tokenResponse verify];
+  XCTAssertNil([_keychainHandler loadAuthState]);
 }
 
 // Verifies disconnect calls callback with no errors if refresh token is present.
 - (void)testDisconnect_refreshToken {
-  [[[_authorization expect] andReturn:_authState] authState];
+  [_keychainHandler saveAuthState:_authState];
   [[[_authState expect] andReturn:_tokenResponse] lastTokenResponse];
   [[[_tokenResponse expect] andReturn:nil] accessToken];
   [[[_authState expect] andReturn:_tokenResponse] lastTokenResponse];
@@ -822,12 +790,12 @@ static NSString *const kNewScope = @"newScope";
   [self verifyAndRevokeToken:kRefreshToken hasCallback:YES];
   [_authorization verify];
   [_authState verify];
-  [_tokenResponse verify];
+  XCTAssertNil([_keychainHandler loadAuthState]);
 }
 
 // Verifies disconnect errors are passed along to the callback.
 - (void)testDisconnect_errors {
-  [[[_authorization expect] andReturn:_authState] authState];
+  [_keychainHandler saveAuthState:_authState];
   [[[_authState expect] andReturn:_tokenResponse] lastTokenResponse];
   [[[_tokenResponse expect] andReturn:kAccessToken] accessToken];
   [[[_authorization expect] andReturn:_fetcherService] fetcherService];
@@ -846,11 +814,12 @@ static NSString *const kNewScope = @"newScope";
   [_authorization verify];
   [_authState verify];
   [_tokenResponse verify];
+  XCTAssertNotNil([_keychainHandler loadAuthState]);
 }
 
 // Verifies disconnect with errors
 - (void)testDisconnectNoCallback_errors {
-  [[[_authorization expect] andReturn:_authState] authState];
+  [_keychainHandler saveAuthState:_authState];
   [[[_authState expect] andReturn:_tokenResponse] lastTokenResponse];
   [[[_tokenResponse expect] andReturn:kAccessToken] accessToken];
   [[[_authorization expect] andReturn:_fetcherService] fetcherService];
@@ -862,12 +831,12 @@ static NSString *const kNewScope = @"newScope";
   [_authorization verify];
   [_authState verify];
   [_tokenResponse verify];
+  XCTAssertNotNil([_keychainHandler loadAuthState]);
 }
 
-
 // Verifies disconnect calls callback with no errors and clears keychain if no tokens are present.
 - (void)testDisconnect_noTokens {
-  [[[_authorization expect] andReturn:_authState] authState];
+  [_keychainHandler saveAuthState:_authState];
   [[[_authState expect] andReturn:_tokenResponse] lastTokenResponse];
   [[[_tokenResponse expect] andReturn:nil] accessToken];
   [[[_authState expect] andReturn:_tokenResponse] lastTokenResponse];
@@ -881,25 +850,25 @@ static NSString *const kNewScope = @"newScope";
   }];
   [self waitForExpectationsWithTimeout:1 handler:nil];
   XCTAssertFalse([self isFetcherStarted], @"should not fetch");
-  XCTAssertTrue(_keychainRemoved, @"keychain should be removed");
   [_authorization verify];
   [_authState verify];
   [_tokenResponse verify];
+  XCTAssertNil([_keychainHandler loadAuthState]);
 }
 
 // Verifies disconnect clears keychain if no tokens are present.
 - (void)testDisconnectNoCallback_noTokens {
-  [[[_authorization expect] andReturn:_authState] authState];
+  [_keychainHandler saveAuthState:_authState];
   [[[_authState expect] andReturn:_tokenResponse] lastTokenResponse];
   [[[_tokenResponse expect] andReturn:nil] accessToken];
   [[[_authState expect] andReturn:_tokenResponse] lastTokenResponse];
   [[[_tokenResponse expect] andReturn:nil] refreshToken];
   [_signIn disconnectWithCompletion:nil];
   XCTAssertFalse([self isFetcherStarted], @"should not fetch");
-  XCTAssertTrue(_keychainRemoved, @"keychain should be removed");
   [_authorization verify];
   [_authState verify];
   [_tokenResponse verify];
+  XCTAssertNil([_keychainHandler loadAuthState]);
 }
 
 - (void)testPresentingViewControllerException {
@@ -1066,7 +1035,6 @@ static NSString *const kNewScope = @"newScope";
 
   [self waitForExpectationsWithTimeout:1 handler:nil];
 
-  XCTAssertFalse(_keychainSaved, @"should not save to keychain");
   XCTAssertTrue(_completionCalled, @"should call delegate");
   XCTAssertNotNil(_authError, @"should have error");
   XCTAssertEqualObjects(_authError.domain, kGIDSignInErrorDomain);
@@ -1106,7 +1074,6 @@ static NSString *const kNewScope = @"newScope";
   [self waitForExpectationsWithTimeout:1 handler:nil];
 
   [emmSupport verify];
-  XCTAssertFalse(_keychainSaved, @"should not save to keychain");
   XCTAssertTrue(_completionCalled, @"should call delegate");
   XCTAssertNotNil(_authError, @"should have error");
   XCTAssertEqualObjects(_authError.domain, kGIDSignInErrorDomain);
@@ -1167,7 +1134,6 @@ static NSString *const kNewScope = @"newScope";
   if (hasCallback) {
     [self waitForExpectationsWithTimeout:1 handler:nil];
   }
-  XCTAssertTrue(_keychainRemoved, @"should clear saved keychain name");
 }
 
 - (void)OAuthLoginWithAddScopesFlow:(BOOL)addScopesFlow
@@ -1203,7 +1169,7 @@ static NSString *const kNewScope = @"newScope";
                    additionalScopes:(NSArray *)additionalScopes {
   if (restoredSignIn) {
     // clearAndAuthenticateWithOptions
-    [[[_authorization expect] andReturn:_authState] authState];
+    [_keychainHandler saveAuthState:_authState];
     BOOL isAuthorized = restoredSignIn ? YES : NO;
     [[[_authState expect] andReturnValue:[NSNumber numberWithBool:isAuthorized]] isAuthorized];
   }
@@ -1290,7 +1256,6 @@ static NSString *const kNewScope = @"newScope";
       }
     }
 
-    [_authorization verify];
     [_authState verify];
 
     XCTAssertNotNil(_savedAuthorizationRequest);
@@ -1369,7 +1334,7 @@ static NSString *const kNewScope = @"newScope";
   __block GIDProfileData *profileData;
 
   if (keychainError) {
-    _saveAuthorizationReturnValue = NO;
+    _keychainHandler.failToSave = YES;
   } else {
     if (addScopesFlow) {
       [[[_authState expect] andReturn:authResponse] lastAuthorizationResponse];
@@ -1408,7 +1373,6 @@ static NSString *const kNewScope = @"newScope";
   
   [_authState verify];
   
-  XCTAssertTrue(_keychainSaved, @"should save to keychain");
   if (addScopesFlow) {
     XCTAssertNotNil(updatedTokenResponse);
     XCTAssertNotNil(updatedAuthorizationResponse);
@@ -1423,8 +1387,6 @@ static NSString *const kNewScope = @"newScope";
 
   // If attempt to authenticate again, will reuse existing auth object.
   _completionCalled = NO;
-  _keychainRemoved = NO;
-  _keychainSaved = NO;
   _authError = nil;
 
   __block GIDGoogleUserCompletion completion;
@@ -1441,16 +1403,6 @@ static NSString *const kNewScope = @"newScope";
   completion(_user, nil);
 
   [self waitForExpectationsWithTimeout:1 handler:nil];
-  XCTAssertFalse(_keychainRemoved, @"should not remove keychain");
-  XCTAssertFalse(_keychainSaved, @"should not save to keychain again");
-  
-  if (restoredSignIn) {
-    OCMVerify([_authorization authorizationFromKeychainForName:kKeychainName
-                                     useDataProtectionKeychain:YES]);
-    OCMVerify([_authorization saveAuthorization:OCMOCK_ANY
-                              toKeychainForName:kKeychainName
-                      useDataProtectionKeychain:YES]);
-  }
 }
 
 @end