From f4651ceadf8a78b40fab259d626db0245945a7da Mon Sep 17 00:00:00 2001
From: Reema Bajwa <reemabajwa@google.com>
Date: Fri, 5 May 2023 18:23:29 +0000
Subject: Hook up binderDied callback to cancel session when service dies

Test: Built & deplyed locally
Bug: 280869004

Change-Id: Id3e23b94b9169c3d5cca3d5b2695c25cb2758934
---
 .../server/credentials/ProviderClearSession.java   |   3 +-
 .../server/credentials/ProviderCreateSession.java  |   3 +-
 .../server/credentials/ProviderGetSession.java     |   3 +-
 .../credentials/RemoteCredentialService.java       | 133 +++++++++++++++------
 4 files changed, 103 insertions(+), 39 deletions(-)

diff --git a/services/credentials/java/com/android/server/credentials/ProviderClearSession.java b/services/credentials/java/com/android/server/credentials/ProviderClearSession.java
index 1f0346a2b9d6..d4b88001111e 100644
--- a/services/credentials/java/com/android/server/credentials/ProviderClearSession.java
+++ b/services/credentials/java/com/android/server/credentials/ProviderClearSession.java
@@ -132,7 +132,8 @@ public final class ProviderClearSession extends ProviderSession<ClearCredentialS
     protected void invokeSession() {
         if (mRemoteCredentialService != null) {
             startCandidateMetrics();
-            mRemoteCredentialService.onClearCredentialState(mProviderRequest, this);
+            mRemoteCredentialService.setCallback(this);
+            mRemoteCredentialService.onClearCredentialState(mProviderRequest);
         }
     }
 }
diff --git a/services/credentials/java/com/android/server/credentials/ProviderCreateSession.java b/services/credentials/java/com/android/server/credentials/ProviderCreateSession.java
index 16beaa4eb13e..409806a21679 100644
--- a/services/credentials/java/com/android/server/credentials/ProviderCreateSession.java
+++ b/services/credentials/java/com/android/server/credentials/ProviderCreateSession.java
@@ -248,7 +248,8 @@ public final class ProviderCreateSession extends ProviderSession<
     protected void invokeSession() {
         if (mRemoteCredentialService != null) {
             startCandidateMetrics();
-            mRemoteCredentialService.onBeginCreateCredential(mProviderRequest, this);
+            mRemoteCredentialService.setCallback(this);
+            mRemoteCredentialService.onBeginCreateCredential(mProviderRequest);
         }
     }
 
diff --git a/services/credentials/java/com/android/server/credentials/ProviderGetSession.java b/services/credentials/java/com/android/server/credentials/ProviderGetSession.java
index 1e80703378e0..64438e338e5b 100644
--- a/services/credentials/java/com/android/server/credentials/ProviderGetSession.java
+++ b/services/credentials/java/com/android/server/credentials/ProviderGetSession.java
@@ -309,7 +309,8 @@ public final class ProviderGetSession extends ProviderSession<BeginGetCredential
     protected void invokeSession() {
         if (mRemoteCredentialService != null) {
             startCandidateMetrics();
-            mRemoteCredentialService.onBeginGetCredential(mProviderRequest, this);
+            mRemoteCredentialService.setCallback(this);
+            mRemoteCredentialService.onBeginGetCredential(mProviderRequest);
         }
     }
 
diff --git a/services/credentials/java/com/android/server/credentials/RemoteCredentialService.java b/services/credentials/java/com/android/server/credentials/RemoteCredentialService.java
index 83e7e939e064..4bcf8be0d21e 100644
--- a/services/credentials/java/com/android/server/credentials/RemoteCredentialService.java
+++ b/services/credentials/java/com/android/server/credentials/RemoteCredentialService.java
@@ -48,6 +48,7 @@ import java.util.concurrent.CancellationException;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 
 /**
@@ -66,6 +67,10 @@ public class RemoteCredentialService extends ServiceConnector.Impl<ICredentialPr
 
     private final ComponentName mComponentName;
 
+    private AtomicBoolean mOngoingRequest = new AtomicBoolean(false);
+
+    @Nullable private ProviderCallbacks mCallback;
+
     /**
      * Callbacks to be invoked when the provider remote service responds with a
      * success or failure.
@@ -94,12 +99,35 @@ public class RemoteCredentialService extends ServiceConnector.Impl<ICredentialPr
         mComponentName = componentName;
     }
 
+    public void setCallback(ProviderCallbacks callback) {
+        mCallback = callback;
+    }
+
     /** Unbinds automatically after this amount of time. */
     @Override
     protected long getAutoDisconnectTimeoutMs() {
         return TIMEOUT_IDLE_SERVICE_CONNECTION_MILLIS;
     }
 
