Refactor the App Check internal call logic used by other products (#1215)

* Refactor AppCheck internal call logic

* Formatting

* Update app_check_desktop.h

* Just use the internal class directly

* Update app_check.h

* Update app_check_desktop.cc

* Update app_check_desktop.cc

Author: a-maurice

app_check/src/common/common.h

@@ -22,6 +22,7 @@ namespace internal {
 // Used by App Check functions that return a future
 enum AppCheckFn {
   kAppCheckFnGetAppCheckToken = 0,
+  kAppCheckFnGetAppCheckStringInternal,
   kAppCheckFnCount,
 };
 

app_check/src/desktop/app_check_desktop.cc

@@ -28,7 +28,10 @@ namespace internal {
 static AppCheckProviderFactory* g_provider_factory = nullptr;
 
 AppCheckInternal::AppCheckInternal(App* app)
-    : app_(app), cached_token_(), cached_provider_() {
+    : app_(app),
+      cached_token_(),
+      cached_provider_(),
+      is_token_auto_refresh_enabled_(true) {
   future_manager().AllocFutureApi(this, kAppCheckFnCount);
   InitRegistryCalls();
 }
@@ -76,7 +79,9 @@ void AppCheckInternal::SetAppCheckProviderFactory(
 }
 
 void AppCheckInternal::SetTokenAutoRefreshEnabled(
-    bool is_token_auto_refresh_enabled) {}
+    bool is_token_auto_refresh_enabled) {
+  is_token_auto_refresh_enabled_ = is_token_auto_refresh_enabled;
+}
 
 Future<AppCheckToken> AppCheckInternal::GetAppCheckToken(bool force_refresh) {
   auto handle = future()->SafeAlloc<AppCheckToken>(kAppCheckFnGetAppCheckToken);
@@ -112,6 +117,42 @@ Future<AppCheckToken> AppCheckInternal::GetAppCheckTokenLastResult() {
       future()->LastResult(kAppCheckFnGetAppCheckToken));
 }
 
+Future<std::string> AppCheckInternal::GetAppCheckTokenStringInternal() {
+  auto handle =
+      future()->SafeAlloc<std::string>(kAppCheckFnGetAppCheckStringInternal);
+  if (HasValidCacheToken()) {
+    future()->CompleteWithResult(handle, 0, cached_token_.token);
+  } else if (is_token_auto_refresh_enabled_) {
+    // Only refresh the token if it is enabled
+    AppCheckProvider* provider = GetProvider();
+    if (provider != nullptr) {
+      // Get a new token, and pass the result into the future.
+      // Note that this is slightly different from the one above, as the
+      // Future result is just the string token, and not the full struct.
+      auto token_callback{
+          [this, handle](firebase::app_check::AppCheckToken token,
+                         int error_code, const std::string& error_message) {
+            if (error_code == firebase::app_check::kAppCheckErrorNone) {
+              UpdateCachedToken(token);
+              future()->CompleteWithResult(handle, 0, token.token);
+            } else {
+              future()->Complete(handle, error_code, error_message.c_str());
+            }
+          }};
+      provider->GetToken(token_callback);
+    } else {
+      future()->Complete(
+          handle, firebase::app_check::kAppCheckErrorInvalidConfiguration,
+          "No AppCheckProvider installed.");
+    }
+  } else {
+    future()->Complete(
+        handle, kAppCheckErrorUnknown,
+        "No AppCheck token available, and auto refresh is disabled");
+  }
+  return MakeFuture(future(), handle);
+}
+
 void AppCheckInternal::AddAppCheckListener(AppCheckListener* listener) {
   if (listener) {
     token_listeners_.push_back(listener);
@@ -153,17 +194,14 @@ void AppCheckInternal::CleanupRegistryCalls() {
 bool AppCheckInternal::GetAppCheckTokenAsyncForRegistry(App* app,
                                                         void* /*unused*/,
                                                         void* out) {
-  Future<AppCheckToken>* out_future = static_cast<Future<AppCheckToken>*>(out);
+  Future<std::string>* out_future = static_cast<Future<std::string>*>(out);
   if (!app || !out_future) {
     return false;
   }
 
   AppCheck* app_check = AppCheck::GetInstance(app);
-  if (app_check) {
-    // TODO(amaurice): This should call some internal function instead of the
-    // public one, since this will change the *LastResult value behind the
-    // scenes.
-    *out_future = app_check->GetAppCheckToken(false);
+  if (app_check && app_check->internal_) {
+    *out_future = app_check->internal_->GetAppCheckTokenStringInternal();
     return true;
   }
   return false;

app_check/src/desktop/app_check_desktop.h

@@ -16,6 +16,7 @@
 #define FIREBASE_APP_CHECK_SRC_DESKTOP_APP_CHECK_DESKTOP_H_
 
 #include <list>
+#include <string>
 
 #include "app/src/future_manager.h"
 #include "app/src/include/firebase/app.h"
@@ -42,6 +43,10 @@ class AppCheckInternal {
 
   Future<AppCheckToken> GetAppCheckTokenLastResult();
 
+  // Gets the App Check token as just the string, to be used by
+  // internal methods to not conflict with the publicly returned future.
+  Future<std::string> GetAppCheckTokenStringInternal();
+
   void AddAppCheckListener(AppCheckListener* listener);
 
   void RemoveAppCheckListener(AppCheckListener* listener);
@@ -82,6 +87,9 @@ class AppCheckInternal {
   AppCheckToken cached_token_;
   // List of registered listeners for Token changes.
   std::list<AppCheckListener*> token_listeners_;
+  // Should it automatically get an App Check token if there is not a valid
+  // cached token.
+  bool is_token_auto_refresh_enabled_;
 };
 
 }  // namespace internal

app_check/src/include/firebase/app_check.h

@@ -157,6 +157,8 @@ class AppCheck {
 
   void DeleteInternal();
 
+  // Make the Internal version a friend class, so that it can access itself.
+  friend class internal::AppCheckInternal;
   internal::AppCheckInternal* internal_;
 };
 

storage/src/desktop/storage_reference_desktop.cc

@@ -29,7 +29,6 @@
 #include "app/src/function_registry.h"
 #include "app/src/include/firebase/app.h"
 #include "app/src/thread.h"
-#include "app_check/src/include/firebase/app_check.h"
 #include "storage/src/common/common_internal.h"
 #include "storage/src/desktop/controller_desktop.h"
 #include "storage/src/desktop/metadata_desktop.h"
@@ -259,15 +258,14 @@ void StorageReferenceInternal::PrepareRequestBlocking(
                       storage_->user_agent().c_str());
 
   // Use the function registry to get the App Check token.
-  Future<::firebase::app_check::AppCheckToken> app_check_future;
+  Future<std::string> app_check_future;
   bool succeeded = storage_->app()->function_registry()->CallFunction(
       ::firebase::internal::FnAppCheckGetTokenAsync, storage_->app(), nullptr,
       &app_check_future);
   if (succeeded && app_check_future.status() != kFutureStatusInvalid) {
-    const ::firebase::app_check::AppCheckToken* token =
-        app_check_future.Await(kAppCheckTokenTimeoutMs);
+    const std::string* token = app_check_future.Await(kAppCheckTokenTimeoutMs);
     if (token) {
-      request->add_header("X-Firebase-AppCheck", token->token.c_str());
+      request->add_header("X-Firebase-AppCheck", token->c_str());
     }
   }
 }