Fix race condition between VDM and ProxyManager when registering a proxy

Exo may register a proxy before VDM begins tracking the streamed,
but after this app's AccessibilityManagerClient is created in AMS.
This race conditions means the app is not notified that
it is proxied, since VDM#getDevicesForUid will not return the proxy
device id when called.

Add a AppsOnVirtualDeviceListener to the ProxyManager. If the running
apps are changed on any virtual device, force notifying the apps if
any of these belong to a proxied device. We add `forceUpdate` since the
last and current client state may be the same (no proxy connections
have changed) in this case.

Also move any apps that are no longer proxied back to the default device
if needed when updating the client device ids, since Exo will use the
same device and launch activities on the same display. (We were
previously resetting clients only when a device id
had no registered proxies, but Exo should be able to keep the proxy
registered)

Bug: 286587811
Test: manual with Exo, atest ProxyManagerTest, atest
AccessibilityDisplayProxyTest

Change-Id: I1a1b600c9e71d7baf1f9647475ac63580f7f01c2
diff --git a/services/accessibility/Android.bp b/services/accessibility/Android.bp
index bf8a9af..e9bb763 100644
--- a/services/accessibility/Android.bp
+++ b/services/accessibility/Android.bp
@@ -27,4 +27,20 @@
         "services.core",
         "androidx.annotation_annotation",
     ],
+    static_libs: [
+        "com_android_server_accessibility_flags_lib",
+    ],
+}
+
+aconfig_declarations {
+    name: "com_android_server_accessibility_flags",
+    package: "com.android.server.accessibility",
+    srcs: [
+        "accessibility.aconfig",
+    ],
+}
+
+java_aconfig_library {
+    name: "com_android_server_accessibility_flags_lib",
+    aconfig_declarations: "com_android_server_accessibility_flags",
 }
diff --git a/services/accessibility/accessibility.aconfig b/services/accessibility/accessibility.aconfig
new file mode 100644
index 0000000..b5fc2b6
--- /dev/null
+++ b/services/accessibility/accessibility.aconfig
@@ -0,0 +1,7 @@
+package: "com.android.server.accessibility"
+flag {
+    name: "proxy_use_apps_on_virtual_device_listener"
+    namespace: "accessibility"
+    description: "Fixes race condition described in b/286587811"
+    bug: "286587811"
+}
diff --git a/services/accessibility/java/com/android/server/accessibility/ProxyManager.java b/services/accessibility/java/com/android/server/accessibility/ProxyManager.java
index 119f575..ed77476 100644
--- a/services/accessibility/java/com/android/server/accessibility/ProxyManager.java
+++ b/services/accessibility/java/com/android/server/accessibility/ProxyManager.java
@@ -33,6 +33,7 @@
 import android.os.IBinder;
 import android.os.RemoteCallbackList;
 import android.os.RemoteException;
+import android.util.ArraySet;
 import android.util.IntArray;
 import android.util.Slog;
 import android.util.SparseArray;
@@ -42,6 +43,7 @@
 import android.view.accessibility.AccessibilityManager;
 import android.view.accessibility.IAccessibilityManagerClient;
 
+import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.util.IntPair;
 import com.android.server.LocalServices;
 import com.android.server.companion.virtual.VirtualDeviceManagerInternal;
@@ -96,6 +98,9 @@
 
     private final SystemSupport mSystemSupport;
 
+    private VirtualDeviceManagerInternal.AppsOnVirtualDeviceListener
+            mAppsOnVirtualDeviceListener;
+
     /**
      * Callbacks into AccessibilityManagerService.
      */
@@ -174,6 +179,16 @@
 
         synchronized (mLock) {
             mProxyA11yServiceConnections.put(displayId, connection);
+            if (Flags.proxyUseAppsOnVirtualDeviceListener()) {
+                if (mAppsOnVirtualDeviceListener == null) {
+                    mAppsOnVirtualDeviceListener = allRunningUids ->
+                            notifyProxyOfRunningAppsChange(allRunningUids);
+                    final VirtualDeviceManagerInternal localVdm = getLocalVdm();
+                    if (localVdm != null) {
+                        localVdm.registerAppsOnVirtualDeviceListener(mAppsOnVirtualDeviceListener);
+                    }
+                }
+            }
         }
 
         // If the client dies, make sure to remove the connection.