+    @Override
+    public void onBindingDied(ComponentName name) {
+        super.onBindingDied(name);
+
+        Slog.w(TAG, "binding died for: " + name);
+    }
+
+    @Override
+    public void binderDied() {
+        super.binderDied();
+        Slog.w(TAG, "binderDied");
+
+        if (mCallback != null) {
+            mOngoingRequest.set(false);
+            mCallback.onProviderServiceDied(this);
+        }
+
+    }
+
     /** Return the componentName of the service to be connected. */
     @NonNull
     public ComponentName getComponentName() {
@@ -116,11 +144,14 @@ public class RemoteCredentialService extends ServiceConnector.Impl<ICredentialPr
      * provider service.
      *
      * @param request  the request to be sent to the provider
-     * @param callback the callback to be used to send back the provider response to the
-     *                 {@link ProviderGetSession} class that maintains provider state
      */
-    public void onBeginGetCredential(@NonNull BeginGetCredentialRequest request,
-            ProviderCallbacks<BeginGetCredentialResponse> callback) {
+    public void onBeginGetCredential(@NonNull BeginGetCredentialRequest request) {
+        if (mCallback == null) {
+            Slog.w(TAG, "Callback is not set");
+            return;
+        }
+        mOngoingRequest.set(true);
+
         AtomicReference<ICancellationSignal> cancellationSink = new AtomicReference<>();
         AtomicReference<CompletableFuture<BeginGetCredentialResponse>> futureRef =
                 new AtomicReference<>();
@@ -154,7 +185,9 @@ public class RemoteCredentialService extends ServiceConnector.Impl<ICredentialPr
                                     dispatchCancellationSignal(cancellation);
                                 } else {
                                     cancellationSink.set(cancellation);
-                                    callback.onProviderCancellable(cancellation);
+                                    if (mCallback != null) {
+                                        mCallback.onProviderCancellable(cancellation);
+                                    }
                                 }
                             }
                         });
@@ -166,7 +199,7 @@ public class RemoteCredentialService extends ServiceConnector.Impl<ICredentialPr
         futureRef.set(connectThenExecute);
 
         connectThenExecute.whenComplete((result, error) -> Handler.getMain().post(() ->
-                handleExecutionResponse(result, error, cancellationSink, callback)));
+                handleExecutionResponse(result, error, cancellationSink)));
     }
 
     /**
@@ -174,11 +207,14 @@ public class RemoteCredentialService extends ServiceConnector.Impl<ICredentialPr
      * provider service.
      *
      * @param request  the request to be sent to the provider
-     * @param callback the callback to be used to send back the provider response to the
-     *                 {@link ProviderCreateSession} class that maintains provider state
      */
-    public void onBeginCreateCredential(@NonNull BeginCreateCredentialRequest request,
-            ProviderCallbacks<BeginCreateCredentialResponse> callback) {
+    public void onBeginCreateCredential(@NonNull BeginCreateCredentialRequest request) {
+        if (mCallback == null) {
+            Slog.w(TAG, "Callback is not set");
+            return;
+        }
+        mOngoingRequest.set(true);
+
         AtomicReference<ICancellationSignal> cancellationSink = new AtomicReference<>();
         AtomicReference<CompletableFuture<BeginCreateCredentialResponse>> futureRef =
                 new AtomicReference<>();
@@ -212,7 +248,9 @@ public class RemoteCredentialService extends ServiceConnector.Impl<ICredentialPr
                                             dispatchCancellationSignal(cancellation);
                                         } else {
                                             cancellationSink.set(cancellation);
-                                            callback.onProviderCancellable(cancellation);
+                                            if (mCallback != null) {
+                                                mCallback.onProviderCancellable(cancellation);
+                                            }
                                         }
                                     }
                                 });
@@ -224,7 +262,7 @@ public class RemoteCredentialService extends ServiceConnector.Impl<ICredentialPr
         futureRef.set(connectThenExecute);
 
         connectThenExecute.whenComplete((result, error) -> Handler.getMain().post(() ->
-                handleExecutionResponse(result, error, cancellationSink, callback)));
+                handleExecutionResponse(result, error, cancellationSink)));
     }
 
     /**
@@ -232,11 +270,14 @@ public class RemoteCredentialService extends ServiceConnector.Impl<ICredentialPr
      * provider service.
      *
      * @param request  the request to be sent to the provider
-     * @param callback the callback to be used to send back the provider response to the
-     *                 {@link ProviderClearSession} class that maintains provider state
      */
