[Auth] protect against auth deletion before schedule the ensure token task.

There are two bugs to fix here, both are race conditions while creating / destroying auth in very short time.

The fix includes the following changes:

1. Previous Auth destruction is waiting for all async operation to complete, which would easily cause hang during destruction. In this change actually cancels all on going async operations, which live in two different schedulers.
One scheduler is the auth's own scheduler, that can be cleaned up for no consequence.
The other is persistence scheduler, which is globally shared among all FirebaseApp instances.

2. To clean up async ops in persistence scheduler, we restore the op handler in UserSecureManager, mapped by operation type. We make sure each operation type should only have one valid task. In destruction we can rely on active handlers to cancel all operations.

3. Even with cancelling the operation, the already running tasks will still finish, and thus are going to trigger the callbacks, which would still cause issue during destruction. To avoid this issue we introduced destructing flag for auth instance, and callbacks will abort if already destructing.

PiperOrigin-RevId: 261781967

Author: cynthiajiang

app/src/secure/user_secure_manager.cc

@@ -78,6 +78,7 @@ UserSecureManager::UserSecureManager(
 }
 
 UserSecureManager::~UserSecureManager() {
+  CancelScheduledTasks();
   // Clear safe reference immediately so that scheduled callback can skip
   // executing code which requires reference to this.
   safe_this_.ClearReference();
@@ -115,7 +116,9 @@ Future<std::string> UserSecureManager::LoadUserData(
         }
       },
       safe_this_, data_handle, user_secure_.get());
-  s_scheduler_->Schedule(callback);
+
+  CancelOperation(kLoadUserData);
+  operation_handles_[kLoadUserData] = s_scheduler_->Schedule(callback);
   return MakeFuture(&future_api_, future_handle);
 }
 
@@ -137,7 +140,9 @@ Future<void> UserSecureManager::SaveUserData(const std::string& app_name,
         }
       },
       safe_this_, data_handle, user_secure_.get());
-  s_scheduler_->Schedule(callback);
+
+  CancelOperation(kSaveUserData);
+  operation_handles_[kSaveUserData] = s_scheduler_->Schedule(callback);
   return MakeFuture(&future_api_, future_handle);
 }
 
@@ -158,7 +163,9 @@ Future<void> UserSecureManager::DeleteUserData(const std::string& app_name) {
         }
       },
       safe_this_, data_handle, user_secure_.get());
-  s_scheduler_->Schedule(callback);
+
+  CancelOperation(kDeleteUserData);
+  operation_handles_[kDeleteUserData] = s_scheduler_->Schedule(callback);
   return MakeFuture(&future_api_, future_handle);
 }
 
@@ -179,7 +186,9 @@ Future<void> UserSecureManager::DeleteAllData() {
         }
       },
       safe_this_, data_handle, user_secure_.get());
-  s_scheduler_->Schedule(callback);
+
+  CancelOperation(kDeleteAllData);
+  operation_handles_[kDeleteAllData] = s_scheduler_->Schedule(callback);
   return MakeFuture(&future_api_, future_handle);
 }
 
@@ -278,6 +287,23 @@ void UserSecureManager::BinaryToAscii(const std::string& original,
   }
 }
 
+void UserSecureManager::CancelScheduledTasks() {
+  for (auto it = operation_handles_.begin(); it != operation_handles_.end();
+       ++it) {
+    it->second.Cancel();
+  }
+  operation_handles_.clear();
+}
+
+void UserSecureManager::CancelOperation(SecureOperationType operation_type) {
+  // Cancel and remove existing handle.
+  auto it = operation_handles_.find(operation_type);
+  if (it != operation_handles_.end()) {
+    it->second.Cancel();
+    operation_handles_.erase(it);
+  }
+}
+
 }  // namespace secure
 }  // namespace app
 }  // namespace firebase

app/src/secure/user_secure_manager.h

