Update BlobStoreMS to augment storage stats with blobs data.

- Any pending sessions data is attributed to the apps
  which contributed them.
- Any commited blobs data is attributed to the app which
  has a lease on it. If multiple apps have lease on a blob, don't
  attribute the blob to those apps for now.
- Remove StorageStatsAugmenter.augmentStatsForUser as it
  is not used for anything currently.
- Fix an issue in how we override existing committers and leasees.

Bug: 148694869
Test: atest cts/tests/BlobStore/src/com/android/cts/blob/BlobStoreManagerTest.java
Test: atest tests/tests/os/src/android/os/storage/cts/StorageStatsManagerTest.java
Test: atest hostsidetests/appsecurity/src/android/appsecurity/cts/StorageHostTest.java
Test: manual
Change-Id: Ia4af0a2549c75db66741f2d1979de95d2d150bc8
diff --git a/apex/blobstore/service/Android.bp b/apex/blobstore/service/Android.bp
index d8caa5a..22b0cbe 100644
--- a/apex/blobstore/service/Android.bp
+++ b/apex/blobstore/service/Android.bp
@@ -23,5 +23,6 @@
     libs: [
         "framework",
         "services.core",
+        "services.usage",
     ],
 }
diff --git a/apex/blobstore/service/java/com/android/server/blob/BlobMetadata.java b/apex/blobstore/service/java/com/android/server/blob/BlobMetadata.java
index c7d803c..4a85a69 100644
--- a/apex/blobstore/service/java/com/android/server/blob/BlobMetadata.java
+++ b/apex/blobstore/service/java/com/android/server/blob/BlobMetadata.java
@@ -26,6 +26,7 @@
 import static android.app.blob.XmlTags.TAG_BLOB_HANDLE;
 import static android.app.blob.XmlTags.TAG_COMMITTER;
 import static android.app.blob.XmlTags.TAG_LEASEE;
+import static android.os.Process.INVALID_UID;
 import static android.system.OsConstants.O_RDONLY;
 
 import static com.android.server.blob.BlobStoreConfig.TAG;
@@ -88,7 +89,7 @@
     private final ArrayMap<String, ArraySet<RevocableFileDescriptor>> mRevocableFds =
             new ArrayMap<>();
 
-    // Do not access this directly, instead use getSessionFile().
+    // Do not access this directly, instead use #getBlobFile().
     private File mBlobFile;
 
     BlobMetadata(Context context, long blobId, BlobHandle blobHandle, int userId) {
@@ -112,12 +113,16 @@
 
     void addCommitter(@NonNull Committer committer) {
         synchronized (mMetadataLock) {
+            // We need to override the committer data, so first remove any existing
+            // committer before adding the new one.
+            mCommitters.remove(committer);
             mCommitters.add(committer);
         }
     }
 
     void addCommitters(ArraySet<Committer> committers) {
         synchronized (mMetadataLock) {
+            mCommitters.clear();
             mCommitters.addAll(committers);
         }
     }
@@ -147,13 +152,18 @@
     void addLeasee(String callingPackage, int callingUid, int descriptionResId,
             CharSequence description, long leaseExpiryTimeMillis) {
         synchronized (mMetadataLock) {
-            mLeasees.add(new Leasee(callingPackage, callingUid,
-                    descriptionResId, description, leaseExpiryTimeMillis));
+            // We need to override the leasee data, so first remove any existing
+            // leasee before adding the new one.
+            final Leasee leasee = new Leasee(callingPackage, callingUid,
+                    descriptionResId, description, leaseExpiryTimeMillis);
+            mLeasees.remove(leasee);
+            mLeasees.add(leasee);
         }
     }
 
     void addLeasees(ArraySet<Leasee> leasees) {
         synchronized (mMetadataLock) {
+            mLeasees.clear();
             mLeasees.addAll(leasees);
         }
     }
@@ -178,7 +188,11 @@
         }
     }
 
