Add APIs to update and listen to deviceId changes in Context

- Adds a updateDeviceId @hide API that can be used by the system
  to update the device association of a Context that is not
  explicitly created as a Device Context.
- Adds a listener that would be notified whenever the deviceId
  of the Context has changed.
- Adds an isDeviceContext() API that can be used to determine
  if the device ID returned by getDeviceId() is reliable for
  this Context instance.

These APIs will enable implicit deviceId association for Contexts
that are not Device Contexts.

Test: atest DeviceAssociationTest
Bug: 253201821
Change-Id: I8ef7a5f7a82ee341fb98236a01940b4be1e4fb23

6 files changed, 281 insertions, 10 deletions

diff --git a/core/api/current.txt b/core/api/current.txt
index 1ba99f995ec7..38db3ee16d1e 100644
--- a/core/api/current.txt
+++ b/core/api/current.txt
@@ -9851,6 +9851,7 @@ package android.content {
     method @Deprecated public abstract int getWallpaperDesiredMinimumHeight();
     method @Deprecated public abstract int getWallpaperDesiredMinimumWidth();
     method public abstract void grantUriPermission(String, android.net.Uri, int);
+    method public boolean isDeviceContext();
     method public abstract boolean isDeviceProtectedStorage();
     method public boolean isRestricted();
     method public boolean isUiContext();
@@ -9866,6 +9867,7 @@ package android.content {
     method public abstract android.database.sqlite.SQLiteDatabase openOrCreateDatabase(String, int, android.database.sqlite.SQLiteDatabase.CursorFactory, @Nullable android.database.DatabaseErrorHandler);
     method @Deprecated public abstract android.graphics.drawable.Drawable peekWallpaper();
     method public void registerComponentCallbacks(android.content.ComponentCallbacks);
+    method public void registerDeviceIdChangeListener(@NonNull java.util.concurrent.Executor, @NonNull java.util.function.IntConsumer);
     method @Nullable public abstract android.content.Intent registerReceiver(@Nullable android.content.BroadcastReceiver, android.content.IntentFilter);
     method @Nullable public abstract android.content.Intent registerReceiver(@Nullable android.content.BroadcastReceiver, android.content.IntentFilter, int);
     method @Nullable public abstract android.content.Intent registerReceiver(android.content.BroadcastReceiver, android.content.IntentFilter, @Nullable String, @Nullable android.os.Handler);
@@ -9905,6 +9907,7 @@ package android.content {
     method public abstract boolean stopService(android.content.Intent);
     method public abstract void unbindService(@NonNull android.content.ServiceConnection);
     method public void unregisterComponentCallbacks(android.content.ComponentCallbacks);
+    method public void unregisterDeviceIdChangeListener(@NonNull java.util.function.IntConsumer);
     method public abstract void unregisterReceiver(android.content.BroadcastReceiver);
     method public void updateServiceGroup(@NonNull android.content.ServiceConnection, int, int);
     field public static final String ACCESSIBILITY_SERVICE = "accessibility";
diff --git a/core/java/android/app/ContextImpl.java b/core/java/android/app/ContextImpl.java
index 39f71539b380..48a4daeddbcd 100644
--- a/core/java/android/app/ContextImpl.java
+++ b/core/java/android/app/ContextImpl.java
@@ -21,12 +21,14 @@ import static android.content.pm.PackageManager.PERMISSION_GRANTED;
 import static android.os.StrictMode.vmIncorrectContextUseEnabled;
 import static android.view.WindowManager.LayoutParams.WindowType;
 
+import android.annotation.CallbackExecutor;
 import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.UiContext;
 import android.companion.virtual.VirtualDevice;
 import android.companion.virtual.VirtualDeviceManager;
+import android.companion.virtual.VirtualDeviceParams;
 import android.compat.annotation.UnsupportedAppUsage;
 import android.content.AttributionSource;
 import android.content.AutofillOptions;
@@ -121,6 +123,7 @@ import java.util.List;
 import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.Executor;
+import java.util.function.IntConsumer;
 
 class ReceiverRestrictedContext extends ContextWrapper {
     @UnsupportedAppUsage
@@ -257,6 +260,13 @@ class ContextImpl extends Context {
     /** @see Context#isConfigurationContext() */
     private boolean mIsConfigurationBasedContext;
 
+    /**
+     *  Indicates that this context was created with an explicit device ID association via
+     *  Context#createDeviceContext and under no circumstances will it ever change, even if
+     *  this context is not associated with a display id, or if the associated display id changes.
+     */
+    private boolean mIsExplicitDeviceId = false;
+
     @UnsupportedAppUsage(maxTargetSdk = Build.VERSION_CODES.P, trackingBug = 115609023)
     private final int mFlags;
 
@@ -372,6 +382,24 @@ class ContextImpl extends Context {
     @ServiceInitializationState
     final int[] mServiceInitializationStateArray = new int[mServiceCache.length];
 
+    private final Object mDeviceIdListenerLock = new Object();
+    /**
+     * List of listeners for deviceId changes and their associated Executor.
+     * List is lazy-initialized on first registration
+     */
+    @GuardedBy("mDeviceIdListenerLock")
+    @Nullable
+    private ArrayList<DeviceIdChangeListenerDelegate> mDeviceIdChangeListeners;
+
+    private static class DeviceIdChangeListenerDelegate {
+        final @NonNull IntConsumer mListener;
+        final @NonNull Executor mExecutor;
+        DeviceIdChangeListenerDelegate(IntConsumer listener, Executor executor) {
+            mListener = listener;
+            mExecutor = executor;
+        }
+    }
+
     @UnsupportedAppUsage
     static ContextImpl getImpl(Context context) {
         Context nextContext;
@@ -2699,15 +2727,7 @@ class ContextImpl extends Context {
 
     @Override
     public @NonNull Context createDeviceContext(int deviceId) {
-        boolean validDeviceId = deviceId == VirtualDeviceManager.DEVICE_ID_DEFAULT;
-        if (deviceId > VirtualDeviceManager.DEVICE_ID_DEFAULT) {
-            VirtualDeviceManager vdm = getSystemService(VirtualDeviceManager.class);
-            if (vdm != null) {
-                List<VirtualDevice> virtualDevices = vdm.getVirtualDevices();
-                validDeviceId = virtualDevices.stream().anyMatch(d -> d.getDeviceId() == deviceId);
-            }
-        }
-        if (!validDeviceId) {
+        if (!isValidDeviceId(deviceId)) {
             throw new IllegalArgumentException(
                     "Not a valid ID of the default device or any virtual device: " + deviceId);
         }
@@ -2718,9 +2738,35 @@ class ContextImpl extends Context {
                 mSplitName, mToken, mUser, mFlags, mClassLoader, null);
 
         context.mDeviceId = deviceId;
+        context.mIsExplicitDeviceId = true;
         return context;
     }
 
+    /**
+     * Checks whether the passed {@code deviceId} is valid or not.
+     * {@link VirtualDeviceManager#DEVICE_ID_DEFAULT} is valid as it is the ID of the default
+     * device when no additional virtual devices exist. If {@code deviceId} is the id of
+     * a virtual device, it should correspond to a virtual device created by
+     * {@link VirtualDeviceManager#createVirtualDevice(int, VirtualDeviceParams)}.
+     */
+    private boolean isValidDeviceId(int deviceId) {
+        if (deviceId == VirtualDeviceManager.DEVICE_ID_DEFAULT) {
+            return true;
+        }
+        if (deviceId > VirtualDeviceManager.DEVICE_ID_DEFAULT) {
+            VirtualDeviceManager vdm = getSystemService(VirtualDeviceManager.class);
+            if (vdm != null) {
+                List<VirtualDevice> virtualDevices = vdm.getVirtualDevices();
+                for (int i = 0; i < virtualDevices.size(); i++) {
+                    if (virtualDevices.get(i).getDeviceId() == deviceId) {
+                        return true;
+                    }
+                }
+            }
+        }
+        return false;
+    }
+
     @NonNull
     @Override
     public WindowContext createWindowContext(@WindowType int type,
@@ -2965,6 +3011,21 @@ class ContextImpl extends Context {
         if (mContextType == CONTEXT_TYPE_NON_UI) {
             mContextType = CONTEXT_TYPE_DISPLAY_CONTEXT;
         }
+        // TODO(b/253201821): Update deviceId when display is updated.
+    }
+
+    @Override
+    public void updateDeviceId(int updatedDeviceId) {
+        if (!isValidDeviceId(updatedDeviceId)) {
+            throw new IllegalArgumentException(
+                    "Not a valid ID of the default device or any virtual device: " + mDeviceId);
+        }
+        if (mIsExplicitDeviceId) {
+            throw new UnsupportedOperationException(
+                    "Cannot update device ID on a Context created with createDeviceContext()");
+        }
+        mDeviceId = updatedDeviceId;
+        notifyOnDeviceChangedListeners(updatedDeviceId);
     }
 
     @Override
@@ -2973,6 +3034,69 @@ class ContextImpl extends Context {
     }
 
     @Override
+    public boolean isDeviceContext() {
+        return mIsExplicitDeviceId || isAssociatedWithDisplay();
+    }
+
+    @Override
+    public void registerDeviceIdChangeListener(@NonNull @CallbackExecutor Executor executor,
+            @NonNull IntConsumer listener) {
+        Objects.requireNonNull(executor, "executor cannot be null");
+        Objects.requireNonNull(listener, "listener cannot be null");
+
+        synchronized (mDeviceIdListenerLock) {
+            if (getDeviceIdListener(listener) != null) {
+                throw new IllegalArgumentException(
+                        "attempt to call registerDeviceIdChangeListener() "
+                                + "on a previously registered listener");
+            }
+            // lazy initialization
+            if (mDeviceIdChangeListeners == null) {
+                mDeviceIdChangeListeners = new ArrayList<>();
+            }
+            mDeviceIdChangeListeners.add(new DeviceIdChangeListenerDelegate(listener, executor));
+        }
+    }
+
+    @Override
+    public void unregisterDeviceIdChangeListener(@NonNull IntConsumer listener) {
+        Objects.requireNonNull(listener, "listener cannot be null");
+        synchronized (mDeviceIdListenerLock) {
+            DeviceIdChangeListenerDelegate listenerToRemove = getDeviceIdListener(listener);
+            if (listenerToRemove != null) {
+                mDeviceIdChangeListeners.remove(listenerToRemove);
+            }
+        }
+    }
+
+    @GuardedBy("mDeviceIdListenerLock")
+    @Nullable
+    private DeviceIdChangeListenerDelegate getDeviceIdListener(
+            @Nullable IntConsumer listener) {
+        if (mDeviceIdChangeListeners == null) {
+            return null;
+        }
+        for (int i = 0; i < mDeviceIdChangeListeners.size(); i++) {
+            DeviceIdChangeListenerDelegate delegate = mDeviceIdChangeListeners.get(i);
+            if (delegate.mListener == listener) {
+                return delegate;
+            }
+        }
+        return null;
+    }
+
+    private void notifyOnDeviceChangedListeners(int deviceId) {
+        synchronized (mDeviceIdListenerLock) {
+            if (mDeviceIdChangeListeners != null) {
+                for (DeviceIdChangeListenerDelegate delegate : mDeviceIdChangeListeners) {
+                    delegate.mExecutor.execute(() ->
+                            delegate.mListener.accept(deviceId));
+                }
+            }
+        }
+    }
+
+    @Override
     public DisplayAdjustments getDisplayAdjustments(int displayId) {
         return mResources.getDisplayAdjustments();
     }
@@ -3227,6 +3351,8 @@ class ContextImpl extends Context {
             opPackageName = container.mOpPackageName;
             setResources(container.mResources);
             mDisplay = container.mDisplay;
+            mDeviceId = container.mDeviceId;
+            mIsExplicitDeviceId = container.mIsExplicitDeviceId;
             mForceDisplayOverrideInResources = container.mForceDisplayOverrideInResources;
             mIsConfigurationBasedContext = container.mIsConfigurationBasedContext;
             mContextType = container.mContextType;
diff --git a/core/java/android/content/Context.java b/core/java/android/content/Context.java
index 382e2bb6ee43..7e6574197ed2 100644
--- a/core/java/android/content/Context.java
+++ b/core/java/android/content/Context.java
@@ -109,6 +109,7 @@ import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.Executor;
 import java.util.function.Consumer;
+import java.util.function.IntConsumer;
 
 /**
  * Interface to global information about an application environment.  This is
@@ -6906,6 +6907,10 @@ public abstract class Context {
      * {@link android.companion.virtual.VirtualDeviceManager#DEVICE_ID_DEFAULT}. Similarly,
      * applications running on the default device may access the functionality of virtual devices.
      * </p>
+     * <p>
+     * Note that the newly created instance will be associated with the same display as the parent
+     * Context, regardless of the device ID passed here.
+     * </p>
      * @param deviceId The ID of the device to associate with this context.
      * @return A context associated with the given device ID.
      *
@@ -7241,20 +7246,116 @@ public abstract class Context {
     public abstract void updateDisplay(int displayId);
 
     /**
-     * Get the device ID this context is associated with. Applications can use this method to
+     * Updates the device ID association of this Context. Since a Context created with
+     * {@link #createDeviceContext} cannot change its device association, this method must
+     * not be called for instances created with {@link #createDeviceContext}.
+     *
+     * @param deviceId The new device ID to assign to this Context.
+     * @throws UnsupportedOperationException if the method is called on an instance that was
+     *         created with {@link Context#createDeviceContext(int)}
+     * @throws IllegalArgumentException if the given device ID is not a valid ID of the default
+     *         device or a virtual device.
+     *
+     * @see #isDeviceContext()
+     * @see #createDeviceContext(int)
+     * @hide
+     */
+    public void updateDeviceId(int deviceId) {
+        throw new RuntimeException("Not implemented. Must override in a subclass.");
+    }
+
+    /**
+     * Gets the device ID this context is associated with. Applications can use this method to
      * determine whether they are running on a virtual device and identify that device.
      *
      * The device ID of the host device is
      * {@link android.companion.virtual.VirtualDeviceManager#DEVICE_ID_DEFAULT}
      *
+     * <p>
+     * If the underlying device ID is changed by the system, for example, when an
+     * {@link Activity} is moved to a different virtual device, applications can register to listen
+     * to changes by calling
+     * {@link Context#registerDeviceIdChangeListener(Executor, IntConsumer)}.
+     * </p>
+     *
+     * <p>
+     * This method will only return a reliable value for this instance if
+     * {@link Context#isDeviceContext()} is {@code true}. The system can assign an arbitrary device
+     * id value for Contexts not logically associated with a device.
+     * </p>
+     *
      * @return the ID of the device this context is associated with.
+     * @see #isDeviceContext()
      * @see #createDeviceContext(int)
+     * @see #registerDeviceIdChangeListener(Executor, IntConsumer)
      */
     public int getDeviceId() {
         throw new RuntimeException("Not implemented. Must override in a subclass.");
     }
 
     /**
+     * Indicates whether the value of {@link Context#getDeviceId()} can be relied upon for
+     * this instance. It will return {@code true} for Contexts created by
+     * {@link Context#createDeviceContext(int)}, as well as for UI and Display Contexts.
+     * <p>
+     * Contexts created with {@link Context#createDeviceContext(int)} will have an explicit
+     * device association, which will never change. UI Contexts and Display Contexts are
+     * already associated with a display, so if the device association is not explicitly
+     * given, {@link Context#getDeviceId()} will return the ID of the device associated with
+     * the associated display. The system can assign an arbitrary device id value for Contexts not
+     * logically associated with a device.
+     * </p>
+     *
+     * @return {@code true} if {@link Context#getDeviceId()} is reliable, {@code false} otherwise.
+     *
+     * @see #createDeviceContext(int)
+     * @see #getDeviceId()}
+     * @see #createDisplayContext(Display)
+     * @see #isUiContext()
+     */
+
+    public boolean isDeviceContext() {
+        throw new RuntimeException("Not implemented. Must override in a subclass.");
+    }
+
+    /**
+     * Adds a new device ID changed listener to the {@code Context}, which will be called when
+     * the device association is changed by the system.
+     * <p>
+     * The callback can be called when an app is moved to a different device and the {@code Context}
+     * is not explicily associated with a specific device.
+     * </p>
+     * <p> When an application receives a device id update callback, this Context is guaranteed to
+     * also have an updated display ID(if any) and {@link Configuration}.
+     * <p/>
+     * @param executor The Executor on whose thread to execute the callbacks of the {@code listener}
+     *                 object.
+     * @param listener The listener {@code IntConsumer} to call which will receive the updated
+     *                 device ID.
+     *
+     * @see Context#isDeviceContext()
+     * @see Context#getDeviceId()
+     * @see Context#createDeviceContext(int)
+     */
+    public void registerDeviceIdChangeListener(@NonNull @CallbackExecutor Executor executor,
+            @NonNull IntConsumer listener) {
+        throw new RuntimeException("Not implemented. Must override in a subclass.");
+    }
+
+    /**
+     * Removes a device ID changed listener from the Context. It's a no-op if
+     * the listener is not already registered.
+     *
+     * @param listener The {@code Consumer} to remove.
+     *
+     * @see #getDeviceId()
+     * @see #registerDeviceIdChangeListener(Executor, IntConsumer)
+     */
+    public void unregisterDeviceIdChangeListener(@NonNull IntConsumer listener) {
+        throw new RuntimeException("Not implemented. Must override in a subclass.");
+    }
+
+    /**
      * Indicates whether this Context is restricted.
      *
      * @return {@code true} if this Context is restricted, {@code false} otherwise.
diff --git a/core/java/android/content/ContextWrapper.java b/core/java/android/content/ContextWrapper.java
index a1646a172521..0a32dd78092f 100644
--- a/core/java/android/content/ContextWrapper.java
+++ b/core/java/android/content/ContextWrapper.java
@@ -16,6 +16,7 @@
 
 package android.content;
 
+import android.annotation.CallbackExecutor;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.RequiresPermission;
@@ -61,6 +62,7 @@ import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
 import java.util.concurrent.Executor;
+import java.util.function.IntConsumer;
 
 /**
  * Proxying implementation of Context that simply delegates all of its calls to
@@ -1171,12 +1173,36 @@ public class ContextWrapper extends Context {
         mBase.updateDisplay(displayId);
     }
 
+    /**
+     * @hide
+     */
+    @Override
+    public void updateDeviceId(int deviceId) {
+        mBase.updateDeviceId(deviceId);
+    }
+
     @Override
     public int getDeviceId() {
         return mBase.getDeviceId();
     }
 
     @Override
+    public boolean isDeviceContext() {
+        return mBase.isDeviceContext();
+    }
+
+    @Override
+    public void registerDeviceIdChangeListener(@NonNull @CallbackExecutor Executor executor,
+            @NonNull IntConsumer listener) {
+        mBase.registerDeviceIdChangeListener(executor, listener);
+    }
+
+    @Override
+    public void unregisterDeviceIdChangeListener(@NonNull IntConsumer listener) {
+        mBase.unregisterDeviceIdChangeListener(listener);
+    }
+
+    @Override
     public Context createDeviceProtectedStorageContext() {
         return mBase.createDeviceProtectedStorageContext();
     }
diff --git a/core/java/com/android/internal/policy/DecorContext.java b/core/java/com/android/internal/policy/DecorContext.java
index 134a91710c0b..63785f270b59 100644
--- a/core/java/com/android/internal/policy/DecorContext.java
+++ b/core/java/com/android/internal/policy/DecorContext.java
@@ -139,6 +139,15 @@ public class DecorContext extends ContextThemeWrapper {
     }
 
     @Override
+    public boolean isDeviceContext() {
+        Context context = mContext.get();
+        if (context != null) {
+            return context.isDeviceContext();
+        }
+        return false;
+    }
+
+    @Override
     public boolean isConfigurationContext() {
         Context context = mContext.get();
         if (context != null) {
diff --git a/test-mock/src/android/test/mock/MockContext.java b/test-mock/src/android/test/mock/MockContext.java
index 8fc8c7d162f4..b63fbe679e23 100644
--- a/test-mock/src/android/test/mock/MockContext.java
+++ b/test-mock/src/android/test/mock/MockContext.java
@@ -887,6 +887,12 @@ public class MockContext extends Context {
         throw new UnsupportedOperationException();
     }
 
+    /** @hide */
+    @Override
+    public void updateDeviceId(int deviceId) {
+        throw new UnsupportedOperationException();
+    }
+
     @Override
     public int getDeviceId() {
         throw new UnsupportedOperationException();