Fix adb command for multiple users

Add userID as a parameter when calling addAssociation
and removeAssociation since getCallingUserId() always
return 0 for adb associate/disassociate.

Fix: 191699828
Test: manual test
Change-Id: Ia33d38d0d23a2275b22ff7748f74c6e06b754408

1 files changed, 10 insertions, 13 deletions

diff --git a/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java b/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
index 906edb30ff12..9ff1b10c09ed 100644
--- a/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
+++ b/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
@@ -455,7 +455,7 @@ public class CompanionDeviceManagerService extends SystemService implements Bind
 
             }, FgThread.getExecutor()).whenComplete(uncheckExceptions((association, err) -> {
                 if (err == null) {
-                    addAssociation(association);
+                    addAssociation(association, userId);
                 } else {
                     Slog.e(LOG_TAG, "Failed to discover device(s)", err);
                     callback.onFailure("No devices found: " + err.getMessage());
@@ -646,7 +646,7 @@ public class CompanionDeviceManagerService extends SystemService implements Bind
                 } else {
                     return association;
                 }
-            }));
+            }), userId);
 
             restartBleScan();
         }
@@ -664,7 +664,8 @@ public class CompanionDeviceManagerService extends SystemService implements Bind
                     android.Manifest.permission.ASSOCIATE_COMPANION_DEVICES, "createAssociation");
 
             addAssociation(new Association(
-                    userId, macAddress, packageName, null, false, System.currentTimeMillis()));
+                    userId, macAddress, packageName, null, false,
+                    System.currentTimeMillis()), userId);
         }
 
         private void checkCanCallNotificationApi(String callingPackage) throws RemoteException {
@@ -738,9 +739,9 @@ public class CompanionDeviceManagerService extends SystemService implements Bind
         return Binder.getCallingUid() == Process.SYSTEM_UID;
     }
 
-    void addAssociation(Association association) {
+    void addAssociation(Association association, int userId) {
         updateSpecialAccessPermissionForAssociatedPackage(association);
-        recordAssociation(association);
+        recordAssociation(association, userId);
     }
 
     void removeAssociation(int userId, String pkg, String deviceMacAddress) {
@@ -752,7 +753,7 @@ public class CompanionDeviceManagerService extends SystemService implements Bind
                 onAssociationPreRemove(association);
             }
             return notMatch;
-        }));
+        }), userId);
         restartBleScan();
     }
 
@@ -944,13 +945,9 @@ public class CompanionDeviceManagerService extends SystemService implements Bind
         }, getContext(), packageName, userId).recycleOnUse());
     }
 
-    private void recordAssociation(Association association) {
+    private void recordAssociation(Association association, int userId) {
         Slog.i(LOG_TAG, "recordAssociation(" + association + ")");
-        updateAssociations(associations -> CollectionUtils.add(associations, association));
-    }
-
-    private void updateAssociations(Function<Set<Association>, Set<Association>> update) {
-        updateAssociations(update, getCallingUserId());
+        updateAssociations(associations -> CollectionUtils.add(associations, association), userId);
     }
 
     private void updateAssociations(Function<Set<Association>, Set<Association>> update,
@@ -1515,7 +1512,7 @@ public class CompanionDeviceManagerService extends SystemService implements Bind
                         String pkg = getNextArgRequired();
                         String address = getNextArgRequired();
                         addAssociation(new Association(userId, address, pkg, null, false,
-                                System.currentTimeMillis()));
+                                System.currentTimeMillis()), userId);
                     }
                     break;