Feat: Add limited use token suport

Author: AustinBenoit

app_check/src/android/app_check_android.cc

@@ -109,11 +109,17 @@ JNIEXPORT void JNICALL JniAppCheckProvider_nativeGetToken(
     JNIEnv* env, jobject j_provider, jlong c_provider,
     jobject task_completion_source);
 
+JNIEXPORT void JNICALL JniAppCheckProvider_nativeGetLimitedUseToken(
+    JNIEnv* env, jobject j_provider, jlong c_provider,
+    jobject task_completion_source);
+
 static const JNINativeMethod kNativeJniAppCheckProviderMethods[] = {
     {"nativeGetToken",
      "(JLcom/google/android/gms/tasks/TaskCompletionSource;)V",
      reinterpret_cast<void*>(JniAppCheckProvider_nativeGetToken)},
-};
+    {"nativeGetLimitedUseToken",
+     "(JLcom/google/android/gms/tasks/TaskCompletionSource;)V",
+     reinterpret_cast<void*>(JniAppCheckProvider_nativeGetLimitedUseToken)}};
 
 // clang-format off
 #define JNI_APP_CHECK_LISTENER_METHODS(X)                              \
@@ -266,6 +272,39 @@ JNIEXPORT void JNICALL JniAppCheckProvider_nativeGetToken(
   provider->GetToken(token_callback);
 }
 