-    public void onClearCredentialState(@NonNull ClearCredentialStateRequest request,
-            ProviderCallbacks<Void> callback) {
+    public void onClearCredentialState(@NonNull ClearCredentialStateRequest request) {
+        if (mCallback == null) {
+            Slog.w(TAG, "Callback is not set");
+            return;
+        }
+        mOngoingRequest.set(true);
+
         AtomicReference<ICancellationSignal> cancellationSink = new AtomicReference<>();
         AtomicReference<CompletableFuture<Void>> futureRef = new AtomicReference<>();
 
@@ -269,7 +310,9 @@ public class RemoteCredentialService extends ServiceConnector.Impl<ICredentialPr
                                             dispatchCancellationSignal(cancellation);
                                         } else {
                                             cancellationSink.set(cancellation);
-                                            callback.onProviderCancellable(cancellation);
+                                            if (mCallback != null) {
+                                                mCallback.onProviderCancellable(cancellation);
+                                            }
                                         }
                                     }
                                 });
@@ -281,40 +324,58 @@ public class RemoteCredentialService extends ServiceConnector.Impl<ICredentialPr
         futureRef.set(connectThenExecute);
 
         connectThenExecute.whenComplete((result, error) -> Handler.getMain().post(() ->
-                handleExecutionResponse(result, error, cancellationSink, callback)));
+                handleExecutionResponse(result, error, cancellationSink)));
     }
 
     private <T> void handleExecutionResponse(T result,
             Throwable error,
-            AtomicReference<ICancellationSignal> cancellationSink,
-            ProviderCallbacks<T> callback) {
+            AtomicReference<ICancellationSignal> cancellationSink) {
         if (error == null) {
-            callback.onProviderResponseSuccess(result);
+            if (mCallback != null) {
+                mCallback.onProviderResponseSuccess(result);
+            }
         } else {
             if (error instanceof TimeoutException) {
                 Slog.i(TAG, "Remote provider response timed tuo for: " + mComponentName);
+                if (!mOngoingRequest.get()) {
+                    return;
+                }
                 dispatchCancellationSignal(cancellationSink.get());
-                callback.onProviderResponseFailure(
-                        CredentialProviderErrors.ERROR_TIMEOUT,
-                        null);
+                if (mCallback != null) {
+                    mOngoingRequest.set(false);
+                    mCallback.onProviderResponseFailure(
+                            CredentialProviderErrors.ERROR_TIMEOUT, null);
+                }
             } else if (error instanceof CancellationException) {
                 Slog.i(TAG, "Cancellation exception for remote provider: " + mComponentName);
+                if (!mOngoingRequest.get()) {
+                    return;
+                }
                 dispatchCancellationSignal(cancellationSink.get());
-                callback.onProviderResponseFailure(
-                        CredentialProviderErrors.ERROR_TASK_CANCELED,
-                        null);
+                if (mCallback != null) {
+                    mOngoingRequest.set(false);
+                    mCallback.onProviderResponseFailure(
+                            CredentialProviderErrors.ERROR_TASK_CANCELED,
+                            null);
+                }
             } else if (error instanceof GetCredentialException) {
-                callback.onProviderResponseFailure(
-                        CredentialProviderErrors.ERROR_PROVIDER_FAILURE,
-                        (GetCredentialException) error);
+                if (mCallback != null) {
+                    mCallback.onProviderResponseFailure(
+                            CredentialProviderErrors.ERROR_PROVIDER_FAILURE,
+                            (GetCredentialException) error);
+                }
             } else if (error instanceof CreateCredentialException) {
-                callback.onProviderResponseFailure(
-                        CredentialProviderErrors.ERROR_PROVIDER_FAILURE,
-                        (CreateCredentialException) error);
+                if (mCallback != null) {
+                    mCallback.onProviderResponseFailure(
+                            CredentialProviderErrors.ERROR_PROVIDER_FAILURE,
+                            (CreateCredentialException) error);
+                }
             } else {
-                callback.onProviderResponseFailure(
-                        CredentialProviderErrors.ERROR_UNKNOWN,
-                        (Exception) error);
+                if (mCallback != null) {
+                    mCallback.onProviderResponseFailure(
+                            CredentialProviderErrors.ERROR_UNKNOWN,
+                            (Exception) error);
+                }
             }
         }
     }
-- 
cgit v1.2.3-59-g8ed1b