[Companion] Dont store duplicate association records

Fixes: 62675985
Test: Assosiate the same item twice and ensure getAssociations lists it
only once
Change-Id: I028c08010740baaa04464647dffa701fc066a4fe

3 files changed, 137 insertions, 23 deletions

diff --git a/core/java/com/android/internal/util/CollectionUtils.java b/core/java/com/android/internal/util/CollectionUtils.java
index a7fdb27e03b8..dbb6e93de7d4 100644
--- a/core/java/com/android/internal/util/CollectionUtils.java
+++ b/core/java/com/android/internal/util/CollectionUtils.java
@@ -21,13 +21,16 @@ import static com.android.internal.util.ArrayUtils.isEmpty;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.util.ArraySet;
+import android.util.ExceptionUtils;
+
+import com.android.internal.util.FunctionalUtils.ThrowingConsumer;
 
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.Set;
-import java.util.function.Function;
+import java.util.function.*;
 import java.util.stream.Stream;
 
 /**
@@ -58,6 +61,32 @@ public class CollectionUtils {
     }
 
     /**
+     * @see #filter(List, java.util.function.Predicate)
+     */
+    public static @NonNull <T> Set<T> filter(@Nullable Set<T> set,
+            java.util.function.Predicate<? super T> predicate) {
+        if (set == null || set.size() == 0) return Collections.emptySet();
+        ArraySet<T> result = null;
+        if (set instanceof ArraySet) {
+            ArraySet<T> arraySet = (ArraySet<T>) set;
+            int size = arraySet.size();
+            for (int i = 0; i < size; i++) {
+                final T item = arraySet.valueAt(i);
+                if (predicate.test(item)) {
+                    result = ArrayUtils.add(result, item);
+                }
+            }
+        } else {
+            for (T item : set) {
+                if (predicate.test(item)) {
+                    result = ArrayUtils.add(result, item);
+                }
+            }
+        }
+        return emptyIfNull(result);
+    }
+
+    /**
      * Returns a list of items resulting from applying the given function to each element of the
      * provided list.
      *
@@ -77,6 +106,27 @@ public class CollectionUtils {
     }
 
     /**
+     * @see #map(List, Function)
+     */
+    public static @NonNull <I, O> Set<O> map(@Nullable Set<I> cur,
+            Function<? super I, ? extends O> f) {
+        if (isEmpty(cur)) return Collections.emptySet();
+        ArraySet<O> result = new ArraySet<>();
+        if (cur instanceof ArraySet) {
+            ArraySet<I> arraySet = (ArraySet<I>) cur;
+            int size = arraySet.size();
+            for (int i = 0; i < size; i++) {
+                result.add(f.apply(arraySet.valueAt(i)));
+            }
+        } else {
+            for (I item : cur) {
+                result.add(f.apply(item));
+            }
+        }
+        return result;
+    }
+
+    /**
      * {@link #map(List, Function)} + {@link #filter(List, java.util.function.Predicate)}
      *
      * Calling this is equivalent (but more memory efficient) to:
@@ -180,6 +230,17 @@ public class CollectionUtils {
     }
 
     /**
+     * @see #add(List, Object)
+     */
+    public static @NonNull <T> Set<T> add(@Nullable Set<T> cur, T val) {
+        if (cur == null || cur == Collections.emptySet()) {
+            cur = new ArraySet<>();
+        }
+        cur.add(val);
+        return cur;
+    }
+
+    /**
      * Similar to {@link List#remove}, but with support for list values of {@code null} and
      * {@link Collections#emptyList}
      */
@@ -192,9 +253,52 @@ public class CollectionUtils {
     }
 
     /**
+     * @see #remove(List, Object)
+     */
+    public static @NonNull <T> Set<T> remove(@Nullable Set<T> cur, T val) {
+        if (isEmpty(cur)) {
+            return emptyIfNull(cur);
+        }
+        cur.remove(val);
+        return cur;
+    }
+
+    /**
      * @return a list that will not be affected by mutations to the given original list.
      */
     public static @NonNull <T> List<T> copyOf(@Nullable List<T> cur) {
         return isEmpty(cur) ? Collections.emptyList() : new ArrayList<>(cur);
     }
+
+    /**
+     * @return a list that will not be affected by mutations to the given original list.
+     */
+    public static @NonNull <T> Set<T> copyOf(@Nullable Set<T> cur) {
+        return isEmpty(cur) ? Collections.emptySet() : new ArraySet<>(cur);
+    }
+
+    /**
+     * Applies {@code action} to each element in {@code cur}
+     *
+     * This avoids creating an iterator if the given set is an {@link ArraySet}
+     */
+    public static <T> void forEach(@Nullable Set<T> cur, @Nullable ThrowingConsumer<T> action) {
+        if (cur == null || action == null) return;
+        int size = cur.size();
+        if (size == 0) return;
+        try {
+            if (cur instanceof ArraySet) {
+                ArraySet<T> arraySet = (ArraySet<T>) cur;
+                for (int i = 0; i < size; i++) {
+                    action.accept(arraySet.valueAt(i));
+                }
+            } else {
+                for (T t : cur) {
+                    action.accept(t);
+                }
+            }
+        } catch (Exception e) {
+            throw ExceptionUtils.propagate(e);
+        }
+    }
 }
