Add bundle validations on request/response and state params.

This is the implementation for approach#2 in go/isolated-bundle-validation to ensure that bundle being passed to isolated process is essentially read-only.

AndroidFuture<T> does not support any raw Binder interfaces.  For
instance, AndroidFuture<ICancellationCallback> would crash due to
java.lang.ClassCastException at run time. You must use IBinder
instead and manually create a stub object in the receiver side. Hence the generic usage of AndroidFuture across.

Bug: 326618319
Change-Id: Idb7aad4dca45d29683ae9b85532f790f3cf643ad
diff --git a/core/java/android/app/ondeviceintelligence/OnDeviceIntelligenceManager.java b/core/java/android/app/ondeviceintelligence/OnDeviceIntelligenceManager.java
index bc50d2e4..5e1c1e0 100644
--- a/core/java/android/app/ondeviceintelligence/OnDeviceIntelligenceManager.java
+++ b/core/java/android/app/ondeviceintelligence/OnDeviceIntelligenceManager.java
@@ -365,7 +365,6 @@
      *                           associated params.
      */
     @RequiresPermission(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE)
-
     public void processRequest(@NonNull Feature feature, @NonNull @InferenceParams Bundle request,
             @RequestType int requestType,
             @Nullable CancellationSignal cancellationSignal,
@@ -423,16 +422,16 @@
      * when the final response contains an enhanced aggregation of the contents already
      * streamed.
      *
-     * @param feature                   feature associated with the request.
-     * @param request                   request and associated params represented by the Bundle
-     *                                  data.
-     * @param requestType               type of request being sent for processing the content.
-     * @param cancellationSignal        signal to invoke cancellation.
-     * @param processingSignal          signal to send custom signals in the
-     *                                  remote implementation.
-     * @param streamingResponseCallback streaming callback to populate the response content and
-     *                                  associated params.
-     * @param callbackExecutor          executor to run the callback on.
+     * @param feature                     feature associated with the request.
+     * @param request                     request and associated params represented by the Bundle
+     *                                    data.
+     * @param requestType                 type of request being sent for processing the content.
+     * @param cancellationSignal          signal to invoke cancellation.
+     * @param processingSignal            signal to send custom signals in the
+     *                                    remote implementation.
+     * @param streamingProcessingCallback streaming callback to populate the response content and
+     *                                    associated params.
+     * @param callbackExecutor            executor to run the callback on.
      */
     @RequiresPermission(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE)
     public void processRequestStreaming(@NonNull Feature feature,
@@ -528,19 +527,18 @@
 
     /**
      * {@link Bundle}s annotated with this type will be validated that they are in-effect read-only
-     * when passed to inference service via Binder IPC. Following restrictions apply :
+     * when passed via Binder IPC. Following restrictions apply :
      * <ul>
+     * <li> No Nested Bundles are allowed.</li>
+     * <li> {@link PersistableBundle}s are allowed.</li>
      * <li> Any primitive types or their collections can be added as usual.</li>
      * <li>IBinder objects should *not* be added.</li>
      * <li>Parcelable data which has no active-objects, should be added as
      * {@link Bundle#putByteArray}</li>
      * <li>Parcelables have active-objects, only following types will be allowed</li>
      * <ul>
-     *  <li>{@link Bitmap} set as {@link Bitmap#setImmutable()}</li>
-     *  <li>{@link android.database.CursorWindow}</li>
      *  <li>{@link android.os.ParcelFileDescriptor} opened in
      *  {@link android.os.ParcelFileDescriptor#MODE_READ_ONLY}</li>
-     *  <li>{@link android.os.SharedMemory} set to {@link OsConstants#PROT_READ}</li>
      * </ul>
      * </ul>
      *
@@ -550,9 +548,40 @@
      * @hide
      */
     @Target({ElementType.PARAMETER, ElementType.FIELD})
+    public @interface StateParams {
+    }
+
+    /**
+     * This is an extension of {@link StateParams} but for purpose of inference few other types are
+     * also allowed as read-only, as listed below.
+     *
+     * <li>{@link Bitmap} set as immutable.</li>
+     * <li>{@link android.database.CursorWindow}</li>
+     * <li>{@link android.os.SharedMemory} set to {@link OsConstants#PROT_READ}</li>
+     * </ul>
+     * </ul>
+     *
+     * In all other scenarios the system-server might throw a
+     * {@link android.os.BadParcelableException} if the Bundle validation fails.
+     *
+     * @hide
+     */
+    @Target({ElementType.PARAMETER, ElementType.FIELD, ElementType.TYPE_USE})
     public @interface InferenceParams {
     }
 
+    /**
+     * This is an extension of {@link StateParams} with the exception that it allows writing
+     * {@link Bitmap} as part of the response.
+     *
+     * In all other scenarios the system-server might throw a
+     * {@link android.os.BadParcelableException} if the Bundle validation fails.
+     *
+     * @hide
+     */
+    @Target({ElementType.PARAMETER, ElementType.FIELD})
+    public @interface ResponseParams {
+    }
 
     @Nullable
     private static AndroidFuture<IBinder> configureRemoteCancellationFuture(
diff --git a/core/java/android/app/ondeviceintelligence/ProcessingCallback.java b/core/java/android/app/ondeviceintelligence/ProcessingCallback.java
index 4d936ea..e50d6b1 100644
--- a/core/java/android/app/ondeviceintelligence/ProcessingCallback.java
+++ b/core/java/android/app/ondeviceintelligence/ProcessingCallback.java
@@ -23,6 +23,7 @@
 import android.annotation.SystemApi;
 import android.os.Bundle;
 import android.app.ondeviceintelligence.OnDeviceIntelligenceManager.InferenceParams;
+import android.app.ondeviceintelligence.OnDeviceIntelligenceManager.ResponseParams;
 
 import java.util.function.Consumer;
 
@@ -43,7 +44,7 @@
      *
      * @param result Response to be passed as a result.
      */
-    void onResult(@NonNull @InferenceParams Bundle result);
+    void onResult(@NonNull @ResponseParams Bundle result);
 
     /**
      * Called when the request processing fails. The failure details are indicated by the
@@ -64,8 +65,8 @@
      *                         expected to be non-null or EMPTY when there is no response.
      */
     default void onDataAugmentRequest(
-            @NonNull @InferenceParams Bundle processedContent,
-            @NonNull Consumer<Bundle> contentConsumer) {
+            @NonNull @ResponseParams Bundle processedContent,
+            @NonNull Consumer<@InferenceParams Bundle> contentConsumer) {
         contentConsumer.accept(Bundle.EMPTY);
     }
 }
diff --git a/core/java/android/app/ondeviceintelligence/StreamingProcessingCallback.java b/core/java/android/app/ondeviceintelligence/StreamingProcessingCallback.java
index 41f1807..7ee2af7 100644
--- a/core/java/android/app/ondeviceintelligence/StreamingProcessingCallback.java
+++ b/core/java/android/app/ondeviceintelligence/StreamingProcessingCallback.java
@@ -22,7 +22,7 @@
 import android.annotation.NonNull;
 import android.annotation.SystemApi;
 import android.os.Bundle;
-import android.app.ondeviceintelligence.OnDeviceIntelligenceManager.InferenceParams;
+import android.app.ondeviceintelligence.OnDeviceIntelligenceManager.ResponseParams;
 
 /**
  * Streaming variant of {@link ProcessingCallback} to populate response while processing a given
@@ -37,5 +37,5 @@
      * Callback that would be invoked when a part of the response i.e. some response is
      * already processed, and needs to be passed onto the caller.
      */
-    void onPartialResult(@NonNull @InferenceParams Bundle partialResult);
+    void onPartialResult(@NonNull @ResponseParams Bundle partialResult);
 }
diff --git a/core/java/android/service/ondeviceintelligence/IOnDeviceIntelligenceService.aidl b/core/java/android/service/ondeviceintelligence/IOnDeviceIntelligenceService.aidl
index a708718..45c4350 100644
--- a/core/java/android/service/ondeviceintelligence/IOnDeviceIntelligenceService.aidl
+++ b/core/java/android/service/ondeviceintelligence/IOnDeviceIntelligenceService.aidl
@@ -42,7 +42,7 @@
     void getReadOnlyFileDescriptor(in String fileName, in AndroidFuture<ParcelFileDescriptor> future);
     void getReadOnlyFeatureFileDescriptorMap(in Feature feature, in RemoteCallback remoteCallback);
     void requestFeatureDownload(int callerUid, in Feature feature,
-                                in AndroidFuture<ICancellationSignal> cancellationSignal,
+                                in AndroidFuture cancellationSignal,
                                 in IDownloadCallback downloadCallback);
     void registerRemoteServices(in IRemoteProcessingService remoteProcessingService);
     void notifyInferenceServiceConnected();
diff --git a/core/java/android/service/ondeviceintelligence/IOnDeviceSandboxedInferenceService.aidl b/core/java/android/service/ondeviceintelligence/IOnDeviceSandboxedInferenceService.aidl
index 4213a09..2aa17c4 100644
--- a/core/java/android/service/ondeviceintelligence/IOnDeviceSandboxedInferenceService.aidl
+++ b/core/java/android/service/ondeviceintelligence/IOnDeviceSandboxedInferenceService.aidl
@@ -36,15 +36,15 @@
 oneway interface IOnDeviceSandboxedInferenceService {
     void registerRemoteStorageService(in IRemoteStorageService storageService);
     void requestTokenInfo(int callerUid, in Feature feature, in Bundle request,
-                            in AndroidFuture<ICancellationSignal> cancellationSignal,
+                            in AndroidFuture cancellationSignal,
                             in ITokenInfoCallback tokenInfoCallback);
     void processRequest(int callerUid, in Feature feature, in Bundle request, in int requestType,
-                        in AndroidFuture<ICancellationSignal> cancellationSignal,
-                        in AndroidFuture<IProcessingSignal> processingSignal,
+                        in AndroidFuture cancellationSignal,
+                        in AndroidFuture processingSignal,
                         in IResponseCallback callback);
     void processRequestStreaming(int callerUid, in Feature feature, in Bundle request, in int requestType,
-                                in AndroidFuture<ICancellationSignal> cancellationSignal,
-                                in AndroidFuture<IProcessingSignal> processingSignal,
+                                in AndroidFuture cancellationSignal,
+                                in AndroidFuture processingSignal,
                                 in IStreamingResponseCallback callback);
     void updateProcessingState(in Bundle processingState,
                                      in IProcessingUpdateStatusCallback callback);
diff --git a/core/java/android/service/ondeviceintelligence/OnDeviceIntelligenceService.java b/core/java/android/service/ondeviceintelligence/OnDeviceIntelligenceService.java
index 5dc540a..793e58a 100644
--- a/core/java/android/service/ondeviceintelligence/OnDeviceIntelligenceService.java
+++ b/core/java/android/service/ondeviceintelligence/OnDeviceIntelligenceService.java
@@ -35,7 +35,7 @@
 import android.app.ondeviceintelligence.IListFeaturesCallback;
 import android.app.ondeviceintelligence.OnDeviceIntelligenceException;
 import android.app.ondeviceintelligence.OnDeviceIntelligenceManager;
-import android.app.ondeviceintelligence.OnDeviceIntelligenceManager.InferenceParams;
+import android.app.ondeviceintelligence.OnDeviceIntelligenceManager.StateParams;
 import android.content.Intent;
 import android.os.Binder;
 import android.os.Bundle;
@@ -245,13 +245,13 @@
      * service if there is a state change to be performed. State change could be config updates,
      * performing initialization or cleanup tasks in the remote inference service.
      * The Bundle passed in here is expected to be read-only and will be rejected if it has any
-     * writable fields as detailed under {@link InferenceParams}.
+     * writable fields as detailed under {@link StateParams}.
      *
      * @param processingState  the updated state to be applied.
      * @param callbackExecutor executor to the run status callback on.
      * @param statusReceiver   receiver to get status of the update state operation.
      */
-    public final void updateProcessingState(@NonNull @InferenceParams Bundle processingState,
+    public final void updateProcessingState(@NonNull @StateParams Bundle processingState,
             @NonNull @CallbackExecutor Executor callbackExecutor,
             @NonNull OutcomeReceiver<PersistableBundle, OnDeviceIntelligenceException> statusReceiver) {
         Objects.requireNonNull(callbackExecutor);
diff --git a/core/java/android/service/ondeviceintelligence/OnDeviceSandboxedInferenceService.java b/core/java/android/service/ondeviceintelligence/OnDeviceSandboxedInferenceService.java
index 96c45ee..29a6db6 100644
--- a/core/java/android/service/ondeviceintelligence/OnDeviceSandboxedInferenceService.java
+++ b/core/java/android/service/ondeviceintelligence/OnDeviceSandboxedInferenceService.java
@@ -35,6 +35,7 @@
 import android.app.ondeviceintelligence.OnDeviceIntelligenceException;
 import android.app.ondeviceintelligence.OnDeviceIntelligenceManager;
 import android.app.ondeviceintelligence.OnDeviceIntelligenceManager.InferenceParams;
+import android.app.ondeviceintelligence.OnDeviceIntelligenceManager.StateParams;
 import android.app.ondeviceintelligence.ProcessingCallback;
 import android.app.ondeviceintelligence.ProcessingSignal;
 import android.app.ondeviceintelligence.StreamingProcessingCallback;
@@ -293,7 +294,7 @@
      * @param callback        callback to populate the update status and if there are params
      *                        associated with the status.
      */
-    public abstract void onUpdateProcessingState(@NonNull @InferenceParams Bundle processingState,
+    public abstract void onUpdateProcessingState(@NonNull @StateParams Bundle processingState,
             @NonNull OutcomeReceiver<PersistableBundle,
                     OnDeviceIntelligenceException> callback);
 
diff --git a/services/core/java/com/android/server/ondeviceintelligence/BundleUtil.java b/services/core/java/com/android/server/ondeviceintelligence/BundleUtil.java
new file mode 100644
index 0000000..681dd0b
--- /dev/null
+++ b/services/core/java/com/android/server/ondeviceintelligence/BundleUtil.java
@@ -0,0 +1,352 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.ondeviceintelligence;
+
+import static android.system.OsConstants.F_GETFL;
+import static android.system.OsConstants.O_ACCMODE;
+import static android.system.OsConstants.O_RDONLY;
+import static android.system.OsConstants.PROT_READ;
+
+import android.app.ondeviceintelligence.IResponseCallback;
+import android.app.ondeviceintelligence.IStreamingResponseCallback;
+import android.app.ondeviceintelligence.ITokenInfoCallback;
+import android.app.ondeviceintelligence.OnDeviceIntelligenceManager.InferenceParams;
+import android.app.ondeviceintelligence.OnDeviceIntelligenceManager.ResponseParams;
+import android.app.ondeviceintelligence.OnDeviceIntelligenceManager.StateParams;
+import android.app.ondeviceintelligence.TokenInfo;
+import android.database.CursorWindow;
+import android.graphics.Bitmap;
+import android.os.BadParcelableException;
+import android.os.Bundle;
+import android.os.ParcelFileDescriptor;
+import android.os.PersistableBundle;
+import android.os.RemoteCallback;
+import android.os.RemoteException;
+import android.os.SharedMemory;
+import android.system.ErrnoException;
+import android.system.Os;
+import android.util.Log;
+
+import java.util.concurrent.Executor;
+
+/**
+ * Util methods for ensuring the Bundle passed in various methods are read-only and restricted to
+ * some known types.
+ */
+public class BundleUtil {
+    private static final String TAG = "BundleUtil";
+
+    /**
+     * Validation of the inference request payload as described in {@link InferenceParams}
+     * description.
+     *
+     * @throws BadParcelableException when the bundle does not meet the read-only requirements.
+     */
+    public static void sanitizeInferenceParams(
+            @InferenceParams Bundle bundle) {
+        ensureValidBundle(bundle);
+
+        if (!bundle.hasFileDescriptors()) {
+            return; //safe to exit if there are no FDs and Binders
+        }
+
+        for (String key : bundle.keySet()) {
+            Object obj = bundle.get(key);
+            if (obj == null) {
+                /* Null value here could also mean deserializing a custom parcelable has failed,
+                 *  and since {@link Bundle} is marked as defusable in system-server - the
+                 * {@link ClassNotFoundException} exception is swallowed and `null` is returned
+                 * instead. We want to ensure cleanup of null entries in such case.
+                 */
+                bundle.putObject(key, null);
+                continue;
+            }
+            if (canMarshall(obj) || obj instanceof CursorWindow) {
+                continue;
+            }
+
+            if (obj instanceof ParcelFileDescriptor) {
+                validatePfdReadOnly((ParcelFileDescriptor) obj);
+            } else if (obj instanceof SharedMemory) {
+                ((SharedMemory) obj).setProtect(PROT_READ);
+            } else if (obj instanceof Bitmap) {
+                if (((Bitmap) obj).isMutable()) {
+                    throw new BadParcelableException(
+                            "Encountered a mutable Bitmap in the Bundle at key : " + key);
+                }
+            } else {
+                throw new BadParcelableException(
+                        "Unsupported Parcelable type encountered in the Bundle: "
+                                + obj.getClass().getSimpleName());
+            }
+        }
+    }
+
+    /**
+     * Validation of the inference request payload as described in {@link ResponseParams}
+     * description.
+     *
+     * @throws BadParcelableException when the bundle does not meet the read-only requirements.
+     */
+    public static void sanitizeResponseParams(
+            @ResponseParams Bundle bundle) {
+        ensureValidBundle(bundle);
+
+        if (!bundle.hasFileDescriptors()) {
+            return; //safe to exit if there are no FDs and Binders
+        }
+
+        for (String key : bundle.keySet()) {
+            Object obj = bundle.get(key);
+            if (obj == null) {
+                /* Null value here could also mean deserializing a custom parcelable has failed,
+                 *  and since {@link Bundle} is marked as defusable in system-server - the
+                 * {@link ClassNotFoundException} exception is swallowed and `null` is returned
+                 * instead. We want to ensure cleanup of null entries in such case.
+                 */
+                bundle.putObject(key, null);
+                continue;
+            }
+            if (canMarshall(obj)) {
+                continue;
+            }
+
+            if (obj instanceof ParcelFileDescriptor) {
+                validatePfdReadOnly((ParcelFileDescriptor) obj);
+            } else if (obj instanceof Bitmap) {
+                if (((Bitmap) obj).isMutable()) {
+                    throw new BadParcelableException(
+                            "Encountered a mutable Bitmap in the Bundle at key : " + key);
+                }
+            } else {
+                throw new BadParcelableException(
+                        "Unsupported Parcelable type encountered in the Bundle: "
+                                + obj.getClass().getSimpleName());
+            }
+        }
+        Log.e(TAG, "validateResponseParams : Finished");
+    }
+
+    /**
+     * Validation of the inference request payload as described in {@link StateParams}
+     * description.
+     *
+     * @throws BadParcelableException when the bundle does not meet the read-only requirements.
+     */
+    public static void sanitizeStateParams(
+            @StateParams Bundle bundle) {
+        ensureValidBundle(bundle);
+
+        if (!bundle.hasFileDescriptors()) {
+            return; //safe to exit if there are no FDs and Binders
+        }
+
+        for (String key : bundle.keySet()) {
+            Object obj = bundle.get(key);
+            if (obj == null) {
+                /* Null value here could also mean deserializing a custom parcelable has failed,
+                 *  and since {@link Bundle} is marked as defusable in system-server - the
+                 * {@link ClassNotFoundException} exception is swallowed and `null` is returned
+                 * instead. We want to ensure cleanup of null entries in such case.
+                 */
+                bundle.putObject(key, null);
+                continue;
+            }
+            if (canMarshall(obj)) {
+                continue;
+            }
+
+            if (obj instanceof ParcelFileDescriptor) {
+                validatePfdReadOnly((ParcelFileDescriptor) obj);
+            } else {
+                throw new BadParcelableException(
+                        "Unsupported Parcelable type encountered in the Bundle: "
+                                + obj.getClass().getSimpleName());
+            }
+        }
+    }
+
+
+    public static IStreamingResponseCallback wrapWithValidation(
+            IStreamingResponseCallback streamingResponseCallback,
+            Executor resourceClosingExecutor) {
+        return new IStreamingResponseCallback.Stub() {
+            @Override
+            public void onNewContent(Bundle processedResult) throws RemoteException {
+                try {
+                    sanitizeResponseParams(processedResult);
+                    streamingResponseCallback.onNewContent(processedResult);
+                } finally {
+                    resourceClosingExecutor.execute(() -> tryCloseResource(processedResult));
+                }
+            }
+
+            @Override
+            public void onSuccess(Bundle resultBundle)
+                    throws RemoteException {
+                try {
+                    sanitizeResponseParams(resultBundle);
+                    streamingResponseCallback.onSuccess(resultBundle);
+                } finally {
+                    resourceClosingExecutor.execute(() -> tryCloseResource(resultBundle));
+                }
+            }
+
+            @Override
+            public void onFailure(int errorCode, String errorMessage,
+                    PersistableBundle errorParams) throws RemoteException {
+                streamingResponseCallback.onFailure(errorCode, errorMessage, errorParams);
+            }
+
+            @Override
+            public void onDataAugmentRequest(Bundle processedContent,
+                    RemoteCallback remoteCallback)
+                    throws RemoteException {
+                try {
+                    sanitizeResponseParams(processedContent);
+                    streamingResponseCallback.onDataAugmentRequest(processedContent,
+                            new RemoteCallback(
+                                    augmentedData -> {
+                                        try {
+                                            sanitizeInferenceParams(augmentedData);
+                                            remoteCallback.sendResult(augmentedData);
+                                        } finally {
+                                            resourceClosingExecutor.execute(
+                                                    () -> tryCloseResource(augmentedData));
+                                        }
+                                    }));
+                } finally {
+                    resourceClosingExecutor.execute(() -> tryCloseResource(processedContent));
+                }
+            }
+        };
+    }
+
+    public static IResponseCallback wrapWithValidation(IResponseCallback responseCallback,
+            Executor resourceClosingExecutor) {
+        return new IResponseCallback.Stub() {
+            @Override
+            public void onSuccess(Bundle resultBundle)
+                    throws RemoteException {
+                try {
+                    sanitizeResponseParams(resultBundle);
+                    responseCallback.onSuccess(resultBundle);
+                } finally {
+                    resourceClosingExecutor.execute(() -> tryCloseResource(resultBundle));
+                }
+            }
+
+            @Override
+            public void onFailure(int errorCode, String errorMessage,
+                    PersistableBundle errorParams) throws RemoteException {
+                responseCallback.onFailure(errorCode, errorMessage, errorParams);
+            }
+
+            @Override
+            public void onDataAugmentRequest(Bundle processedContent,
+                    RemoteCallback remoteCallback)
+                    throws RemoteException {
+                try {
+                    sanitizeResponseParams(processedContent);
+                    responseCallback.onDataAugmentRequest(processedContent, new RemoteCallback(
+                            augmentedData -> {
+                                try {
+                                    sanitizeInferenceParams(augmentedData);
+                                    remoteCallback.sendResult(augmentedData);
+                                } finally {
+                                    resourceClosingExecutor.execute(
+                                            () -> tryCloseResource(augmentedData));
+                                }
+                            }));
+                } finally {
+                    resourceClosingExecutor.execute(() -> tryCloseResource(processedContent));
+                }
+            }
+        };
+    }
+
+
+    public static ITokenInfoCallback wrapWithValidation(ITokenInfoCallback responseCallback) {
+        return new ITokenInfoCallback.Stub() {
+            @Override
+            public void onSuccess(TokenInfo tokenInfo) throws RemoteException {
+                responseCallback.onSuccess(tokenInfo);
+            }
+
+            @Override
+            public void onFailure(int errorCode, String errorMessage, PersistableBundle errorParams)
+                    throws RemoteException {
+                responseCallback.onFailure(errorCode, errorMessage, errorParams);
+            }
+        };
+    }
+
+    private static boolean canMarshall(Object obj) {
+        return obj instanceof byte[] || obj instanceof PersistableBundle
+                || PersistableBundle.isValidType(obj);
+    }
+
+    private static void ensureValidBundle(Bundle bundle) {
+        if (bundle == null) {
+            throw new IllegalArgumentException("Request passed is expected to be non-null");
+        }
+
+        if (bundle.hasBinders() != Bundle.STATUS_BINDERS_NOT_PRESENT) {
+            throw new BadParcelableException("Bundle should not contain IBinder objects.");
+        }
+    }
+
+    public static void validatePfdReadOnly(ParcelFileDescriptor pfd) {
+        if (pfd == null) {
+            return;
+        }
+        try {
+            int readMode = Os.fcntlInt(pfd.getFileDescriptor(), F_GETFL, 0) & O_ACCMODE;
+            if (readMode != O_RDONLY) {
+                throw new BadParcelableException(
+                        "Bundle contains a parcel file descriptor which is not read-only.");
+            }
+        } catch (ErrnoException e) {
+            throw new BadParcelableException(
+                    "Invalid File descriptor passed in the Bundle.", e);
+        }
+    }
+
+    public static void tryCloseResource(Bundle bundle) {
+        if (bundle == null || bundle.isEmpty() || !bundle.hasFileDescriptors()) {
+            return;
+        }
+
+        for (String key : bundle.keySet()) {
+            Object obj = bundle.get(key);
+
+            try {
+                // TODO(b/329898589) : This can be cleaned up after the flag passing is fixed.
+                if (obj instanceof ParcelFileDescriptor) {
+                    ((ParcelFileDescriptor) obj).close();
+                } else if (obj instanceof CursorWindow) {
+                    ((CursorWindow) obj).close();
+                } else if (obj instanceof SharedMemory) {
+                    // TODO(b/331796886) : Shared memory should honour parcelable flags.
+                    ((SharedMemory) obj).close();
+                }
+            } catch (Exception e) {
+                Log.e(TAG, "Error closing resource with key: " + key, e);
+            }
+        }
+    }
+}
diff --git a/services/core/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerService.java b/services/core/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerService.java
index a4a6221a..af339df 100644
--- a/services/core/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerService.java
+++ b/services/core/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerService.java
@@ -16,6 +16,12 @@
 
 package com.android.server.ondeviceintelligence;
 
+import static com.android.server.ondeviceintelligence.BundleUtil.sanitizeInferenceParams;
+import static com.android.server.ondeviceintelligence.BundleUtil.validatePfdReadOnly;
+import static com.android.server.ondeviceintelligence.BundleUtil.sanitizeStateParams;
+import static com.android.server.ondeviceintelligence.BundleUtil.wrapWithValidation;
+
+
 import android.Manifest;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
@@ -61,6 +67,7 @@
 import android.service.ondeviceintelligence.OnDeviceIntelligenceService;
 import android.service.ondeviceintelligence.OnDeviceSandboxedInferenceService;
 import android.text.TextUtils;
+import android.util.Log;
 import android.util.Slog;
 
 import com.android.internal.R;
@@ -72,8 +79,11 @@
 import com.android.server.SystemService;
 
 import java.io.FileDescriptor;
+import java.io.IOException;
 import java.util.Objects;
 import java.util.Set;
+import java.util.concurrent.Executor;
+import java.util.concurrent.Executors;
 
 /**
  * This is the system service for handling calls on the
@@ -99,6 +109,9 @@
     private static final boolean DEFAULT_SERVICE_ENABLED = true;
     private static final String NAMESPACE_ON_DEVICE_INTELLIGENCE = "ondeviceintelligence";
 
+    private final Executor resourceClosingExecutor = Executors.newCachedThreadPool();
+    private final Executor callbackExecutor = Executors.newCachedThreadPool();
+
     private final Context mContext;
     protected final Object mLock = new Object();
 
@@ -261,7 +274,7 @@
                 ensureRemoteIntelligenceServiceInitialized();
                 mRemoteOnDeviceIntelligenceService.run(
                         service -> service.requestFeatureDownload(Binder.getCallingUid(), feature,
-                                cancellationSignalFuture,
+                                wrapCancellationFuture(cancellationSignalFuture),
                                 downloadCallback));
             }
 
@@ -272,26 +285,35 @@
                     AndroidFuture cancellationSignalFuture,
                     ITokenInfoCallback tokenInfoCallback) throws RemoteException {
                 Slog.i(TAG, "OnDeviceIntelligenceManagerInternal requestTokenInfo");
-                Objects.requireNonNull(feature);
-                Objects.requireNonNull(request);
-                Objects.requireNonNull(tokenInfoCallback);
+                AndroidFuture<Void> result = null;
+                try {
+                    Objects.requireNonNull(feature);
+                    sanitizeInferenceParams(request);
+                    Objects.requireNonNull(tokenInfoCallback);
 
-                mContext.enforceCallingOrSelfPermission(
-                        Manifest.permission.USE_ON_DEVICE_INTELLIGENCE, TAG);
-                if (!mIsServiceEnabled) {
-                    Slog.w(TAG, "Service not available");
-                    tokenInfoCallback.onFailure(
-                            OnDeviceIntelligenceException.ON_DEVICE_INTELLIGENCE_SERVICE_UNAVAILABLE,
-                            "OnDeviceIntelligenceManagerService is unavailable",
-                            PersistableBundle.EMPTY);
+                    mContext.enforceCallingOrSelfPermission(
+                            Manifest.permission.USE_ON_DEVICE_INTELLIGENCE, TAG);
+                    if (!mIsServiceEnabled) {
+                        Slog.w(TAG, "Service not available");
+                        tokenInfoCallback.onFailure(
+                                OnDeviceIntelligenceException.ON_DEVICE_INTELLIGENCE_SERVICE_UNAVAILABLE,
+                                "OnDeviceIntelligenceManagerService is unavailable",
+                                PersistableBundle.EMPTY);
+                    }
+                    ensureRemoteInferenceServiceInitialized();
+
+                    result = mRemoteInferenceService.post(
+                            service -> service.requestTokenInfo(Binder.getCallingUid(), feature,
+                                    request,
+                                    wrapCancellationFuture(cancellationSignalFuture),
+                                    wrapWithValidation(tokenInfoCallback)));
+                    result.whenCompleteAsync((c, e) -> BundleUtil.tryCloseResource(request),
+                            resourceClosingExecutor);
+                } finally {
+                    if (result == null) {
+                        resourceClosingExecutor.execute(() -> BundleUtil.tryCloseResource(request));
+                    }
                 }
-                ensureRemoteInferenceServiceInitialized();
-
-                mRemoteInferenceService.run(
-                        service -> service.requestTokenInfo(Binder.getCallingUid(), feature,
-                                request,
-                                cancellationSignalFuture,
-                                tokenInfoCallback));
             }
 
             @Override
@@ -302,24 +324,36 @@
                     AndroidFuture processingSignalFuture,
                     IResponseCallback responseCallback)
                     throws RemoteException {
-                Slog.i(TAG, "OnDeviceIntelligenceManagerInternal processRequest");
-                Objects.requireNonNull(feature);
-                Objects.requireNonNull(responseCallback);
-                mContext.enforceCallingOrSelfPermission(
-                        Manifest.permission.USE_ON_DEVICE_INTELLIGENCE, TAG);
-                if (!mIsServiceEnabled) {
-                    Slog.w(TAG, "Service not available");
-                    responseCallback.onFailure(
-                            OnDeviceIntelligenceException.PROCESSING_ERROR_SERVICE_UNAVAILABLE,
-                            "OnDeviceIntelligenceManagerService is unavailable",
-                            PersistableBundle.EMPTY);
+                AndroidFuture<Void> result = null;
+                try {
+                    Slog.i(TAG, "OnDeviceIntelligenceManagerInternal processRequest");
+                    Objects.requireNonNull(feature);
+                    sanitizeInferenceParams(request);
+                    Objects.requireNonNull(responseCallback);
+                    mContext.enforceCallingOrSelfPermission(
+                            Manifest.permission.USE_ON_DEVICE_INTELLIGENCE, TAG);
+                    if (!mIsServiceEnabled) {
+                        Slog.w(TAG, "Service not available");
+                        responseCallback.onFailure(
+                                OnDeviceIntelligenceException.PROCESSING_ERROR_SERVICE_UNAVAILABLE,
+                                "OnDeviceIntelligenceManagerService is unavailable",
+                                PersistableBundle.EMPTY);
+                    }
+                    ensureRemoteInferenceServiceInitialized();
+                    result = mRemoteInferenceService.post(
+                            service -> service.processRequest(Binder.getCallingUid(), feature,
+                                    request,
+                                    requestType,
+                                    wrapCancellationFuture(cancellationSignalFuture),
+                                    wrapProcessingFuture(processingSignalFuture),
+                                    wrapWithValidation(responseCallback, resourceClosingExecutor)));
+                    result.whenCompleteAsync((c, e) -> BundleUtil.tryCloseResource(request),
+                            resourceClosingExecutor);
+                } finally {
+                    if (result == null) {
+                        resourceClosingExecutor.execute(() -> BundleUtil.tryCloseResource(request));
+                    }
                 }
-                ensureRemoteInferenceServiceInitialized();
-                mRemoteInferenceService.run(
-                        service -> service.processRequest(Binder.getCallingUid(), feature, request,
-                                requestType,
-                                cancellationSignalFuture, processingSignalFuture,
-                                responseCallback));
             }
 
             @Override
@@ -329,24 +363,36 @@
                     AndroidFuture cancellationSignalFuture,
                     AndroidFuture processingSignalFuture,
                     IStreamingResponseCallback streamingCallback) throws RemoteException {
-                Slog.i(TAG, "OnDeviceIntelligenceManagerInternal processRequestStreaming");
-                Objects.requireNonNull(feature);
-                Objects.requireNonNull(streamingCallback);
-                mContext.enforceCallingOrSelfPermission(
-                        Manifest.permission.USE_ON_DEVICE_INTELLIGENCE, TAG);
-                if (!mIsServiceEnabled) {
-                    Slog.w(TAG, "Service not available");
-                    streamingCallback.onFailure(
-                            OnDeviceIntelligenceException.PROCESSING_ERROR_SERVICE_UNAVAILABLE,
-                            "OnDeviceIntelligenceManagerService is unavailable",
-                            PersistableBundle.EMPTY);
+                AndroidFuture<Void> result = null;
+                try {
+                    Slog.i(TAG, "OnDeviceIntelligenceManagerInternal processRequestStreaming");
+                    Objects.requireNonNull(feature);
+                    sanitizeInferenceParams(request);
+                    Objects.requireNonNull(streamingCallback);
+                    mContext.enforceCallingOrSelfPermission(
+                            Manifest.permission.USE_ON_DEVICE_INTELLIGENCE, TAG);
+                    if (!mIsServiceEnabled) {
+                        Slog.w(TAG, "Service not available");
+                        streamingCallback.onFailure(
+                                OnDeviceIntelligenceException.PROCESSING_ERROR_SERVICE_UNAVAILABLE,
+                                "OnDeviceIntelligenceManagerService is unavailable",
+                                PersistableBundle.EMPTY);
+                    }
+                    ensureRemoteInferenceServiceInitialized();
+                    result = mRemoteInferenceService.post(
+                            service -> service.processRequestStreaming(Binder.getCallingUid(),
+                                    feature,
+                                    request, requestType,
+                                    wrapCancellationFuture(cancellationSignalFuture),
+                                    wrapProcessingFuture(processingSignalFuture),
+                                    streamingCallback));
+                    result.whenCompleteAsync((c, e) -> BundleUtil.tryCloseResource(request),
+                            resourceClosingExecutor);
+                } finally {
+                    if (result == null) {
+                        resourceClosingExecutor.execute(() -> BundleUtil.tryCloseResource(request));
+                    }
                 }
-                ensureRemoteInferenceServiceInitialized();
-                mRemoteInferenceService.run(
-                        service -> service.processRequestStreaming(Binder.getCallingUid(), feature,
-                                request, requestType,
-                                cancellationSignalFuture, processingSignalFuture,
-                                streamingCallback));
             }
 
             @Override
@@ -372,9 +418,9 @@
                             public void onConnected(
                                     @NonNull IOnDeviceIntelligenceService service) {
                                 try {
-                                    service.ready();
                                     service.registerRemoteServices(
                                             getRemoteProcessingService());
+                                    service.ready();
                                 } catch (RemoteException ex) {
                                     Slog.w(TAG, "Failed to send connected event", ex);
                                 }
@@ -391,10 +437,24 @@
             public void updateProcessingState(
                     Bundle processingState,
                     IProcessingUpdateStatusCallback callback) {
-                ensureRemoteInferenceServiceInitialized();
-                mRemoteInferenceService.run(
-                        service -> service.updateProcessingState(
-                                processingState, callback));
+                callbackExecutor.execute(() -> {
+                    AndroidFuture<Void> result = null;
+                    try {
+                        sanitizeStateParams(processingState);
+                        ensureRemoteInferenceServiceInitialized();
+                        result = mRemoteInferenceService.post(
+                                service -> service.updateProcessingState(
+                                        processingState, callback));
+                        result.whenCompleteAsync(
+                                (c, e) -> BundleUtil.tryCloseResource(processingState),
+                                resourceClosingExecutor);
+                    } finally {
+                        if (result == null) {
+                            resourceClosingExecutor.execute(
+                                    () -> BundleUtil.tryCloseResource(processingState));
+                        }
+                    }
+                });
             }
         };
     }
@@ -415,7 +475,7 @@
                                 try {
                                     ensureRemoteIntelligenceServiceInitialized();
                                     mRemoteOnDeviceIntelligenceService.run(
-                                            intelligenceService -> intelligenceService.notifyInferenceServiceConnected());
+                                            IOnDeviceIntelligenceService::notifyInferenceServiceConnected);
                                     service.registerRemoteStorageService(
                                             getIRemoteStorageService());
                                 } catch (RemoteException ex) {
@@ -434,18 +494,49 @@
             public void getReadOnlyFileDescriptor(
                     String filePath,
                     AndroidFuture<ParcelFileDescriptor> future) {
+                ensureRemoteIntelligenceServiceInitialized();
+                AndroidFuture<ParcelFileDescriptor> pfdFuture = new AndroidFuture<>();
                 mRemoteOnDeviceIntelligenceService.run(
                         service -> service.getReadOnlyFileDescriptor(
-                                filePath, future));
+                                filePath, pfdFuture));
+                pfdFuture.whenCompleteAsync((pfd, error) -> {
+                    try {
+                        if (error != null) {
+                            future.completeExceptionally(error);
+                        } else {
+                            validatePfdReadOnly(pfd);
+                            future.complete(pfd);
+                        }
+                    } finally {
+                        tryClosePfd(pfd);
+                    }
+                }, callbackExecutor);
             }
 
             @Override
             public void getReadOnlyFeatureFileDescriptorMap(
                     Feature feature,
                     RemoteCallback remoteCallback) {
+                ensureRemoteIntelligenceServiceInitialized();
                 mRemoteOnDeviceIntelligenceService.run(
                         service -> service.getReadOnlyFeatureFileDescriptorMap(
-                                feature, remoteCallback));
+                                feature,
+                                new RemoteCallback(result -> callbackExecutor.execute(() -> {
+                                    try {
+                                        if (result == null) {
+                                            remoteCallback.sendResult(null);
+                                        }
+                                        for (String key : result.keySet()) {
+                                            ParcelFileDescriptor pfd = result.getParcelable(key,
+                                                    ParcelFileDescriptor.class);
+                                            validatePfdReadOnly(pfd);
+                                        }
+                                        remoteCallback.sendResult(result);
+                                    } finally {
+                                        resourceClosingExecutor.execute(
+                                                () -> BundleUtil.tryCloseResource(result));
+                                    }
+                                }))));
             }
         };
     }
@@ -461,7 +552,8 @@
             ServiceInfo serviceInfo = AppGlobals.getPackageManager().getServiceInfo(
                     serviceComponent,
                     PackageManager.MATCH_DIRECT_BOOT_AWARE
-                            | PackageManager.MATCH_DIRECT_BOOT_UNAWARE, 0);
+                            | PackageManager.MATCH_DIRECT_BOOT_UNAWARE,
+                    UserHandle.SYSTEM.getIdentifier());
             if (serviceInfo != null) {
                 if (!checkIsolated) {
                     checkServiceRequiresPermission(serviceInfo,
@@ -596,4 +688,57 @@
 
         throw new SecurityException(message + ": Only shell user can call it");
     }
+
+    private AndroidFuture<IBinder> wrapCancellationFuture(
+            AndroidFuture future) {
+        if (future == null) {
+            return null;
+        }
+        AndroidFuture<IBinder> cancellationFuture = new AndroidFuture<>();
+        cancellationFuture.whenCompleteAsync((c, e) -> {
+            if (e != null) {
+                Log.e(TAG, "Error forwarding ICancellationSignal to manager layer", e);
+                future.completeExceptionally(e);
+            } else {
+                future.complete(new ICancellationSignal.Stub() {
+                    @Override
+                    public void cancel() throws RemoteException {
+                        ICancellationSignal.Stub.asInterface(c).cancel();
+                    }
+                });
+            }
+        });
+        return cancellationFuture;
+    }
+
+    private AndroidFuture<IBinder> wrapProcessingFuture(
+            AndroidFuture future) {
+        if (future == null) {
+            return null;
+        }
+        AndroidFuture<IBinder> processingSignalFuture = new AndroidFuture<>();
+        processingSignalFuture.whenCompleteAsync((c, e) -> {
+            if (e != null) {
+                future.completeExceptionally(e);
+            } else {
+                future.complete(new IProcessingSignal.Stub() {
+                    @Override
+                    public void sendSignal(PersistableBundle actionParams) throws RemoteException {
+                        IProcessingSignal.Stub.asInterface(c).sendSignal(actionParams);
+                    }
+                });
+            }
+        });
+        return processingSignalFuture;
+    }
+
+    private static void tryClosePfd(ParcelFileDescriptor pfd) {
+        if (pfd != null) {
+            try {
+                pfd.close();
+            } catch (IOException e) {
+                Log.e(TAG, "Failed to close parcel file descriptor ", e);
+            }
+        }
+    }
 }