@@ -276,11 +291,21 @@
                         }
                     }
                 });
-        // If there isn't an existing proxy for the device id, reset clients. Resetting
+        // If there isn't an existing proxy for the device id, reset app clients. Resetting
         // will usually happen, since in most cases there will only be one proxy for a
         // device.
         if (!isProxyedDeviceId(deviceId)) {
             synchronized (mLock) {
+                if (Flags.proxyUseAppsOnVirtualDeviceListener()) {
+                    if (mProxyA11yServiceConnections.size() == 0) {
+                        final VirtualDeviceManagerInternal localVdm = getLocalVdm();
+                        if (localVdm != null && mAppsOnVirtualDeviceListener != null) {
+                            localVdm.unregisterAppsOnVirtualDeviceListener(
+                                    mAppsOnVirtualDeviceListener);
+                            mAppsOnVirtualDeviceListener = null;
+                        }
+                    }
+                }
                 mSystemSupport.removeDeviceIdLocked(deviceId);
                 mLastStates.delete(deviceId);
             }
@@ -307,7 +332,7 @@
      * Returns {@code true} if {@code deviceId} is being proxy-ed.
      */
     public boolean isProxyedDeviceId(int deviceId) {
-        if (deviceId == DEVICE_ID_DEFAULT && deviceId == DEVICE_ID_INVALID) {
+        if (deviceId == DEVICE_ID_DEFAULT || deviceId == DEVICE_ID_INVALID) {
             return false;
         }
         boolean isTrackingDeviceId;
@@ -566,7 +591,7 @@
      * This is similar to onUserStateChangeLocked and onClientChangeLocked, but does not require an
      * A11yUserState and only checks proxy-relevant settings.
      */
-    public void onProxyChanged(int deviceId) {
+    private void onProxyChanged(int deviceId, boolean forceUpdate) {
         if (DEBUG) {
             Slog.v(LOG_TAG, "onProxyChanged called for deviceId: " + deviceId);
         }
@@ -584,7 +609,7 @@
             // Calls A11yManager#setRelevantEventTypes (test these)
             updateRelevantEventTypesLocked(deviceId);
             // Calls A11yManager#setState
-            scheduleUpdateProxyClientsIfNeededLocked(deviceId);
+            scheduleUpdateProxyClientsIfNeededLocked(deviceId, forceUpdate);
             //Calls A11yManager#notifyServicesStateChanged(timeout)
             scheduleNotifyProxyClientsOfServicesStateChangeLocked(deviceId);
             // Calls A11yManager#setFocusAppearance
@@ -594,16 +619,25 @@
     }
 
     /**
+     * Handles proxy changes, but does not force an update of app clients.
+     */
+    public void onProxyChanged(int deviceId) {
+        onProxyChanged(deviceId, false);
+    }
+
+    /**
      * Updates the states of the app AccessibilityManagers.
      */
-    private void scheduleUpdateProxyClientsIfNeededLocked(int deviceId) {
+    private void scheduleUpdateProxyClientsIfNeededLocked(int deviceId, boolean forceUpdate) {
         final int proxyState = getStateLocked(deviceId);
         if (DEBUG) {
             Slog.v(LOG_TAG, "State for device id " + deviceId + " is " + proxyState);
             Slog.v(LOG_TAG, "Last state for device id " + deviceId + " is "
                     + getLastSentStateLocked(deviceId));
+            Slog.v(LOG_TAG, "force update: " + forceUpdate);
         }
-        if ((getLastSentStateLocked(deviceId)) != proxyState) {
+        if ((getLastSentStateLocked(deviceId)) != proxyState
+                || (Flags.proxyUseAppsOnVirtualDeviceListener() && forceUpdate)) {
             setLastStateLocked(deviceId, proxyState);
             mMainHandler.post(() -> {
                 synchronized (mLock) {
@@ -792,7 +826,7 @@
     }
 
     /**
-     * Updates the device ids of IAccessibilityManagerClients if needed.
+     * Updates the device ids of IAccessibilityManagerClients if needed after a proxy change.
      */
     private void updateDeviceIdsIfNeededLocked(int deviceId,
             @NonNull RemoteCallbackList<IAccessibilityManagerClient> clients) {
@@ -804,13 +838,66 @@
         for (int i = 0; i < clients.getRegisteredCallbackCount(); i++) {
             final AccessibilityManagerService.Client client =
                     ((AccessibilityManagerService.Client) clients.getRegisteredCallbackCookie(i));
-            if (deviceId != DEVICE_ID_DEFAULT && deviceId != DEVICE_ID_INVALID
-                    && localVdm.getDeviceIdsForUid(client.mUid).contains(deviceId)) {
-                if (DEBUG) {
-                    Slog.v(LOG_TAG, "Packages moved to device id " + deviceId + " are "
-                            + Arrays.toString(client.mPackageNames));
+            if (Flags.proxyUseAppsOnVirtualDeviceListener()) {
+                if (deviceId == DEVICE_ID_DEFAULT || deviceId == DEVICE_ID_INVALID) {
+                    continue;
                 }
-                client.mDeviceId = deviceId;
+                boolean uidBelongsToDevice =
+                        localVdm.getDeviceIdsForUid(client.mUid).contains(deviceId);
+                if (client.mDeviceId != deviceId && uidBelongsToDevice) {
+                    if (DEBUG) {
+                        Slog.v(LOG_TAG, "Packages moved to device id " + deviceId + " are "
+                                + Arrays.toString(client.mPackageNames));
+                    }
+                    client.mDeviceId = deviceId;
+                } else if (client.mDeviceId == deviceId && !uidBelongsToDevice) {
+                    client.mDeviceId = DEVICE_ID_DEFAULT;
+                    if (DEBUG) {
+                        Slog.v(LOG_TAG, "Packages moved to the default device from device id "
+                                + deviceId + " are " + Arrays.toString(client.mPackageNames));
+                    }
+                }
+            } else {
+                if (deviceId != DEVICE_ID_DEFAULT && deviceId != DEVICE_ID_INVALID
+                    && localVdm.getDeviceIdsForUid(client.mUid).contains(deviceId)) {
+                    if (DEBUG) {
+                        Slog.v(LOG_TAG, "Packages moved to device id " + deviceId + " are "
+                                + Arrays.toString(client.mPackageNames));
+                    }
+                    client.mDeviceId = deviceId;
+                }
+            }
+        }
+    }
+
+    @VisibleForTesting
+    void notifyProxyOfRunningAppsChange(Set<Integer> allRunningUids) {
+        if (DEBUG) {
+            Slog.v(LOG_TAG, "notifyProxyOfRunningAppsChange: " + allRunningUids);
+        }
+        synchronized (mLock) {
+            if (mProxyA11yServiceConnections.size() == 0) {
+                return;
+            }
+            final VirtualDeviceManagerInternal localVdm = getLocalVdm();
+            if  (localVdm == null) {
+                return;
+            }
+            final ArraySet<Integer> deviceIdsToUpdate = new ArraySet<>();
+            for (int i = 0; i < mProxyA11yServiceConnections.size(); i++) {
+                final ProxyAccessibilityServiceConnection proxy =
+                        mProxyA11yServiceConnections.valueAt(i);
+                if (proxy != null) {
+                    final int proxyDeviceId = proxy.getDeviceId();
+                    for (Integer uid : allRunningUids) {
+                        if (localVdm.getDeviceIdsForUid(uid).contains(proxyDeviceId)) {
+                            deviceIdsToUpdate.add(proxyDeviceId);
+                        }
+                    }
+                }
+            }
+            for (Integer proxyDeviceId : deviceIdsToUpdate) {
+                onProxyChanged(proxyDeviceId, true);
             }
         }
     }
@@ -843,6 +930,11 @@
         return mLocalVdm;
     }
 
+    @VisibleForTesting
+    void setLocalVirtualDeviceManager(VirtualDeviceManagerInternal localVdm) {
+        mLocalVdm = localVdm;
+    }
+
     /**
      * Prints information belonging to each display that is controlled by an
      * AccessibilityDisplayProxy.