Unhide the storage attribution API with system visibility.

Bug: 179160886
Test: StorageStatsManagerLocalTest
Change-Id: I3a8994e215baff123df65d89d2d71b9e12de2af5

8 files changed, 202 insertions, 73 deletions

diff --git a/apex/appsearch/service/java/com/android/server/appsearch/AppSearchManagerService.java b/apex/appsearch/service/java/com/android/server/appsearch/AppSearchManagerService.java
index a4188a2733b8..6e90655dfc99 100644
--- a/apex/appsearch/service/java/com/android/server/appsearch/AppSearchManagerService.java
+++ b/apex/appsearch/service/java/com/android/server/appsearch/AppSearchManagerService.java
@@ -41,6 +41,7 @@ import android.content.BroadcastReceiver;
 import android.content.Context;
 import android.content.Intent;
 import android.content.IntentFilter;
+import android.content.pm.PackageInfo;
 import android.content.pm.PackageManager;
 import android.content.pm.PackageManagerInternal;
 import android.content.pm.PackageStats;
@@ -56,14 +57,15 @@ import android.util.ArraySet;
 import android.util.Log;
 
 import com.android.internal.annotations.GuardedBy;
+import com.android.server.LocalManagerRegistry;
 import com.android.server.LocalServices;
 import com.android.server.SystemService;
 import com.android.server.appsearch.external.localstorage.AppSearchImpl;
 import com.android.server.appsearch.external.localstorage.stats.CallStats;
 import com.android.server.appsearch.stats.LoggerInstanceManager;
 import com.android.server.appsearch.stats.PlatformLogger;
-import com.android.server.usage.StorageStatsManagerInternal;
-import com.android.server.usage.StorageStatsManagerInternal.StorageStatsAugmenter;
+import com.android.server.usage.StorageStatsManagerLocal;
+import com.android.server.usage.StorageStatsManagerLocal.StorageStatsAugmenter;
 
 import com.google.android.icing.proto.PersistType;
 
@@ -120,7 +122,7 @@ public class AppSearchManagerService extends SystemService {
         mUserManager = mContext.getSystemService(UserManager.class);
         mLoggerInstanceManager = LoggerInstanceManager.getInstance();
         registerReceivers();
-        LocalServices.getService(StorageStatsManagerInternal.class)
+        LocalManagerRegistry.getManager(StorageStatsManagerLocal.class)
                 .registerStorageStatsAugmenter(new AppSearchStorageStatsAugmenter(), TAG);
     }
 