@@ -27,6 +27,13 @@ namespace firebase {
 namespace app {
 namespace secure {
 
+enum SecureOperationType {
+  kLoadUserData,
+  kSaveUserData,
+  kDeleteUserData,
+  kDeleteAllData,
+};
+
 class UserSecureManager {
  public:
   explicit UserSecureManager(const char* domain, const char* app_id);
@@ -55,6 +62,11 @@ class UserSecureManager {
   static void BinaryToAscii(const std::string& original, std::string* encoded);
 
  private:
+  // Cancel already scheduled tasks in destruction.
+  void CancelScheduledTasks();
+
+  void CancelOperation(SecureOperationType operation_type);
+
   UniquePtr<UserSecureInternal> user_secure_;
   ReferenceCountedFutureImpl future_api_;
 
@@ -66,6 +78,10 @@ class UserSecureManager {
   static scheduler::Scheduler* s_scheduler_;
   static int32_t s_scheduler_ref_count_;
 
+  // map from operation type to scheduled request handle. Make sure only one
+  // request exist in scheduler for each type.
+  std::map<SecureOperationType, scheduler::RequestHandle> operation_handles_;
+
   // Safe reference to this.  Set in constructor and cleared in destructor
   // Should be safe to be copied in any thread because the SharedPtr never
   // changes, until safe_this_ is completely destroyed.

auth/src/auth.cc

@@ -140,6 +140,11 @@ void Auth::DeleteInternal() {
 
   if (!auth_data_) return;
 
+  {
+    MutexLock destructing_lock(auth_data_->desctruting_mutex);
+    auth_data_->destructing = true;
+  }
+
   CleanupNotifier* notifier = CleanupNotifier::FindByOwner(auth_data_->app);
   assert(notifier);
   notifier->UnregisterObject(this);

auth/src/data.h

@@ -83,7 +83,8 @@ struct AuthData {
         listener_impl(nullptr),
         id_token_listener_impl(nullptr),
         expect_id_token_listener_callback(false),
-        persistent_cache_load_pending(true) {}
+        persistent_cache_load_pending(true),
+        destructing(false) {}
 
   ~AuthData() {
     ClearUserInfos(this);
@@ -174,6 +175,12 @@ struct AuthData {
   // Mutex protecting `expect_id_token_listener_callback`
   Mutex expect_id_token_mutex;
 
+  // Tracks if auth is being destroyed.
+  bool destructing;
+
+  // Mutex protecting destructing
+  Mutex desctruting_mutex;
+
   // Sets if the Id Token Listener is expecting a callback.
   // Used to workaround an issue where the Id Token is not reset with a new one,
   // and thus not triggered correctly.

auth/src/desktop/auth_desktop.cc

@@ -222,7 +222,8 @@ void Auth::InitPlatformAuth(AuthData* const auth_data) {
 
 void Auth::DestroyPlatformAuth(AuthData* const auth_data) {
   FIREBASE_ASSERT_RETURN_VOID(auth_data);
-  WaitForAllAsyncToComplete(auth_data->auth_impl);
+  auto auth_impl = static_cast<AuthImpl*>(auth_data->auth_impl);
+  auth_impl->scheduler_.CancelAllAndShutdownWorkerThread();
   // Unregister from the function registry.
   auth_data->app->function_registry()->UnregisterFunction(
       internal::FnAuthGetCurrentToken);
@@ -235,14 +236,6 @@ void Auth::DestroyPlatformAuth(AuthData* const auth_data) {
 
   DestroyTokenRefresher(auth_data);
 
-  {
-    // Acquire listeners' mutex to avoid race condition if another thread is
-    // about to notify listeners.
-    MutexLock lock(auth_data->listeners_mutex);
-    auth_data->listeners.clear();
-    auth_data->id_token_listeners.clear();
-  }
-
   DestroyUserDataPersist(auth_data);
 
   UserView::ClearUser(auth_data);
@@ -493,6 +486,11 @@ void DestroyUserDataPersist(AuthData* auth_data) {
 }
 
 void LoadFinishTriggerListeners(AuthData* auth_data) {
+  MutexLock destructing_lock(auth_data->desctruting_mutex);
+  if (auth_data->destructing) {
+    // If auth is destructing, abort.
+    return;
+  }
   // We would have to block other listener changes to protect race condition
   // on how many times a listener should be triggered. We would rely on first
   // listener trigger to flip the persistence loading bit.

auth/src/desktop/auth_desktop.h

@@ -122,7 +122,7 @@ struct AuthCompletionHandle {
 
 // The desktop-specific Auth implementation.
 struct AuthImpl {
-  AuthImpl() : async_sem(0), active_async_calls(0) {}
+  AuthImpl() {}
 
   // The application's API key.
   std::string api_key;
@@ -135,10 +135,6 @@ struct AuthImpl {
   IdTokenRefreshThread token_refresh_thread;
   // Instance responsible for user data persistence.
   UniquePtr<UserDataPersist> user_data_persist;
-  // Synchronization primitives used for managing threads spawned by CallAsync.
-  Semaphore async_sem;
-  Mutex async_mutex;
-  int active_async_calls;
 
   // Serializes all REST call from this object.
   scheduler::Scheduler scheduler_;

auth/src/desktop/auth_util.cc

@@ -47,35 +47,5 @@ void CompletePromise(Promise<void>* const promise,
   promise->Complete();
 }
 
-// Utility functions for making sure that the CallAsync functions end before
-// we destruct, etc.
-void StartAsyncFunction(void* auth_impl_void) {
-  auto auth_impl = static_cast<AuthImpl*>(auth_impl_void);
-  MutexLock lock(auth_impl->async_mutex);
-  auth_impl->active_async_calls++;
-}
-
-void EndAsyncFunction(void* auth_impl_void) {
-  auto auth_impl = static_cast<AuthImpl*>(auth_impl_void);
-  {
-    MutexLock lock(auth_impl->async_mutex);
-    auth_impl->active_async_calls--;
-    auth_impl->async_sem.Post();
-  }
-}
-
-void WaitForAllAsyncToComplete(void* auth_impl_void) {
-  auto auth_impl = static_cast<AuthImpl*>(auth_impl_void);
-  for (;;) {
-    bool transfers_complete = false;
-    {
-      MutexLock lock(auth_impl->async_mutex);
-      transfers_complete = auth_impl->active_async_calls == 0;
-    }
-    if (transfers_complete) break;
-    auth_impl->async_sem.Wait();
-  }
-}
-
 }  // namespace auth
 }  // namespace firebase

auth/src/desktop/auth_util.h

@@ -73,12 +73,6 @@ Future<ResultT> CallAsync(
     std::unique_ptr<RequestT> request,
     const typename AuthDataHandle<ResultT, RequestT>::CallbackT callback);
 
-// Utility functions for making sure that the CallAsync functions end before
-// we destruct, etc.
-void StartAsyncFunction(void* auth_impl_void);
-void EndAsyncFunction(void* auth_impl_void);
-void WaitForAllAsyncToComplete(void* auth_impl_void);
-
 // Sends the given request on the network and returns the response. The response
 // is cast to the specified T without any checks, so it's the caller's
 // responsibility to ensure the correct type is given.
@@ -100,14 +94,12 @@ inline Future<ResultT> CallAsync(
   // doesn't need to access the request anyway.
   FIREBASE_ASSERT_RETURN(Future<ResultT>(), auth_data && callback);
 
-  StartAsyncFunction(auth_data->auth_impl);
   typedef AuthDataHandle<ResultT, RequestT> HandleT;
 
   auto scheduler_callback = NewCallback(
       [](HandleT* const raw_auth_data_handle) {
         std::unique_ptr<HandleT> handle(raw_auth_data_handle);
         handle->callback(handle.get());
-        EndAsyncFunction(handle->auth_data->auth_impl);
       },
       new HandleT(auth_data, promise, std::move(request), callback));
   auto auth_impl = static_cast<AuthImpl*>(auth_data->auth_impl);

auth/src/desktop/user_desktop.cc

@@ -186,7 +186,6 @@ Future<ResultT> CallAsyncWithFreshToken(
     std::unique_ptr<RequestT> request,
     const typename AuthDataHandle<ResultT, RequestT>::CallbackT callback) {
   FIREBASE_ASSERT_RETURN(Future<ResultT>(), auth_data && request && callback);
-  StartAsyncFunction(auth_data->auth_impl);
 
   typedef AuthDataHandle<ResultT, RequestT> HandleT;
 
@@ -198,13 +197,11 @@ Future<ResultT> CallAsyncWithFreshToken(
             EnsureFreshToken(handle->auth_data, false);
         if (!get_token_result.IsValid()) {
           FailPromise(&handle->promise, get_token_result.error());
-          EndAsyncFunction(handle->auth_data->auth_impl);
           return;
         }
         handle->request->SetIdToken(get_token_result.token().c_str());
 
         handle->callback(handle.get());
-        EndAsyncFunction(handle->auth_data->auth_impl);
       },
       new HandleT(auth_data, promise, std::move(request), callback));
   auto auth_impl = static_cast<AuthImpl*>(auth_data->auth_impl);
@@ -482,38 +479,31 @@ void AssignLoadedData(const Future<std::string>& future, AuthData* auth_data) {
 
 void HandleLoadedData(const Future<std::string>& future, void* auth_data) {
   auto cast_auth_data = static_cast<AuthData*>(auth_data);
+  MutexLock destructing_lock(cast_auth_data->desctruting_mutex);
+  if (cast_auth_data->destructing) {
+    // If auth is destructing, abort.
+    return;
+  }
   AssignLoadedData(future, cast_auth_data);
-
-  // End async call for LoadUserData
-  EndAsyncFunction(cast_auth_data->auth_impl);
-
   auto scheduler_callback = NewCallback(
       [](AuthData* callback_auth_data) {
         // Don't trigger token listeners if token get refreshed, since
         // it will get triggered inside LoadFinishTriggerListeners anyways.
         const GetTokenResult get_token_result =
             EnsureFreshToken(callback_auth_data, false, false);
         LoadFinishTriggerListeners(callback_auth_data);
-        EndAsyncFunction(callback_auth_data->auth_impl);
       },
       cast_auth_data);
-  // Start Async call for EnsureFreshToken
-  StartAsyncFunction(cast_auth_data->auth_impl);
+
   auto auth_impl = static_cast<AuthImpl*>(cast_auth_data->auth_impl);
   auth_impl->scheduler_.Schedule(scheduler_callback);
 }
 
-void HandlePersistenceComplete(const Future<void>& future, void* auth_data) {
-  auto cast_auth_data = static_cast<AuthData*>(auth_data);
-  EndAsyncFunction(cast_auth_data->auth_impl);
-}
-
 Future<std::string> UserDataPersist::LoadUserData(AuthData* auth_data) {
   if (auth_data == nullptr) {
     return Future<std::string>();
   }
 
-  StartAsyncFunction(auth_data->auth_impl);
   Future<std::string> future =
       user_secure_manager_->LoadUserData(auth_data->app->name());
   future.OnCompletion(HandleLoadedData, auth_data);
@@ -565,11 +555,7 @@ Future<void> UserDataPersist::SaveUserData(AuthData* auth_data) {
 }
 
 Future<void> UserDataPersist::DeleteUserData(AuthData* auth_data) {
-  StartAsyncFunction(auth_data->auth_impl);
-  Future<void> future =
-      user_secure_manager_->DeleteUserData(auth_data->app->name());
-  future.OnCompletion(HandlePersistenceComplete, auth_data);
-  return future;
+  return user_secure_manager_->DeleteUserData(auth_data->app->name());
 }
 
 User::~User() {