Add cancellation signal death recipient.

Flag: android.app.appfunctions.flags.enable_app_function_manager
Test: cts
Bug: 357551503
Change-Id: I68a2d5acd9d69d181be46ba27d3d6696ae4378ef

3 files changed, 55 insertions, 20 deletions

diff --git a/services/appfunctions/java/com/android/server/appfunctions/AppFunctionManagerServiceImpl.java b/services/appfunctions/java/com/android/server/appfunctions/AppFunctionManagerServiceImpl.java
index 14298f416d87..c87d516d2ab4 100644
--- a/services/appfunctions/java/com/android/server/appfunctions/AppFunctionManagerServiceImpl.java
+++ b/services/appfunctions/java/com/android/server/appfunctions/AppFunctionManagerServiceImpl.java
@@ -53,6 +53,7 @@ import android.content.Context;
 import android.content.Intent;
 import android.os.Binder;
 import android.os.CancellationSignal;
+import android.os.IBinder;
 import android.os.ICancellationSignal;
 import android.os.OutcomeReceiver;
 import android.os.ParcelableException;
@@ -160,7 +161,8 @@ public class AppFunctionManagerServiceImpl extends IAppFunctionManager.Stub {
                                 callingUid,
                                 callingPid,
                                 localCancelTransport,
-                                safeExecuteAppFunctionCallback);
+                                safeExecuteAppFunctionCallback,
+                                executeAppFunctionCallback.asBinder());
                     } catch (Exception e) {
                         safeExecuteAppFunctionCallback.onResult(
                                 mapExceptionToExecuteAppFunctionResponse(e));
@@ -171,11 +173,12 @@ public class AppFunctionManagerServiceImpl extends IAppFunctionManager.Stub {
 
     @WorkerThread
     private void executeAppFunctionInternal(
-            ExecuteAppFunctionAidlRequest requestInternal,
+            @NonNull ExecuteAppFunctionAidlRequest requestInternal,
             int callingUid,
             int callingPid,
-            ICancellationSignal localCancelTransport,
-            SafeOneTimeExecuteAppFunctionCallback safeExecuteAppFunctionCallback) {
+            @NonNull ICancellationSignal localCancelTransport,
+            @NonNull SafeOneTimeExecuteAppFunctionCallback safeExecuteAppFunctionCallback,
+            @NonNull IBinder callerBinder) {
         UserHandle targetUser = requestInternal.getUserHandle();
         // TODO(b/354956319): Add and honor the new enterprise policies.
         if (mCallerValidator.isUserOrganizationManaged(targetUser)) {
@@ -250,7 +253,8 @@ public class AppFunctionManagerServiceImpl extends IAppFunctionManager.Stub {
                                     localCancelTransport,
                                     safeExecuteAppFunctionCallback,
                                     /* bindFlags= */ Context.BIND_AUTO_CREATE
-                                            | Context.BIND_FOREGROUND_SERVICE);
+                                            | Context.BIND_FOREGROUND_SERVICE,
+                                    callerBinder);
                         })
                 .exceptionally(
                         ex -> {
@@ -326,8 +330,9 @@ public class AppFunctionManagerServiceImpl extends IAppFunctionManager.Stub {
 
     /**
      * Sets the enabled status of a specified app function.
-     * <p>
-     * Required to hold a lock to call this function to avoid document changes during the process.
+     *
+     * <p>Required to hold a lock to call this function to avoid document changes during the
+     * process.
      */
     @WorkerThread
     @GuardedBy("mLock")
@@ -370,8 +375,7 @@ public class AppFunctionManagerServiceImpl extends IAppFunctionManager.Stub {
                     newMetadata.setEnabled(false);
                 }
                 default ->
-                        throw new IllegalArgumentException(
-                                "Value of EnabledState is unsupported.");
+                        throw new IllegalArgumentException("Value of EnabledState is unsupported.");
             }
             AppSearchBatchResult<String, Void> putDocumentBatchResult =
                     runtimeMetadataSearchSession
@@ -381,8 +385,8 @@ public class AppFunctionManagerServiceImpl extends IAppFunctionManager.Stub {
                                             .build())
                             .get();
             if (!putDocumentBatchResult.isSuccess()) {
-                throw new IllegalStateException("Failed writing updated doc to AppSearch due to "
-                        + putDocumentBatchResult);
+                throw new IllegalStateException(
+                        "Failed writing updated doc to AppSearch due to " + putDocumentBatchResult);
             }
         }
     }
@@ -414,7 +418,8 @@ public class AppFunctionManagerServiceImpl extends IAppFunctionManager.Stub {
             @NonNull UserHandle targetUser,
             @NonNull ICancellationSignal cancellationSignalTransport,
             @NonNull SafeOneTimeExecuteAppFunctionCallback safeExecuteAppFunctionCallback,
-            int bindFlags) {
+            int bindFlags,
+            @NonNull IBinder callerBinder) {
         CancellationSignal cancellationSignal =
                 CancellationSignal.fromTransport(cancellationSignalTransport);
         ICancellationCallback cancellationCallback =
@@ -480,7 +485,8 @@ public class AppFunctionManagerServiceImpl extends IAppFunctionManager.Stub {
                                 // proceed after initiating a cancellation.
                                 safeExecuteAppFunctionCallback.disable();
                             }
-                        });
+                        },
+                        callerBinder);
 
         if (!bindServiceResult) {
             Slog.e(TAG, "Failed to bind to the AppFunctionService");
diff --git a/services/appfunctions/java/com/android/server/appfunctions/RemoteServiceCaller.java b/services/appfunctions/java/com/android/server/appfunctions/RemoteServiceCaller.java
index 55173c350e71..d0e858e71de4 100644
--- a/services/appfunctions/java/com/android/server/appfunctions/RemoteServiceCaller.java
+++ b/services/appfunctions/java/com/android/server/appfunctions/RemoteServiceCaller.java
@@ -18,6 +18,7 @@ package com.android.server.appfunctions;
 import android.annotation.NonNull;
 import android.content.Intent;
 import android.os.CancellationSignal;
+import android.os.IBinder;
 import android.os.UserHandle;
 
 /**
@@ -30,9 +31,10 @@ import android.os.UserHandle;
 public interface RemoteServiceCaller<T> {
 
     /**
-     * Initiates service binding and executes a provided method when the service connects. Unbinds
-     * the service after execution or upon cancellation timeout. Returns the result of the
-     * bindService API.
+     * Initiates service binding and executes a provided method when the service connects.
+     *
+     * <p>Unbinds the service after execution or upon cancellation timeout or calling process death.
+     * Returns the result of the bindService API.
      *
      * <p>When the service connection was made successfully, it's the caller responsibility to
      * report the usage is completed and can be unbound by calling {@link
@@ -43,6 +45,9 @@ public interface RemoteServiceCaller<T> {
      * return the result within `cancellationTimeoutMillis` after the cancellation signal is sent,
      * this method will unbind the service connection.
      *
+     * <p>This method will also unbind the service after the calling process dies (because a
+     * cancellation signal cannot be sent and system server can become bound forever if otherwise).
+     *
      * @param intent An Intent object that describes the service that should be bound.
      * @param bindFlags Flags used to control the binding process See {@link
      *     android.content.Context#bindService}.
@@ -52,6 +57,7 @@ public interface RemoteServiceCaller<T> {
      * @param cancellationSignal The cancellation signal forwarded to the service.
      * @param callback A callback to be invoked for various events. See {@link
      *     RunServiceCallCallback}.
+     * @param callerBinder The binder of the caller.
      */
     boolean runServiceCall(
             @NonNull Intent intent,
@@ -59,7 +65,8 @@ public interface RemoteServiceCaller<T> {
             @NonNull UserHandle userHandle,
             long cancellationTimeoutMillis,
             @NonNull CancellationSignal cancellationSignal,
-            @NonNull RunServiceCallCallback<T> callback);
+            @NonNull RunServiceCallCallback<T> callback,
+            @NonNull IBinder callerBinder);
 
     /** An interface for clients to signal that they have finished using a bound service. */
     interface ServiceUsageCompleteListener {
diff --git a/services/appfunctions/java/com/android/server/appfunctions/RemoteServiceCallerImpl.java b/services/appfunctions/java/com/android/server/appfunctions/RemoteServiceCallerImpl.java
index 7e3a29afc01c..15081598c301 100644
--- a/services/appfunctions/java/com/android/server/appfunctions/RemoteServiceCallerImpl.java
+++ b/services/appfunctions/java/com/android/server/appfunctions/RemoteServiceCallerImpl.java
@@ -16,6 +16,7 @@
 package com.android.server.appfunctions;
 
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
@@ -24,8 +25,10 @@ import android.os.CancellationSignal;
 import android.os.Handler;
 import android.os.IBinder;
 import android.os.Looper;
+import android.os.RemoteException;
 import android.os.UserHandle;
 import android.util.Log;
+import android.util.Slog;
 
 import java.util.concurrent.Executor;
 import java.util.function.Function;
@@ -66,7 +69,8 @@ public class RemoteServiceCallerImpl<T> implements RemoteServiceCaller<T> {
             @NonNull UserHandle userHandle,
             long cancellationTimeoutMillis,
             @NonNull CancellationSignal cancellationSignal,
-            @NonNull RunServiceCallCallback<T> callback) {
+            @NonNull RunServiceCallCallback<T> callback,
+            @NonNull IBinder callerBinder) {
         OneOffServiceConnection serviceConnection =
                 new OneOffServiceConnection(
                         intent,
@@ -74,7 +78,8 @@ public class RemoteServiceCallerImpl<T> implements RemoteServiceCaller<T> {
                         userHandle,
                         cancellationTimeoutMillis,
                         cancellationSignal,
-                        callback);
+                        callback,
+                        callerBinder);
 
         return serviceConnection.bindAndRun();
     }
@@ -88,6 +93,8 @@ public class RemoteServiceCallerImpl<T> implements RemoteServiceCaller<T> {
         private final long mCancellationTimeoutMillis;
         private final CancellationSignal mCancellationSignal;
         private final Runnable mCancellationTimeoutRunnable;
+        private final IBinder mCallerBinder;
+        @Nullable private IBinder.DeathRecipient mDirectServiceVulture;
 
         OneOffServiceConnection(
                 @NonNull Intent intent,
@@ -95,7 +102,8 @@ public class RemoteServiceCallerImpl<T> implements RemoteServiceCaller<T> {
                 @NonNull UserHandle userHandle,
                 long cancellationTimeoutMillis,
                 @NonNull CancellationSignal cancellationSignal,
-                @NonNull RunServiceCallCallback<T> callback) {
+                @NonNull RunServiceCallCallback<T> callback,
+                @NonNull IBinder callerBinder) {
             mIntent = intent;
             mFlags = flags;
             mCallback = callback;
@@ -103,6 +111,7 @@ public class RemoteServiceCallerImpl<T> implements RemoteServiceCaller<T> {
             mCancellationTimeoutMillis = cancellationTimeoutMillis;
             mCancellationSignal = cancellationSignal;
             mCancellationTimeoutRunnable = this::safeUnbind;
+            mCallerBinder = callerBinder;
         }
 
         public boolean bindAndRun() {
@@ -116,6 +125,16 @@ public class RemoteServiceCallerImpl<T> implements RemoteServiceCaller<T> {
                             mHandler.postDelayed(
                                     mCancellationTimeoutRunnable, mCancellationTimeoutMillis);
                         });
+                mDirectServiceVulture =
+                        () -> {
+                            Slog.w(TAG, "Caller process onDeath signal received");
+                            mCancellationSignal.cancel();
+                        };
+                try {
+                    mCallerBinder.linkToDeath(mDirectServiceVulture, /* flags= */ 0);
+                } catch (RemoteException e) {
+                    Slog.w(TAG, "Failed to link to death on " + mCallerBinder + ": ", e);
+                }
             } else {
                 safeUnbind();
             }
@@ -152,6 +171,9 @@ public class RemoteServiceCallerImpl<T> implements RemoteServiceCaller<T> {
             try {
                 mHandler.removeCallbacks(mCancellationTimeoutRunnable);
                 mContext.unbindService(this);
+                if (mDirectServiceVulture != null) {
+                    mCallerBinder.unlinkToDeath(mDirectServiceVulture, 0);
+                }
             } catch (Exception ex) {
                 Log.w(TAG, "Failed to unbind", ex);
             }