@@ -931,13 +933,15 @@ public class AppSearchManagerService extends SystemService {
     // TODO(b/179160886): Cache the previous storage stats.
     private class AppSearchStorageStatsAugmenter implements StorageStatsAugmenter {
         @Override
-        public void augmentStatsForPackage(
+        public void augmentStatsForPackageForUser(
                 @NonNull PackageStats stats,
                 @NonNull String packageName,
-                @UserIdInt int userId,
-                boolean callerHasStatsPermission) {
+                @NonNull UserHandle userHandle,
+                boolean canCallerAccessAllStats) {
             Objects.requireNonNull(stats);
             Objects.requireNonNull(packageName);
+            Objects.requireNonNull(userHandle);
+            int userId = userHandle.getIdentifier();
             try {
                 verifyUserUnlocked(userId);
                 AppSearchImpl impl = mImplInstanceManager.getOrCreateAppSearchImpl(mContext,
@@ -956,7 +960,7 @@ public class AppSearchManagerService extends SystemService {
 
         @Override
         public void augmentStatsForUid(
-                @NonNull PackageStats stats, int uid, boolean callerHasStatsPermission) {
+                @NonNull PackageStats stats, int uid, boolean canCallerAccessAllStats) {
             Objects.requireNonNull(stats);
             int userId = UserHandle.getUserId(uid);
             try {
@@ -967,12 +971,40 @@ public class AppSearchManagerService extends SystemService {
                 }
                 AppSearchImpl impl = mImplInstanceManager.getOrCreateAppSearchImpl(mContext,
                         userId);
-                for (String packageName : packagesForUid) {
-                    stats.dataSize += impl.getStorageInfoForPackage(packageName).getSizeBytes();
+                for (int i = 0; i < packagesForUid.length; i++) {
+                    stats.dataSize +=
+                            impl.getStorageInfoForPackage(packagesForUid[i]).getSizeBytes();
                 }
             } catch (Throwable t) {
                 Log.e(TAG, "Unable to augment storage stats for uid " + uid, t);
             }
         }
+
+        @Override
+        public void augmentStatsForUser(
+                @NonNull PackageStats stats, @NonNull UserHandle userHandle) {
+            // TODO(b/179160886): this implementation could incur many jni calls and a lot of
+            //  in-memory processing from getStorageInfoForPackage. Instead, we can just compute the
+            //  size of the icing dir (or use the overall StorageInfo without interpolating it).
+            Objects.requireNonNull(stats);
+            Objects.requireNonNull(userHandle);
+            int userId = userHandle.getIdentifier();
+            try {
+                verifyUserUnlocked(userId);
+                List<PackageInfo> packagesForUser =
+                        mPackageManager.getInstalledPackagesAsUser(/*flags=*/0, userId);
+                if (packagesForUser == null) {
+                    return;
+                }
+                AppSearchImpl impl =
+                        mImplInstanceManager.getOrCreateAppSearchImpl(mContext, userId);
+                for (int i = 0; i < packagesForUser.size(); i++) {
+                    String packageName = packagesForUser.get(i).packageName;
+                    stats.dataSize += impl.getStorageInfoForPackage(packageName).getSizeBytes();
+                }
+            } catch (Throwable t) {
+                Log.e(TAG, "Unable to augment storage stats for user " + userId, t);
+            }
+        }
     }
 }
diff --git a/apex/appsearch/testing/Android.bp b/apex/appsearch/testing/Android.bp
index ec649410872f..5407cb4ccec7 100644
--- a/apex/appsearch/testing/Android.bp
+++ b/apex/appsearch/testing/Android.bp
@@ -33,7 +33,7 @@ java_library {
     visibility: [
         "//frameworks/base/core/tests/coretests",
         "//cts/hostsidetests/appsearch",
-        "//cts/tests/appsearch",
+        "//cts/tests:__subpackages__",
         "//vendor:__subpackages__",
     ],
 }
diff --git a/apex/blobstore/service/java/com/android/server/blob/BlobStoreManagerService.java b/apex/blobstore/service/java/com/android/server/blob/BlobStoreManagerService.java
index cc5e31a91123..731ba9254ea4 100644
--- a/apex/blobstore/service/java/com/android/server/blob/BlobStoreManagerService.java
+++ b/apex/blobstore/service/java/com/android/server/blob/BlobStoreManagerService.java
@@ -101,14 +101,15 @@ import com.android.internal.util.FrameworkStatsLog;
 import com.android.internal.util.Preconditions;
 import com.android.internal.util.XmlUtils;
 import com.android.internal.util.function.pooled.PooledLambda;
+import com.android.server.LocalManagerRegistry;
 import com.android.server.LocalServices;
 import com.android.server.ServiceThread;
 import com.android.server.SystemService;
 import com.android.server.Watchdog;
 import com.android.server.blob.BlobMetadata.Committer;
 import com.android.server.pm.UserManagerInternal;
-import com.android.server.usage.StorageStatsManagerInternal;
-import com.android.server.usage.StorageStatsManagerInternal.StorageStatsAugmenter;
+import com.android.server.usage.StorageStatsManagerLocal;
+import com.android.server.usage.StorageStatsManagerLocal.StorageStatsAugmenter;
 
 import org.xmlpull.v1.XmlPullParser;
 import org.xmlpull.v1.XmlSerializer;
@@ -208,7 +209,7 @@ public class BlobStoreManagerService extends SystemService {
         mPackageManagerInternal = LocalServices.getService(PackageManagerInternal.class);
         mStatsManager = getContext().getSystemService(StatsManager.class);
         registerReceivers();
-        LocalServices.getService(StorageStatsManagerInternal.class)
+        LocalManagerRegistry.getManager(StorageStatsManagerLocal.class)
                 .registerStorageStatsAugmenter(new BlobStorageStatsAugmenter(), TAG);
     }
 
@@ -1281,17 +1282,20 @@ public class BlobStoreManagerService extends SystemService {
 
     private class BlobStorageStatsAugmenter implements StorageStatsAugmenter {
         @Override
-        public void augmentStatsForPackage(@NonNull PackageStats stats, @NonNull String packageName,
-                @UserIdInt int userId, boolean callerHasStatsPermission) {
+        public void augmentStatsForPackageForUser(
+                @NonNull PackageStats stats,
+                @NonNull String packageName,
+                @NonNull UserHandle userHandle,
+                boolean callerHasStatsPermission) {
             final AtomicLong blobsDataSize = new AtomicLong(0);
             forEachSessionInUser(session -> {
                 if (session.getOwnerPackageName().equals(packageName)) {
                     blobsDataSize.getAndAdd(session.getSize());
                 }
-            }, userId);
+            }, userHandle.getIdentifier());
 
             forEachBlob(blobMetadata -> {
-                if (blobMetadata.shouldAttributeToLeasee(packageName, userId,
+                if (blobMetadata.shouldAttributeToLeasee(packageName, userHandle.getIdentifier(),
                         callerHasStatsPermission)) {
                     blobsDataSize.getAndAdd(blobMetadata.getSize());
                 }
@@ -1320,6 +1324,22 @@ public class BlobStoreManagerService extends SystemService {
 
             stats.dataSize += blobsDataSize.get();
         }
+
+        @Override
+        public void augmentStatsForUser(
+                @NonNull PackageStats stats, @NonNull UserHandle userHandle) {
+            final AtomicLong blobsDataSize = new AtomicLong(0);
+            forEachSessionInUser(session -> {
+                blobsDataSize.getAndAdd(session.getSize());
+            }, userHandle.getIdentifier());
+
+            // TODO(http://b/187460239): Update this to only include blobs available to userId.
+            forEachBlob(blobMetadata -> {
+                blobsDataSize.getAndAdd(blobMetadata.getSize());
+            });
+
+            stats.dataSize += blobsDataSize.get();
+        }
     }
 
     private void forEachSessionInUser(Consumer<BlobStoreSession> consumer, int userId) {
@@ -1900,4 +1920,4 @@ public class BlobStoreManagerService extends SystemService {
             return BackgroundThread.getHandler();
         }
     }
-}
\ No newline at end of file
+}
diff --git a/services/api/current.txt b/services/api/current.txt
index 475dcf5eb20e..50f00524676c 100644
--- a/services/api/current.txt
+++ b/services/api/current.txt
@@ -60,6 +60,20 @@ package com.android.server.stats {
 
 }
 
+package com.android.server.usage {
+
+  public interface StorageStatsManagerLocal {
+    method public void registerStorageStatsAugmenter(@NonNull com.android.server.usage.StorageStatsManagerLocal.StorageStatsAugmenter, @NonNull String);
+  }
+
+  public static interface StorageStatsManagerLocal.StorageStatsAugmenter {
+    method public void augmentStatsForPackageForUser(@NonNull android.content.pm.PackageStats, @NonNull String, @NonNull android.os.UserHandle, boolean);
+    method public void augmentStatsForUid(@NonNull android.content.pm.PackageStats, int, boolean);
+    method public void augmentStatsForUser(@NonNull android.content.pm.PackageStats, @NonNull android.os.UserHandle);
+  }
+
+}
+
 package com.android.server.wifi {
 
   public class SupplicantManager {
diff --git a/services/usage/java/com/android/server/usage/StorageStatsManagerInternal.java b/services/usage/java/com/android/server/usage/StorageStatsManagerInternal.java
deleted file mode 100644
index 47760efe2709..000000000000
--- a/services/usage/java/com/android/server/usage/StorageStatsManagerInternal.java
+++ /dev/null
@@ -1,50 +0,0 @@
-/*
- * Copyright (C) 2020 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.usage;
-
-import android.annotation.NonNull;
-import android.annotation.UserIdInt;
-import android.content.pm.PackageStats;
-
-/**
- * StorageStatsManager local system service interface.
- *
- * Only for use within the system server.
- */
-public abstract class StorageStatsManagerInternal {
-    /**
-     * Class used to augment {@link PackageStats} with the data stored by the system on
-     * behalf of apps in system specific directories
-     * ({@link android.os.Environment#getDataSystemDirectory},
-     * {@link android.os.Environment#getDataSystemCeDirectory}, etc).
-     */
-    public interface StorageStatsAugmenter {
-        void augmentStatsForPackage(@NonNull PackageStats stats,
-                @NonNull String packageName, @UserIdInt int userId,
-                boolean callerHasStatsPermission);
-        void augmentStatsForUid(@NonNull PackageStats stats, int uid,
-                boolean callerHasStatsPermission);
-    }
-
-    /**
-     * Register a {@link StorageStatsAugmenter}.
-     *
-     * @param augmenter the {@link StorageStatsAugmenter} object to be registered.
-     * @param tag the identifier to be used for debugging in logs/trace.
-     */
-    public abstract void registerStorageStatsAugmenter(@NonNull StorageStatsAugmenter augmenter,
-            @NonNull String tag);
-}
diff --git a/services/usage/java/com/android/server/usage/StorageStatsManagerLocal.java b/services/usage/java/com/android/server/usage/StorageStatsManagerLocal.java
new file mode 100644
index 000000000000..7cac820b0665
--- /dev/null
+++ b/services/usage/java/com/android/server/usage/StorageStatsManagerLocal.java
@@ -0,0 +1,87 @@
+/*
+ * Copyright (C) 2021 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.usage;
+
+import android.annotation.NonNull;
+import android.annotation.SystemApi;
+import android.content.pm.PackageStats;
+import android.os.UserHandle;
+
+/**
+ * StorageStatsManager local system service interface.
+ *
+ * @hide Only for use within the system server.
+ */
+@SystemApi(client = SystemApi.Client.SYSTEM_SERVER)
+public interface StorageStatsManagerLocal {
+    /**
+     * Class used to augment {@link PackageStats} with the data stored by the system on
+     * behalf of apps in system specific directories
+     * ({@link android.os.Environment#getDataSystemDirectory},
+     * {@link android.os.Environment#getDataSystemCeDirectory}, etc).
+     */
+    interface StorageStatsAugmenter {
+        /**
+         * Augments {@link PackageStats} with data stored by the system for the given package.
+         *
+         * @param stats                   Structure to modify with usage data
+         * @param packageName             Package name of the app whose data is stored by the
+         *                                system and needs to be added to {@code stats}.
+         * @param userHandle              Device user for which usage stats are being requested.
+         * @param canCallerAccessAllStats Whether the caller who is requesting the storage stats
+         *                                can query stats for packages other than itself. For
+         *                                example, holding the PACKAGE_USAGE_STATS permission is one
+         *                                way to accomplish this.
+         */
+        void augmentStatsForPackageForUser(
+                @NonNull PackageStats stats,
+                @NonNull String packageName,
+                @NonNull UserHandle userHandle,
+                boolean canCallerAccessAllStats);
+
+        /**
+         * Augments {@link PackageStats} with data stored by the system for the given uid.
+         *
+         * @param stats                   Structure to modify with usage data
+         * @param uid                     Unique app ID for the app instance whose stats are being
+         *                                requested.
+         * @param canCallerAccessAllStats Whether the caller who is requesting the storage stats
+         *                                can query stats for packages other than itself. For
+         *                                example, holding the PACKAGE_USAGE_STATS permission is one
+         *                                way to accomplish this.
+         */
+        void augmentStatsForUid(
+                @NonNull PackageStats stats, int uid, boolean canCallerAccessAllStats);
+
+        /**
+         * Augments {@link PackageStats} with data stored by the system for the given device user.
+         *
+         * @param stats      Structure to modify with usage data
+         * @param userHandle Device user whose data is stored by the system and needs to be added to
+         *                   {@code stats}.
+         */
+        void augmentStatsForUser(@NonNull PackageStats stats, @NonNull UserHandle userHandle);
+    }
+
+    /**
+     * Register a {@link StorageStatsAugmenter}.
+     *
+     * @param augmenter the {@link StorageStatsAugmenter} object to be registered.
+     * @param tag       the identifier to be used for debugging in logs/trace.
+     */
+    void registerStorageStatsAugmenter(
+            @NonNull StorageStatsAugmenter augmenter, @NonNull String tag);
+}
diff --git a/services/usage/java/com/android/server/usage/StorageStatsService.java b/services/usage/java/com/android/server/usage/StorageStatsService.java
index a0a390947822..b056de075c36 100644
--- a/services/usage/java/com/android/server/usage/StorageStatsService.java
+++ b/services/usage/java/com/android/server/usage/StorageStatsService.java
@@ -20,7 +20,7 @@ import static android.content.pm.PackageManager.PERMISSION_GRANTED;
 
 import static com.android.internal.util.ArrayUtils.defeatNullable;
 import static com.android.server.pm.PackageManagerService.PLATFORM_PACKAGE_NAME;
-import static com.android.server.usage.StorageStatsManagerInternal.StorageStatsAugmenter;
+import static com.android.server.usage.StorageStatsManagerLocal.StorageStatsAugmenter;
 
 import android.Manifest;
 import android.annotation.NonNull;
@@ -71,6 +71,7 @@ import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.util.ArrayUtils;
 import com.android.internal.util.Preconditions;
 import com.android.server.IoThread;
+import com.android.server.LocalManagerRegistry;
 import com.android.server.LocalServices;
 import com.android.server.SystemService;
 import com.android.server.pm.Installer;
@@ -152,7 +153,7 @@ public class StorageStatsService extends IStorageStatsManager.Stub {
             }
         });
 
-        LocalServices.addService(StorageStatsManagerInternal.class, new LocalService());
+        LocalManagerRegistry.addManager(StorageStatsManagerLocal.class, new LocalService());
     }
 
     private void invalidateMounts() {
@@ -335,9 +336,10 @@ public class StorageStatsService extends IStorageStatsManager.Stub {
                 throw new ParcelableException(new IOException(e.getMessage()));
             }
             if (volumeUuid == StorageManager.UUID_PRIVATE_INTERNAL) {
+                UserHandle userHandle = UserHandle.of(userId);
                 forEachStorageStatsAugmenter((storageStatsAugmenter) -> {
-                    storageStatsAugmenter.augmentStatsForPackage(stats,
-                            packageName, userId, callerHasStatsPermission);
+                    storageStatsAugmenter.augmentStatsForPackageForUser(stats,
+                            packageName, userHandle, callerHasStatsPermission);
                 }, "queryStatsForPackage");
             }
             return translate(stats);
@@ -430,6 +432,12 @@ public class StorageStatsService extends IStorageStatsManager.Stub {
         } catch (InstallerException e) {
             throw new ParcelableException(new IOException(e.getMessage()));
         }
+        if (volumeUuid == StorageManager.UUID_PRIVATE_INTERNAL) {
+            UserHandle userHandle = UserHandle.of(userId);
+            forEachStorageStatsAugmenter((storageStatsAugmenter) -> {
+                storageStatsAugmenter.augmentStatsForUser(stats, userHandle);
+            }, "queryStatsForUser");
+        }
         return translate(stats);
     }
 
@@ -784,7 +792,7 @@ public class StorageStatsService extends IStorageStatsManager.Stub {
         }
     }
 
-    private class LocalService extends StorageStatsManagerInternal {
+    private class LocalService implements StorageStatsManagerLocal {
         @Override
         public void registerStorageStatsAugmenter(
                 @NonNull StorageStatsAugmenter storageStatsAugmenter,
diff --git a/services/usage/java/com/android/server/usage/package-info.java b/services/usage/java/com/android/server/usage/package-info.java
new file mode 100644
index 000000000000..0719df418fd8
--- /dev/null
+++ b/services/usage/java/com/android/server/usage/package-info.java
@@ -0,0 +1,18 @@
+/*
+ * Copyright (C) 2021 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.
+ */
+
+/** @hide */
+package com.android.server.usage;