+JNIEXPORT void JNICALL JniAppCheckProvider_nativeGetLimitedUseToken(
+    JNIEnv* env, jobject j_provider, jlong c_provider,
+    jobject task_completion_source) {
+  // Create GlobalReferences to the provider and task. These references will be
+  // deleted in the completion callback.
+  jobject j_provider_global = env->NewGlobalRef(j_provider);
+  jobject task_completion_source_global =
+      env->NewGlobalRef(task_completion_source);
+
+  // Defines a C++ callback method to call
+  // JniAppCheckProvider.HandleGetTokenResult with the resulting token
+  auto token_callback{[j_provider_global, task_completion_source_global](
+                          firebase::app_check::AppCheckToken token,
+                          int error_code, const std::string& error_message) {
+    // util::GetJNIEnvFromApp returns a threadsafe instance of JNIEnv.
+    JNIEnv* env = firebase::util::GetJNIEnvFromApp();
+    jstring error_string = env->NewStringUTF(error_message.c_str());
+    jstring token_string = env->NewStringUTF(token.token.c_str());
+    env->CallVoidMethod(
+        j_provider_global,
+        jni_provider::GetMethodId(jni_provider::kHandleGetTokenResult),
+        task_completion_source_global, token_string, token.expire_time_millis,
+        error_code, error_string);
+    FIREBASE_ASSERT(!util::CheckAndClearJniExceptions(env));
+    env->DeleteLocalRef(token_string);
+    env->DeleteLocalRef(error_string);
+    env->DeleteGlobalRef(j_provider_global);
+    env->DeleteGlobalRef(task_completion_source_global);
+  }};
+  AppCheckProvider* provider = reinterpret_cast<AppCheckProvider*>(c_provider);
+  provider->GetLimitedUseToken(token_callback);
+}
+
 JNIEXPORT void JNICALL JniAppCheckListener_nativeOnAppCheckTokenChanged(
     JNIEnv* env, jobject clazz, jlong c_app_check, jobject token) {
   auto app_check_internal = reinterpret_cast<AppCheckInternal*>(c_app_check);

app_check/src/common/app_check.cc

@@ -67,6 +67,15 @@ static std::map<::firebase::App*, AppCheck*>* g_app_check_map = nullptr;
 // Define the destructors for the virtual listener/provider/factory classes.
 AppCheckListener::~AppCheckListener() {}
 AppCheckProvider::~AppCheckProvider() {}
+void AppCheckProvider::GetLimitedUseToken(
+    std::function<void(AppCheckToken, int, const std::string&)>
+        completion_callback) {
+  LogWarning(
+      "A limited-use token was requested, but the custom provider did not "
+      "implement the GetLimitedUseToken method. The default implementation is "
+      "triggered as a result, and GetToken has been invoked instead.");
+  GetToken(completion_callback);
+}
 AppCheckProviderFactory::~AppCheckProviderFactory() {}
 
 namespace internal {

app_check/src/desktop/debug_provider_desktop.cc

@@ -41,7 +41,16 @@ class DebugAppCheckProvider : public AppCheckProvider {
   void GetToken(std::function<void(AppCheckToken, int, const std::string&)>
                     completion_callback) override;
 
+  void GetLimitedUseToken(
+      std::function<void(AppCheckToken, int, const std::string&)>
+          completion_callback) override;
+
  private:
+  void GetTokenInternal(
+      bool limited_use,
+      std::function<void(AppCheckToken, int, const std::string&)>
+          completion_callback);
+
   App* app_;
 
   scheduler::Scheduler scheduler_;
@@ -92,6 +101,19 @@ void GetTokenAsync(std::shared_ptr<DebugTokenRequest> request,
 void DebugAppCheckProvider::GetToken(
     std::function<void(AppCheckToken, int, const std::string&)>
         completion_callback) {
+  GetTokenInternal(false, completion_callback);
+}
+
+void DebugAppCheckProvider::GetLimitedUseToken(
+    std::function<void(AppCheckToken, int, const std::string&)>
+        completion_callback) {
+  GetTokenInternal(true, completion_callback);
+}
+
+void DebugAppCheckProvider::GetTokenInternal(
+    bool limited_use,
+    std::function<void(AppCheckToken, int, const std::string&)>
+        completion_callback) {
   // Identify the user's debug token
   const char* debug_token_cstr;
   if (!debug_token_.empty()) {
@@ -109,6 +131,7 @@ void DebugAppCheckProvider::GetToken(
   // Exchange debug token with the backend to get a proper attestation token.
   auto request = std::make_shared<DebugTokenRequest>(app_);
   request->SetDebugToken(debug_token_cstr);
+  request->SetLimitedUse(limited_use);
 
   // Use an async call, since we don't want to block on the server response.
   auto async_call =

app_check/src/desktop/debug_token_request.fbs

@@ -16,6 +16,7 @@ namespace firebase.app_check.fbs;
 
 table DebugTokenRequest {
   debug_token:string;
+  limited_use:bool;
 }
 
 root_type DebugTokenRequest;

app_check/src/desktop/debug_token_request.h

@@ -55,6 +55,11 @@ class DebugTokenRequest
     application_data_->debug_token = std::move(debug_token);
     UpdatePostFields();
   }
+
+  void SetLimitedUse(bool limited_use) {
+    application_data_->limited_use = limited_use;
+    UpdatePostFields();
+  }
 };
 
 }  // namespace internal

app_check/src/include/firebase/app_check.h

@@ -78,6 +78,16 @@ class AppCheckProvider {
   virtual void GetToken(
       std::function<void(AppCheckToken, int, const std::string&)>
           completion_callback) = 0;
+
+  /// Fetches an AppCheckToken suitable for consumption in limited-use scenarios
+  /// and then calls the provided callback function with the token or with an
+  /// error code and error message.
+  ///
+  /// If you do not implement this method, the default implementation invokes
+  /// the GetToken method whenever a limited-use token is requested.
+  virtual void GetLimitedUseToken(
+      std::function<void(AppCheckToken, int, const std::string&)>
+          completion_callback);
 };
 
 /// Interface for a factory that generates {@link AppCheckProvider}s.

app_check/src/ios/app_check_ios.mm

@@ -59,6 +59,19 @@ - (void)getTokenWithCompletion:(nonnull void (^)(FIRAppCheckToken* _Nullable,
   _cppProvider->GetToken(token_callback);
 }
 
+- (void)getLimitedUseTokenWithCompletion:(nonnull void (^)(FIRAppCheckToken* _Nullable,
+                                                           NSError* _Nullable))handler {
+  auto token_callback{[handler](firebase::app_check::AppCheckToken token, int error_code,
+                                const std::string& error_message) {
+    NSError* ios_error = firebase::app_check::internal::AppCheckErrorToNSError(
+        static_cast<firebase::app_check::AppCheckError>(error_code), error_message);
+    FIRAppCheckToken* ios_token =
+        firebase::app_check::internal::AppCheckTokenToFIRAppCheckToken(token);
+    handler(ios_token, ios_error);
+  }};
+  _cppProvider->GetLimitedUseToken(token_callback);
+}
+
 @end
 
 // Defines an iOS AppCheckProviderFactory that wraps a given C++ Factory.

app_check/src_java/com/google/firebase/appcheck/internal/cpp/JniAppCheckProvider.java

@@ -41,6 +41,15 @@ public Task<AppCheckToken> getToken() {
     return taskCompletionSource.getTask();
   }
 
+  public Task<AppCheckToken> getLimitedUseToken() {
+    TaskCompletionSource<AppCheckToken> taskCompletionSource =
+        new TaskCompletionSource<AppCheckToken>();
+    // Call the C++ provider to get an AppCheckToken and set the task result.
+    // The C++ code will call handleGetTokenResult with the resulting token.
+    nativeGetLimitedUseToken(cProvider, taskCompletionSource);
+    return taskCompletionSource.getTask();
+  }
+
   /**
    *  Called by C++ with a token in order to complete the java task.
    */
@@ -58,4 +67,10 @@ public void handleGetTokenResult(TaskCompletionSource<AppCheckToken> taskComplet
    */
   private native void nativeGetToken(
       long cProvider, TaskCompletionSource<AppCheckToken> task_completion_source);
+
+  /**
+   * This function is implemented in the AppCheck C++ library (app_check_android.cc).
+   */
+  private native void nativeGetLimitedUseToken(
+      long cProvider, TaskCompletionSource<AppCheckToken> task_completion_source);
 }