diff --git a/core/java/com/android/internal/util/FunctionalUtils.java b/core/java/com/android/internal/util/FunctionalUtils.java
index 9aeb0415b5fc..cdef97e84f62 100644
--- a/core/java/com/android/internal/util/FunctionalUtils.java
+++ b/core/java/com/android/internal/util/FunctionalUtils.java
@@ -45,4 +45,15 @@ public class FunctionalUtils {
     public interface ThrowingSupplier<T> {
         T get() throws Exception;
     }
+
+    /**
+     * An equivalent of {@link java.util.function.Consumer} that allows throwing checked exceptions
+     *
+     * This can be used to specify a lambda argument without forcing all the checked exceptions
+     * to be handled within it
+     */
+    @FunctionalInterface
+    public interface ThrowingConsumer<T> {
+        void accept(T t) throws Exception;
+    }
 }
diff --git a/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java b/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
index f47b0d3c6e73..f2f01cfa19b0 100644
--- a/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
+++ b/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
@@ -57,6 +57,7 @@ import android.os.UserHandle;
 import android.provider.Settings;
 import android.provider.SettingsStringUtil.ComponentNameSet;
 import android.text.BidiFormatter;
+import android.util.ArraySet;
 import android.util.AtomicFile;
 import android.util.ExceptionUtils;
 import android.util.Log;
@@ -83,6 +84,7 @@ import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
+import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.function.Function;
@@ -247,9 +249,9 @@ public class CompanionDeviceManagerService extends SystemService implements Bind
                 throws RemoteException {
             checkCallerIsSystemOr(callingPackage, userId);
             checkUsesFeature(callingPackage, getCallingUserId());
-            return CollectionUtils.map(
+            return new ArrayList<>(CollectionUtils.map(
                     readAllAssociations(userId, callingPackage),
-                    a -> a.deviceAddress);
+                    a -> a.deviceAddress));
         }
 
         //TODO also revoke notification access
@@ -495,20 +497,20 @@ public class CompanionDeviceManagerService extends SystemService implements Bind
                 new Association(userId, deviceAddress, priviledgedPackage)));
     }
 
-    private void updateAssociations(Function<List<Association>, List<Association>> update) {
+    private void updateAssociations(Function<Set<Association>, Set<Association>> update) {
         updateAssociations(update, getCallingUserId());
     }
 
-    private void updateAssociations(Function<List<Association>, List<Association>> update,
+    private void updateAssociations(Function<Set<Association>, Set<Association>> update,
             int userId) {
         final AtomicFile file = getStorageFileForUser(userId);
         synchronized (file) {
-            List<Association> associations = readAllAssociations(userId);
-            final List<Association> old = CollectionUtils.copyOf(associations);
+            Set<Association> associations = readAllAssociations(userId);
+            final Set<Association> old = CollectionUtils.copyOf(associations);
             associations = update.apply(associations);
             if (size(old) == size(associations)) return;
 
-            List<Association> finalAssociations = associations;
+            Set<Association> finalAssociations = associations;
             file.write((out) -> {
                 XmlSerializer xml = Xml.newSerializer();
                 try {
@@ -517,13 +519,12 @@ public class CompanionDeviceManagerService extends SystemService implements Bind
                     xml.startDocument(null, true);
                     xml.startTag(null, XML_TAG_ASSOCIATIONS);
 
-                    for (int i = 0; i < size(finalAssociations); i++) {
-                        Association association = finalAssociations.get(i);
+                    CollectionUtils.forEach(finalAssociations, association -> {
                         xml.startTag(null, XML_TAG_ASSOCIATION)
-                            .attribute(null, XML_ATTR_PACKAGE, association.companionAppPackage)
-                            .attribute(null, XML_ATTR_DEVICE, association.deviceAddress)
-                            .endTag(null, XML_TAG_ASSOCIATION);
-                    }
+                                .attribute(null, XML_ATTR_PACKAGE, association.companionAppPackage)
+                                .attribute(null, XML_ATTR_DEVICE, association.deviceAddress)
+                                .endTag(null, XML_TAG_ASSOCIATION);
+                    });
 
                     xml.endTag(null, XML_TAG_ASSOCIATIONS);
                     xml.endDocument();
@@ -545,17 +546,17 @@ public class CompanionDeviceManagerService extends SystemService implements Bind
     }
 
     @Nullable
-    private ArrayList<Association> readAllAssociations(int userId) {
+    private Set<Association> readAllAssociations(int userId) {
         return readAllAssociations(userId, null);
     }
 
     @Nullable
-    private ArrayList<Association> readAllAssociations(int userId, @Nullable String packageFilter) {
+    private Set<Association> readAllAssociations(int userId, @Nullable String packageFilter) {
         final AtomicFile file = getStorageFileForUser(userId);
 
         if (!file.getBaseFile().exists()) return null;
 
-        ArrayList<Association> result = null;
+        ArraySet<Association> result = null;
         final XmlPullParser parser = Xml.newPullParser();
         synchronized (file) {
             try (FileInputStream in = file.openRead()) {
@@ -627,12 +628,10 @@ public class CompanionDeviceManagerService extends SystemService implements Bind
         public int onCommand(String cmd) {
             switch (cmd) {
                 case "list": {
-                    ArrayList<Association> associations = readAllAssociations(getNextArgInt());
-                    for (int i = 0; i < size(associations); i++) {
-                        Association a = associations.get(i);
-                        getOutPrintWriter()
-                                .println(a.companionAppPackage + " " + a.deviceAddress);
-                    }
+                    CollectionUtils.forEach(
+                            readAllAssociations(getNextArgInt()),
+                            a -> getOutPrintWriter()
+                                    .println(a.companionAppPackage + " " + a.deviceAddress));
                 } break;
 
                 case "associate": {