-    boolean isAccessAllowedForCaller(String callingPackage, int callingUid) {
+    long getSize() {
+        return getBlobFile().length();
+    }
+
+    boolean isAccessAllowedForCaller(@NonNull String callingPackage, int callingUid) {
         // TODO: verify blob is still valid (expiryTime is not elapsed)
         synchronized (mMetadataLock) {
             // Check if packageName already holds a lease on the blob.
@@ -209,6 +223,61 @@
         return false;
     }
 
+    boolean isALeasee(@NonNull String packageName) {
+        return isALeasee(packageName, INVALID_UID);
+    }
+
+    boolean isALeasee(int uid) {
+        return isALeasee(null, uid);
+    }
+
+    boolean hasOtherLeasees(@NonNull String packageName) {
+        return hasOtherLeasees(packageName, INVALID_UID);
+    }
+
+    boolean hasOtherLeasees(int uid) {
+        return hasOtherLeasees(null, uid);
+    }
+
+    private boolean isALeasee(@Nullable String packageName, int uid) {
+        synchronized (mMetadataLock) {
+            // Check if the package is a leasee of the data blob.
+            for (int i = 0, size = mLeasees.size(); i < size; ++i) {
+                final Leasee leasee = mLeasees.valueAt(i);
+                if (packageName != null && uid != INVALID_UID
+                        && leasee.equals(packageName, uid)) {
+                    return true;
+                } else if (packageName != null && leasee.packageName.equals(packageName)) {
+                    return true;
+                } else if (uid != INVALID_UID && leasee.uid == uid) {
+                    return true;
+                }
+            }
+        }
+        return false;
+    }
+
+    private boolean hasOtherLeasees(@Nullable String packageName, int uid) {
+        synchronized (mMetadataLock) {
+            if (mCommitters.size() > 1 || mLeasees.size() > 1) {
+                return true;
+            }
+            for (int i = 0, size = mLeasees.size(); i < size; ++i) {
+                final Leasee leasee = mLeasees.valueAt(i);
+                // TODO: Also exclude packages which are signed with same cert?
+                if (packageName != null && uid != INVALID_UID
+                        && !leasee.equals(packageName, uid)) {
+                    return true;
+                } else if (packageName != null && !leasee.packageName.equals(packageName)) {
+                    return true;
+                } else if (uid != INVALID_UID && leasee.uid != uid) {
+                    return true;
+                }
+            }
+        }
+        return false;
+    }
+
     File getBlobFile() {
         if (mBlobFile == null) {
             mBlobFile = BlobStoreConfig.getBlobFile(mBlobId);
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 1cc1c5f..97fff24 100644
--- a/apex/blobstore/service/java/com/android/server/blob/BlobStoreManagerService.java
+++ b/apex/blobstore/service/java/com/android/server/blob/BlobStoreManagerService.java
@@ -52,6 +52,7 @@
 import android.content.IntentFilter;
 import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageManagerInternal;
+import android.content.pm.PackageStats;
 import android.content.res.ResourceId;
 import android.os.Binder;
 import android.os.Handler;
@@ -85,6 +86,8 @@
 import com.android.server.SystemService;
 import com.android.server.Watchdog;
 import com.android.server.blob.BlobMetadata.Committer;
+import com.android.server.usage.StorageStatsManagerInternal;
+import com.android.server.usage.StorageStatsManagerInternal.StorageStatsAugmenter;
 
 import org.xmlpull.v1.XmlPullParser;
 import org.xmlpull.v1.XmlSerializer;
@@ -101,6 +104,8 @@
 import java.util.List;
 import java.util.Objects;
 import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
 
 /**
  * Service responsible for maintaining and facilitating access to data blobs published by apps.
@@ -164,6 +169,8 @@
 
         mPackageManagerInternal = LocalServices.getService(PackageManagerInternal.class);
         registerReceivers();
+        LocalServices.getService(StorageStatsManagerInternal.class)
+                .registerStorageStatsAugmenter(new BlobStorageStatsAugmenter(), TAG);
     }
 
     @Override
@@ -946,6 +953,71 @@
         }
     }
 
+    private class BlobStorageStatsAugmenter implements StorageStatsAugmenter {
+        @Override
+        public void augmentStatsForPackage(@NonNull PackageStats stats, @NonNull String packageName,
+                @UserIdInt int userId, boolean callerHasStatsPermission) {
+            final AtomicLong blobsDataSize = new AtomicLong(0);
+            forEachSessionInUser(session -> {
+                if (session.getOwnerPackageName().equals(packageName)) {
+                    blobsDataSize.getAndAdd(session.getSize());
+                }
+            }, userId);
+
+            forEachBlobInUser(blobMetadata -> {
+                if (blobMetadata.isALeasee(packageName)) {
+                    if (!blobMetadata.hasOtherLeasees(packageName) || !callerHasStatsPermission) {
+                        blobsDataSize.getAndAdd(blobMetadata.getSize());
+                    }
+                }
+            }, userId);
+
+            stats.dataSize += blobsDataSize.get();
+        }
+
+        @Override
+        public void augmentStatsForUid(@NonNull PackageStats stats, int uid,
+                boolean callerHasStatsPermission) {
+            final int userId = UserHandle.getUserId(uid);
+            final AtomicLong blobsDataSize = new AtomicLong(0);
+            forEachSessionInUser(session -> {
+                if (session.getOwnerUid() == uid) {
+                    blobsDataSize.getAndAdd(session.getSize());
+                }
+            }, userId);
+
+            forEachBlobInUser(blobMetadata -> {
+                if (blobMetadata.isALeasee(uid)) {
+                    if (!blobMetadata.hasOtherLeasees(uid) || !callerHasStatsPermission) {
+                        blobsDataSize.getAndAdd(blobMetadata.getSize());
+                    }
+                }
+            }, userId);
+
+            stats.dataSize += blobsDataSize.get();
+        }
+    }
+
+    private void forEachSessionInUser(Consumer<BlobStoreSession> consumer, int userId) {
+        synchronized (mBlobsLock) {
+            final LongSparseArray<BlobStoreSession> userSessions = getUserSessionsLocked(userId);
+            for (int i = 0, count = userSessions.size(); i < count; ++i) {
+                final BlobStoreSession session = userSessions.valueAt(i);
+                consumer.accept(session);
+            }
+        }
+    }
+
+    private void forEachBlobInUser(Consumer<BlobMetadata> consumer, int userId) {
+        synchronized (mBlobsLock) {
+            final ArrayMap<BlobHandle, BlobMetadata> userBlobs = getUserBlobsLocked(userId);
+            for (int i = 0, count = userBlobs.size(); i < count; ++i) {
+                final BlobMetadata blobMetadata = userBlobs.valueAt(i);
+                consumer.accept(blobMetadata);
+            }
+        }
+    }
+
     private class PackageChangedReceiver extends BroadcastReceiver {
         @Override
         public void onReceive(Context context, Intent intent) {
diff --git a/api/test-current.txt b/api/test-current.txt
index 957794c..b5cb77c 100644
--- a/api/test-current.txt
+++ b/api/test-current.txt
@@ -2224,6 +2224,8 @@
 
   public final class FileUtils {
     method public static boolean contains(java.io.File, java.io.File);
+    method @NonNull public static byte[] digest(@NonNull java.io.File, @NonNull String) throws java.io.IOException, java.security.NoSuchAlgorithmException;
+    method @NonNull public static byte[] digest(@NonNull java.io.InputStream, @NonNull String) throws java.io.IOException, java.security.NoSuchAlgorithmException;
   }
 
   public class HidlMemory implements java.io.Closeable {
diff --git a/core/java/android/os/FileUtils.java b/core/java/android/os/FileUtils.java
index a47fbba..87e7043 100644
--- a/core/java/android/os/FileUtils.java
+++ b/core/java/android/os/FileUtils.java
@@ -743,6 +743,8 @@
      *            {@link MessageDigest#getInstance(String)}.
      * @hide
      */
+    @TestApi
+    @NonNull
     public static byte[] digest(@NonNull File file, @NonNull String algorithm)
             throws IOException, NoSuchAlgorithmException {
         try (FileInputStream in = new FileInputStream(file)) {
@@ -757,6 +759,8 @@
      *            {@link MessageDigest#getInstance(String)}.
      * @hide
      */
+    @TestApi
+    @NonNull
     public static byte[] digest(@NonNull InputStream in, @NonNull String algorithm)
             throws IOException, NoSuchAlgorithmException {
         // TODO: implement kernel optimizations
diff --git a/services/usage/java/com/android/server/usage/StorageStatsManagerInternal.java b/services/usage/java/com/android/server/usage/StorageStatsManagerInternal.java
index a532548..47760ef 100644
--- a/services/usage/java/com/android/server/usage/StorageStatsManagerInternal.java
+++ b/services/usage/java/com/android/server/usage/StorageStatsManagerInternal.java
@@ -34,11 +34,9 @@
     public interface StorageStatsAugmenter {
         void augmentStatsForPackage(@NonNull PackageStats stats,
                 @NonNull String packageName, @UserIdInt int userId,
-                @NonNull String callingPackage);
+                boolean callerHasStatsPermission);
         void augmentStatsForUid(@NonNull PackageStats stats, int uid,
-                @NonNull String callingPackage);
-        void augmentStatsForUser(@NonNull PackageStats stats, @UserIdInt int userId,
-                @NonNull String callingPackage);
+                boolean callerHasStatsPermission);
     }
 
     /**
diff --git a/services/usage/java/com/android/server/usage/StorageStatsService.java b/services/usage/java/com/android/server/usage/StorageStatsService.java
index 18b640f..42ef78c 100644
--- a/services/usage/java/com/android/server/usage/StorageStatsService.java
+++ b/services/usage/java/com/android/server/usage/StorageStatsService.java
@@ -16,10 +16,13 @@
 
 package com.android.server.usage;
 
+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 android.Manifest;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.UserIdInt;
@@ -160,18 +163,33 @@
     }
 
     private void enforceStatsPermission(int callingUid, String callingPackage) {
-        final int mode = mAppOps.noteOp(AppOpsManager.OP_GET_USAGE_STATS,
-                callingUid, callingPackage);
+        final String errMsg = checkStatsPermission(callingUid, callingPackage, true);
+        if (errMsg != null) {
+            throw new SecurityException(errMsg);
+        }
+    }
+
+    private String checkStatsPermission(int callingUid, String callingPackage, boolean noteOp) {
+        final int mode;
+        if (noteOp) {
+            mode = mAppOps.noteOp(AppOpsManager.OP_GET_USAGE_STATS, callingUid, callingPackage);
+        } else {
+            mode = mAppOps.checkOp(AppOpsManager.OP_GET_USAGE_STATS, callingUid, callingPackage);
+        }
         switch (mode) {
             case AppOpsManager.MODE_ALLOWED:
-                return;
+                return null;
             case AppOpsManager.MODE_DEFAULT:
-                mContext.enforceCallingOrSelfPermission(
-                        android.Manifest.permission.PACKAGE_USAGE_STATS, TAG);
-                return;
+                if (mContext.checkCallingOrSelfPermission(
+                        Manifest.permission.PACKAGE_USAGE_STATS) == PERMISSION_GRANTED) {
+                    return null;
+                } else {
+                    return "Caller does not have " + Manifest.permission.PACKAGE_USAGE_STATS
+                            + "; callingPackage=" + callingPackage + ", callingUid=" + callingUid;
+                }
             default:
-                throw new SecurityException("Package " + callingPackage + " from UID " + callingUid
-                        + " blocked by mode " + mode);
+                return "Package " + callingPackage + " from UID " + callingUid
+                        + " blocked by mode " + mode;
         }
     }
 
@@ -280,10 +298,15 @@
             throw new ParcelableException(e);
         }
 
+        final boolean callerHasStatsPermission;
         if (Binder.getCallingUid() == appInfo.uid) {
-            // No permissions required when asking about themselves
+            // No permissions required when asking about themselves. We still check since it is
+            // needed later on but don't throw if caller doesn't have the permission.
+            callerHasStatsPermission = checkStatsPermission(
+                    Binder.getCallingUid(), callingPackage, false) == null;
         } else {
             enforceStatsPermission(Binder.getCallingUid(), callingPackage);
+            callerHasStatsPermission = true;
         }
 
         if (defeatNullable(mPackage.getPackagesForUid(appInfo.uid)).length == 1) {
@@ -313,7 +336,7 @@
             if (volumeUuid == StorageManager.UUID_PRIVATE_INTERNAL) {
                 forEachStorageStatsAugmenter((storageStatsAugmenter) -> {
                     storageStatsAugmenter.augmentStatsForPackage(stats,
-                            packageName, userId, callingPackage);
+                            packageName, userId, callerHasStatsPermission);
                 }, "queryStatsForPackage");
             }
             return translate(stats);
@@ -330,10 +353,15 @@
                     android.Manifest.permission.INTERACT_ACROSS_USERS, TAG);
         }
 
+        final boolean callerHasStatsPermission;
         if (Binder.getCallingUid() == uid) {
-            // No permissions required when asking about themselves
+            // No permissions required when asking about themselves. We still check since it is
+            // needed later on but don't throw if caller doesn't have the permission.
+            callerHasStatsPermission = checkStatsPermission(
+                    Binder.getCallingUid(), callingPackage, false) == null;
         } else {
             enforceStatsPermission(Binder.getCallingUid(), callingPackage);
+            callerHasStatsPermission = true;
         }
 
         final String[] packageNames = defeatNullable(mPackage.getPackagesForUid(uid));
@@ -372,7 +400,7 @@
 
         if (volumeUuid == StorageManager.UUID_PRIVATE_INTERNAL) {
             forEachStorageStatsAugmenter((storageStatsAugmenter) -> {
-                storageStatsAugmenter.augmentStatsForUid(stats, uid, callingPackage);
+                storageStatsAugmenter.augmentStatsForUid(stats, uid, callerHasStatsPermission);
             }, "queryStatsForUid");
         }
         return translate(stats);
@@ -401,12 +429,6 @@
         } catch (InstallerException e) {
             throw new ParcelableException(new IOException(e.getMessage()));
         }
-
-        if (volumeUuid == StorageManager.UUID_PRIVATE_INTERNAL) {
-            forEachStorageStatsAugmenter((storageStatsAugmenter) -> {
-                storageStatsAugmenter.augmentStatsForUser(stats, userId, callingPackage);
-            }, "queryStatsForUser");
-        }
         return translate(stats);
     }
 
diff --git a/tests/BlobStoreTestUtils/Android.bp b/tests/BlobStoreTestUtils/Android.bp
index edd2b43..829c883 100644
--- a/tests/BlobStoreTestUtils/Android.bp
+++ b/tests/BlobStoreTestUtils/Android.bp
@@ -16,5 +16,5 @@
   name: "BlobStoreTestUtils",
   srcs: ["src/**/*.java"],
   static_libs: ["truth-prebuilt"],
-  platform_apis: true
+  sdk_version: "test_current",
 }
\ No newline at end of file