Merge "Move to non-hidden APIs for handleUser and broadcast registration." into sc-dev

apex/appsearch/service/java/com/android/server/appsearch/AppSearchManagerService.java

@@ -17,11 +17,10 @@ package com.android.server.appsearch;
 
 import static android.app.appsearch.AppSearchResult.throwableToFailedResult;
 import static android.os.Process.INVALID_UID;
-import static android.os.UserHandle.USER_NULL;
 
+import android.Manifest;
 import android.annotation.ElapsedRealtimeLong;
 import android.annotation.NonNull;
-import android.app.ActivityManager;
 import android.app.appsearch.AppSearchBatchResult;
 import android.app.appsearch.AppSearchMigrationHelper;
 import android.app.appsearch.AppSearchResult;
@@ -63,6 +62,7 @@ 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.appsearch.util.PackageUtil;
 import com.android.server.usage.StorageStatsManagerLocal;
 import com.android.server.usage.StorageStatsManagerLocal.StorageStatsAugmenter;
 
@@ -124,8 +124,10 @@ public class AppSearchManagerService extends SystemService {
     }
 
     private void registerReceivers() {
-        mContext.registerReceiverAsUser(new UserActionReceiver(), UserHandle.ALL,
+        mContext.registerReceiverForAllUsers(
-                new IntentFilter(Intent.ACTION_USER_REMOVED), /*broadcastPermission=*/ null,
+                new UserActionReceiver(),
+                new IntentFilter(Intent.ACTION_USER_REMOVED),
+                /*broadcastPermission=*/ null,
                 /*scheduler=*/ null);
 
         //TODO(b/145759910) Add a direct callback when user clears the data instead of relying on
@@ -135,8 +137,10 @@ public class AppSearchManagerService extends SystemService {
         packageChangedFilter.addAction(Intent.ACTION_PACKAGE_DATA_CLEARED);
         packageChangedFilter.addDataScheme("package");
         packageChangedFilter.setPriority(IntentFilter.SYSTEM_HIGH_PRIORITY);
-        mContext.registerReceiverAsUser(new PackageChangedReceiver(), UserHandle.ALL,
+        mContext.registerReceiverForAllUsers(
-                packageChangedFilter, /*broadcastPermission=*/ null,
+                new PackageChangedReceiver(),
+                packageChangedFilter,
+                /*broadcastPermission=*/ null,
                 /*scheduler=*/ null);
     }
 
@@ -148,12 +152,13 @@ public class AppSearchManagerService extends SystemService {
 
             switch (intent.getAction()) {
                 case Intent.ACTION_USER_REMOVED:
-                    int userId = intent.getIntExtra(Intent.EXTRA_USER_HANDLE, USER_NULL);
+                    UserHandle userHandle = intent.getParcelableExtra(Intent.EXTRA_USER);
-                    if (userId == USER_NULL) {
+                    if (userHandle == null) {
-                        Log.e(TAG, "userId is missing in the intent: " + intent);
+                        Log.e(TAG, "Extra "
+                                + Intent.EXTRA_USER + " is missing in the intent: " + intent);
                         return;
                     }
-                    handleUserRemoved(UserHandle.of(userId));
+                    handleUserRemoved(userHandle);
                     break;
                 default:
                     Log.e(TAG, "Received unknown intent: " + intent);
@@ -1183,13 +1188,9 @@ public class AppSearchManagerService extends SystemService {
             Objects.requireNonNull(actualCallingUser);
             Objects.requireNonNull(claimedCallingPackage);
 
-            int claimedCallingUid;
+            int claimedCallingUid = PackageUtil.getPackageUidAsUser(
-            try {
+                    mContext, claimedCallingPackage, actualCallingUser);
-                Context claimedCallingContext =
+            if (claimedCallingUid == INVALID_UID) {
-                        mContext.createContextAsUser(actualCallingUser, /*flags=*/ 0);
-                claimedCallingUid = claimedCallingContext.getPackageManager().getPackageUid(
-                        claimedCallingPackage, /*flags=*/ 0);
-            } catch (PackageManager.NameNotFoundException e) {
                 throw new SecurityException(
                         "Specified calling package [" + claimedCallingPackage + "] not found");
             }
@@ -1257,23 +1258,44 @@ public class AppSearchManagerService extends SystemService {
      *
      * <p>Takes care of checking permissions and converting USER_CURRENT to the actual current user.
      *
+     * @param requestedUser The user which the caller is requesting to execute as.
+     * @param callingUid The actual uid of the caller as determined by Binder.
      * @return the user handle that the call should run as. Will always be a concrete user.
      */
     // TODO(b/173553485) verifying that the caller has permission to access target user's data
     // TODO(b/173553485) Handle ACTION_USER_REMOVED broadcast
     // TODO(b/173553485) Implement SystemService.onUserStopping()
     @NonNull
-    private static UserHandle handleIncomingUser(@NonNull UserHandle userHandle, int callingUid) {
+    private UserHandle handleIncomingUser(@NonNull UserHandle requestedUser, int callingUid) {
         int callingPid = Binder.getCallingPid();
-        int finalUserId = ActivityManager.handleIncomingUser(
+        UserHandle callingUser = UserHandle.getUserHandleForUid(callingUid);
+        if (callingUser.equals(requestedUser)) {
+            return requestedUser;
+        }
+        // Duplicates UserController#ensureNotSpecialUser
+        if (requestedUser.getIdentifier() < 0) {
+            throw new IllegalArgumentException(
+                    "Call does not support special user " + requestedUser);
+        }
+        boolean canInteractAcrossUsers = mContext.checkPermission(
+                Manifest.permission.INTERACT_ACROSS_USERS,
                 callingPid,
-                callingUid,
+                callingUid) == PackageManager.PERMISSION_GRANTED;
-                userHandle.getIdentifier(),
+        if (!canInteractAcrossUsers) {
-                /*allowAll=*/ false,
+            canInteractAcrossUsers = mContext.checkPermission(
-                /*requireFull=*/ false,
+                    Manifest.permission.INTERACT_ACROSS_USERS_FULL,
-                /*name=*/ null,
+                    callingPid,
-                /*callerPackage=*/ null);
+                    callingUid) == PackageManager.PERMISSION_GRANTED;
-        return UserHandle.of(finalUserId);
+        }
+        if (canInteractAcrossUsers) {
+            return requestedUser;
+        }
+        throw new SecurityException(
+                "Permission denied while calling from uid " + callingUid
+                        + " with " + requestedUser + "; Need to run as either the calling user ("
+                        + callingUser + "), or with one of the following permissions: "
+                        + Manifest.permission.INTERACT_ACROSS_USERS + " or "
+                        + Manifest.permission.INTERACT_ACROSS_USERS_FULL);
     }
 
     // TODO(b/179160886): Cache the previous storage stats.

apex/appsearch/service/java/com/android/server/appsearch/ImplInstanceManager.java

@@ -172,11 +172,6 @@ public final class ImplInstanceManager {
             @Nullable AppSearchLogger logger)
             throws AppSearchException {
         File appSearchDir = getAppSearchDir(userHandle);
-        // TODO(b/181787682): Swap AppSearchImpl and VisibilityStore to accept a UserHandle too
+        return AppSearchImpl.create(appSearchDir, userContext, /*logger=*/ null);
-        return AppSearchImpl.create(
-                appSearchDir,
-                userContext,
-                userHandle.getIdentifier(),
-                /*logger=*/ null);
     }
 }

apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/AppSearchImpl.java

@@ -201,7 +201,6 @@ public final class AppSearchImpl implements Closeable {
     public static AppSearchImpl create(
             @NonNull File icingDir,
             @NonNull Context userContext,
-            int userId,
             @Nullable AppSearchLogger logger)
             throws AppSearchException {
         Objects.requireNonNull(icingDir);
@@ -213,8 +212,7 @@ public final class AppSearchImpl implements Closeable {
             initStatsBuilder = new InitializeStats.Builder();
         }
 
-        AppSearchImpl appSearchImpl =
+        AppSearchImpl appSearchImpl = new AppSearchImpl(icingDir, userContext, initStatsBuilder);
-                new AppSearchImpl(icingDir, userContext, userId, initStatsBuilder);
 
         long prepareVisibilityStoreLatencyStartMillis = SystemClock.elapsedRealtime();
         appSearchImpl.initializeVisibilityStore();
@@ -238,7 +236,6 @@ public final class AppSearchImpl implements Closeable {
     private AppSearchImpl(
             @NonNull File icingDir,
             @NonNull Context userContext,
-            int userId,
             @Nullable InitializeStats.Builder initStatsBuilder)
             throws AppSearchException {
         mReadWriteLock.writeLock().lock();
@@ -256,7 +253,7 @@ public final class AppSearchImpl implements Closeable {
                     "Constructing IcingSearchEngine, response",
                     Objects.hashCode(mIcingSearchEngineLocked));
 
-            mVisibilityStoreLocked = new VisibilityStore(this, userContext, userId);
+            mVisibilityStoreLocked = new VisibilityStore(this, userContext);
 
             // The core initialization procedure. If any part of this fails, we bail into
             // resetLocked(), deleting all data (but hopefully allowing AppSearchImpl to come up).

apex/appsearch/service/java/com/android/server/appsearch/stats/PlatformLogger.java

@@ -20,7 +20,6 @@ import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.app.appsearch.exceptions.AppSearchException;
 import android.content.Context;
-import android.content.pm.PackageManager;
 import android.os.Process;
 import android.os.SystemClock;
 import android.os.UserHandle;
@@ -36,6 +35,7 @@ import com.android.server.appsearch.external.localstorage.stats.InitializeStats;
 import com.android.server.appsearch.external.localstorage.stats.PutDocumentStats;
 import com.android.server.appsearch.external.localstorage.stats.RemoveStats;
 import com.android.server.appsearch.external.localstorage.stats.SearchStats;
+import com.android.server.appsearch.util.PackageUtil;
 
 import java.io.UnsupportedEncodingException;
 import java.security.MessageDigest;
@@ -493,21 +493,13 @@ public final class PlatformLogger implements AppSearchLogger {
     @GuardedBy("mLock")
     private int getPackageUidAsUserLocked(@NonNull String packageName) {
         Integer packageUid = mPackageUidCacheLocked.get(packageName);
-        if (packageUid != null) {
+        if (packageUid == null) {
-            return packageUid;
+            packageUid = PackageUtil.getPackageUidAsUser(mContext, packageName, mUserHandle);
+            if (packageUid != Process.INVALID_UID) {
+                mPackageUidCacheLocked.put(packageName, packageUid);
+            }
         }
-
+        return packageUid;
-        // TODO(b/173532925) since VisibilityStore has the same method, we can make this a
-        //  utility function
-        try {
-            packageUid = mContext.getPackageManager().getPackageUidAsUser(
-                    packageName, mUserHandle.getIdentifier());
-            mPackageUidCacheLocked.put(packageName, packageUid);
-            return packageUid;
-        } catch (PackageManager.NameNotFoundException e) {
-            // Package doesn't exist, continue
-        }
-        return Process.INVALID_UID;
     }
 
     //

apex/appsearch/service/java/com/android/server/appsearch/util/PackageUtil.java

@@ -0,0 +1,55 @@
+/*
+ * 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.appsearch.util;
+
+import android.annotation.NonNull;
+import android.content.Context;
+import android.content.pm.PackageManager;
+import android.os.Process;
+import android.os.UserHandle;
+
+/**
+ * Utilities for interacting with {@link android.content.pm.PackageManager},
+ * {@link android.os.UserHandle}, and other parts of dealing with apps and binder.
+ *
+ * @hide
+ */
+public class PackageUtil {
+    private PackageUtil() {}
+
+    /**
+     * Finds the UID of the {@code packageName}. Returns {@link Process#INVALID_UID} if unable to
+     * find the UID.
+     */
+    public static int getPackageUidAsUser(
+            @NonNull Context context, @NonNull String packageName, @NonNull UserHandle user) {
+        Context userContext = context.createContextAsUser(user, /*flags=*/ 0);
+        return getPackageUid(userContext, packageName);
+    }
+
+    /**
+     * Finds the UID of the {@code packageName} in the given {@code context}. Returns
+     * {@link Process#INVALID_UID} if unable to find the UID.
+     */
+    public static int getPackageUid(@NonNull Context context, @NonNull String packageName) {
+        try {
+            return context.getPackageManager().getPackageUid(packageName, /*flags=*/ 0);
+        } catch (PackageManager.NameNotFoundException e) {
+            return Process.INVALID_UID;
+        }
+    }
+}

apex/appsearch/service/java/com/android/server/appsearch/visibilitystore/VisibilityStore.java

@@ -18,7 +18,6 @@ package com.android.server.appsearch.visibilitystore;
 import static android.Manifest.permission.READ_GLOBAL_APP_SEARCH_DATA;
 
 import android.annotation.NonNull;
-import android.annotation.UserIdInt;
 import android.app.appsearch.AppSearchResult;
 import android.app.appsearch.AppSearchSchema;
 import android.app.appsearch.GenericDocument;
@@ -27,7 +26,6 @@ import android.app.appsearch.PackageIdentifier;
 import android.app.appsearch.exceptions.AppSearchException;
 import android.content.Context;
 import android.content.pm.PackageManager;
-import android.os.Process;
 import android.os.UserHandle;
 import android.util.ArrayMap;
 import android.util.ArraySet;
@@ -35,6 +33,7 @@ import android.util.ArraySet;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.server.appsearch.external.localstorage.AppSearchImpl;
 import com.android.server.appsearch.external.localstorage.util.PrefixUtil;
+import com.android.server.appsearch.util.PackageUtil;
 
 import com.google.android.icing.proto.PersistType;
 
@@ -67,9 +66,6 @@ import java.util.Set;
  * @hide
  */
 public class VisibilityStore {
-
-    private static final String TAG = "AppSearchVisibilityStore";
-
     /** No-op user id that won't have any visibility settings. */
     public static final int NO_OP_USER_ID = -1;
 
@@ -95,9 +91,6 @@ public class VisibilityStore {
     // Context of the user that the call is being made as.
     private final Context mUserContext;
 
-    // User ID of the caller who we're checking visibility settings for.
-    private final int mUserId;
-
     /** Stores the schemas that are platform-hidden. All values are prefixed. */
     private final NotPlatformSurfaceableMap mNotPlatformSurfaceableMap =
             new NotPlatformSurfaceableMap();
@@ -112,13 +105,9 @@ public class VisibilityStore {
      * @param appSearchImpl AppSearchImpl instance
      * @param userContext Context of the user that the call is being made as
      */
-    public VisibilityStore(
+    public VisibilityStore(@NonNull AppSearchImpl appSearchImpl, @NonNull Context userContext) {
-            @NonNull AppSearchImpl appSearchImpl,
-            @NonNull Context userContext,
-            @UserIdInt int userId) {
         mAppSearchImpl = appSearchImpl;
-        mUserContext = userContext;
+        mUserContext = Objects.requireNonNull(userContext);
-        mUserId = userId;
     }
 
     /**
@@ -348,6 +337,9 @@ public class VisibilityStore {
         Set<PackageIdentifier> packageIdentifiers =
                 mPackageAccessibleMap.getAccessiblePackages(
                         packageName, databaseName, prefixedSchema);
+        if (packageIdentifiers.isEmpty()) {
+            return false;
+        }
         for (PackageIdentifier packageIdentifier : packageIdentifiers) {
             // TODO(b/169883602): Consider caching the UIDs of packages. Looking this up in the
             // package manager could be costly. We would also need to update the cache on
@@ -357,9 +349,10 @@ public class VisibilityStore {
             // the callerUid since clients can createContextAsUser with some other user, and then
             // make calls to us. So just check if the appId portion of the uid is the same. This is
             // essentially UserHandle.isSameApp, but that's not a system API for us to use.
-            int callerAppId = UserHandle.getAppId((callerUid));
+            int callerAppId = UserHandle.getAppId(callerUid);
-            int userAppId =
+            int packageUid =
-                    UserHandle.getAppId(getPackageUidAsUser(packageIdentifier.getPackageName()));
+                    PackageUtil.getPackageUid(mUserContext, packageIdentifier.getPackageName());
+            int userAppId = UserHandle.getAppId(packageUid);
             if (callerAppId != userAppId) {
                 continue;
             }
@@ -401,17 +394,4 @@ public class VisibilityStore {
             @NonNull String packageName, @NonNull String databaseName) {
         return ID_PREFIX + PrefixUtil.createPrefix(packageName, databaseName);
     }
-
-    /**
-     * Finds the UID of the {@code packageName}. Returns {@link Process#INVALID_UID} if unable to
-     * find the UID.
-     */
-    private int getPackageUidAsUser(@NonNull String packageName) {
-        try {
-            return mUserContext.getPackageManager().getPackageUidAsUser(packageName, mUserId);
-        } catch (PackageManager.NameNotFoundException e) {
-            // Package doesn't exist, continue
-        }
-        return Process.INVALID_UID;
-    }
 }

services/tests/servicestests/src/com/android/server/appsearch/AppSearchImplPlatformTest.java

@@ -25,11 +25,18 @@ import static android.content.pm.PackageManager.PERMISSION_GRANTED;
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.when;
+
+import android.annotation.NonNull;
 import android.app.appsearch.AppSearchSchema;
 import android.app.appsearch.PackageIdentifier;
 import android.content.Context;
 import android.content.ContextWrapper;
 import android.content.pm.PackageManager;
+import android.os.UserHandle;
+import android.util.ArrayMap;
 
 import androidx.test.core.app.ApplicationProvider;
 
@@ -43,13 +50,15 @@ import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.TemporaryFolder;
+import org.mockito.Mockito;
 
 import java.util.Collections;
+import java.util.Map;
 
 /** This tests AppSearchImpl when it's running with a platform-backed VisibilityStore. */
 public class AppSearchImplPlatformTest {
     @Rule public TemporaryFolder mTemporaryFolder = new TemporaryFolder();
-    private MockPackageManager mMockPackageManager = new MockPackageManager();
+    private final Map<UserHandle, PackageManager> mMockPackageManagers = new ArrayMap<>();
     private Context mContext;
     private AppSearchImpl mAppSearchImpl;
     private int mGlobalQuerierUid;
@@ -57,20 +66,28 @@ public class AppSearchImplPlatformTest {
     @Before
     public void setUp() throws Exception {
         Context context = ApplicationProvider.getApplicationContext();
-        mContext =
+        mContext = new ContextWrapper(context) {
-                new ContextWrapper(context) {
+            @Override
+            public Context createContextAsUser(UserHandle user, int flags) {
+                return new ContextWrapper(super.createContextAsUser(user, flags)) {
                     @Override
                     public PackageManager getPackageManager() {
-                        return mMockPackageManager.getMockPackageManager();
+                        return getMockPackageManager(user);
                     }
                 };
+            }
+
+            @Override
+            public PackageManager getPackageManager() {
+                return createContextAsUser(getUser(), /*flags=*/ 0).getPackageManager();
+            }
+        };
 
         // Give ourselves global query permissions
         mAppSearchImpl =
                 AppSearchImpl.create(
                         mTemporaryFolder.newFolder(),
                         mContext,
-                        mContext.getUserId(),
                         /*logger=*/ null);
 
         mGlobalQuerierUid =
@@ -85,14 +102,19 @@ public class AppSearchImplPlatformTest {
         int uidFoo = 1;
 
         // Make sure foo package will pass package manager checks.
-        mMockPackageManager.mockGetPackageUidAsUser(packageNameFoo, mContext.getUserId(), uidFoo);
+        PackageManager mockPackageManager = getMockPackageManager(mContext.getUser());
-        mMockPackageManager.mockAddSigningCertificate(packageNameFoo, sha256CertFoo);
+        when(mockPackageManager.getPackageUid(eq(packageNameFoo), /*flags=*/ anyInt()))
+                .thenReturn(uidFoo);
+        when(mockPackageManager.hasSigningCertificate(
+                packageNameFoo, sha256CertFoo, PackageManager.CERT_INPUT_SHA256))
+                .thenReturn(true);
 
         // Make sure we have global query privileges and "foo" doesn't
-        mMockPackageManager.mockCheckPermission(
+        when(mockPackageManager.checkPermission(
-                READ_GLOBAL_APP_SEARCH_DATA, mContext.getPackageName(), PERMISSION_GRANTED);
+                READ_GLOBAL_APP_SEARCH_DATA, mContext.getPackageName()))
-        mMockPackageManager.mockCheckPermission(
+                .thenReturn(PERMISSION_GRANTED);
-                READ_GLOBAL_APP_SEARCH_DATA, packageNameFoo, PERMISSION_DENIED);
+        when(mockPackageManager.checkPermission(READ_GLOBAL_APP_SEARCH_DATA, packageNameFoo))
+                .thenReturn(PERMISSION_DENIED);
 
         // Set schema1
         String prefix = PrefixUtil.createPrefix("package", "database");
@@ -145,19 +167,16 @@ public class AppSearchImplPlatformTest {
                 /*schemaVersion=*/ 0);
 
         // Check that "schema1" still has the same visibility settings
-        SystemUtil.runWithShellPermissionIdentity(
+        SystemUtil.runWithShellPermissionIdentity(() -> assertThat(
-                () -> {
+                mAppSearchImpl
-                    assertThat(
+                        .getVisibilityStoreLocked()
-                                    mAppSearchImpl
+                        .isSchemaSearchableByCaller(
-                                            .getVisibilityStoreLocked()
+                                "package",
-                                            .isSchemaSearchableByCaller(
+                                "database",
-                                                    "package",
+                                prefix + "schema1",
-                                                    "database",
+                                mContext.getPackageName(),
-                                                    prefix + "schema1",
+                                mGlobalQuerierUid))
-                                                    mContext.getPackageName(),
+                        .isFalse(),
-                                                    mGlobalQuerierUid))
-                            .isFalse();
-                },
                 READ_GLOBAL_APP_SEARCH_DATA);
 
         assertThat(
@@ -172,19 +191,16 @@ public class AppSearchImplPlatformTest {
                 .isTrue();
 
         // "schema2" has default visibility settings
-        SystemUtil.runWithShellPermissionIdentity(
+        SystemUtil.runWithShellPermissionIdentity(() -> assertThat(
-                () -> {
+                mAppSearchImpl
-                    assertThat(
+                        .getVisibilityStoreLocked()
-                                    mAppSearchImpl
+                        .isSchemaSearchableByCaller(
-                                            .getVisibilityStoreLocked()
+                                "package",
-                                            .isSchemaSearchableByCaller(
+                                "database",
-                                                    "package",
+                                prefix + "schema2",
-                                                    "database",
+                                mContext.getPackageName(),
-                                                    prefix + "schema2",
+                                mGlobalQuerierUid))
-                                                    mContext.getPackageName(),
+                        .isTrue(),
-                                                    mGlobalQuerierUid))
-                            .isTrue();
-                },
                 READ_GLOBAL_APP_SEARCH_DATA);
 
         assertThat(
@@ -207,14 +223,19 @@ public class AppSearchImplPlatformTest {
         int uidFoo = 1;
 
         // Make sure foo package will pass package manager checks.
-        mMockPackageManager.mockGetPackageUidAsUser(packageNameFoo, mContext.getUserId(), uidFoo);
+        PackageManager mockPackageManager = getMockPackageManager(mContext.getUser());
-        mMockPackageManager.mockAddSigningCertificate(packageNameFoo, sha256CertFoo);
+        when(mockPackageManager.getPackageUid(eq(packageNameFoo), /*flags=*/ anyInt()))
+                .thenReturn(uidFoo);
+        when(mockPackageManager.hasSigningCertificate(
+                packageNameFoo, sha256CertFoo, PackageManager.CERT_INPUT_SHA256))
+                .thenReturn(true);
 
         // Make sure we have global query privileges and "foo" doesn't
-        mMockPackageManager.mockCheckPermission(
+        when(mockPackageManager.checkPermission(
-                READ_GLOBAL_APP_SEARCH_DATA, mContext.getPackageName(), PERMISSION_GRANTED);
+                READ_GLOBAL_APP_SEARCH_DATA, mContext.getPackageName()))
-        mMockPackageManager.mockCheckPermission(
+                .thenReturn(PERMISSION_GRANTED);
-                READ_GLOBAL_APP_SEARCH_DATA, packageNameFoo, PERMISSION_DENIED);
+        when(mockPackageManager.checkPermission(READ_GLOBAL_APP_SEARCH_DATA, packageNameFoo))
+                .thenReturn(PERMISSION_DENIED);
 
         String prefix = PrefixUtil.createPrefix("package", "database");
         mAppSearchImpl.setSchema(
@@ -320,8 +341,10 @@ public class AppSearchImplPlatformTest {
     @Test
     public void testSetSchema_defaultPlatformVisible() throws Exception {
         // Make sure we have global query privileges
-        mMockPackageManager.mockCheckPermission(
+        PackageManager mockPackageManager = getMockPackageManager(mContext.getUser());
-                READ_GLOBAL_APP_SEARCH_DATA, mContext.getPackageName(), PERMISSION_GRANTED);
+        when(mockPackageManager.checkPermission(
+                READ_GLOBAL_APP_SEARCH_DATA, mContext.getPackageName()))
+                .thenReturn(PERMISSION_GRANTED);
 
         String prefix = PrefixUtil.createPrefix("package", "database");
         mAppSearchImpl.setSchema(
@@ -348,8 +371,10 @@ public class AppSearchImplPlatformTest {
     @Test
     public void testSetSchema_platformHidden() throws Exception {
         // Make sure we have global query privileges
-        mMockPackageManager.mockCheckPermission(
+        PackageManager mockPackageManager = getMockPackageManager(mContext.getUser());
-                READ_GLOBAL_APP_SEARCH_DATA, mContext.getPackageName(), PERMISSION_GRANTED);
+        when(mockPackageManager.checkPermission(
+                READ_GLOBAL_APP_SEARCH_DATA, mContext.getPackageName()))
+                .thenReturn(PERMISSION_GRANTED);
 
         String prefix = PrefixUtil.createPrefix("package", "database");
         mAppSearchImpl.setSchema(
@@ -378,8 +403,9 @@ public class AppSearchImplPlatformTest {
         String packageName = "com.package";
 
         // Make sure package doesn't global query privileges
-        mMockPackageManager.mockCheckPermission(
+        PackageManager mockPackageManager = getMockPackageManager(mContext.getUser());
-                READ_GLOBAL_APP_SEARCH_DATA, packageName, PERMISSION_DENIED);
+        when(mockPackageManager.checkPermission(
+                READ_GLOBAL_APP_SEARCH_DATA, packageName)).thenReturn(PERMISSION_DENIED);
 
         String prefix = PrefixUtil.createPrefix("package", "database");
         mAppSearchImpl.setSchema(
@@ -410,12 +436,16 @@ public class AppSearchImplPlatformTest {
         int uidFoo = 1;
 
         // Make sure foo package will pass package manager checks.
-        mMockPackageManager.mockGetPackageUidAsUser(packageNameFoo, mContext.getUserId(), uidFoo);
+        PackageManager mockPackageManager = getMockPackageManager(mContext.getUser());
-        mMockPackageManager.mockAddSigningCertificate(packageNameFoo, sha256CertFoo);
+        when(mockPackageManager.getPackageUid(eq(packageNameFoo), /*flags=*/ anyInt()))
+                .thenReturn(uidFoo);
+        when(mockPackageManager.hasSigningCertificate(
+                packageNameFoo, sha256CertFoo, PackageManager.CERT_INPUT_SHA256))
+                .thenReturn(true);
 
         // Make sure foo doesn't have global query privileges
-        mMockPackageManager.mockCheckPermission(
+        when(mockPackageManager.checkPermission(READ_GLOBAL_APP_SEARCH_DATA, packageNameFoo))
-                READ_GLOBAL_APP_SEARCH_DATA, packageNameFoo, PERMISSION_DENIED);
+                .thenReturn(PERMISSION_DENIED);
 
         String prefix = PrefixUtil.createPrefix("package", "database");
         mAppSearchImpl.setSchema(
@@ -439,4 +469,14 @@ public class AppSearchImplPlatformTest {
                                         uidFoo))
                 .isTrue();
     }
+
+    @NonNull
+    private PackageManager getMockPackageManager(@NonNull UserHandle user) {
+        PackageManager pm = mMockPackageManagers.get(user);
+        if (pm == null) {
+            pm = Mockito.mock(PackageManager.class);
+            mMockPackageManagers.put(user, pm);
+        }
+        return pm;
+    }
 }

services/tests/servicestests/src/com/android/server/appsearch/MockPackageManager.java

@@ -1,94 +0,0 @@
-/*
- * 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.
- */
-
-// TODO(b/169883602): This is purposely a different package from the path so that AppSearchImplTest
-// can use it without an extra import. This should be moved into a proper package once
-// AppSearchImpl-VisibilityStore's dependencies are refactored.
-package com.android.server.appsearch.external.localstorage;
-
-import static org.mockito.ArgumentMatchers.anyInt;
-import static org.mockito.ArgumentMatchers.eq;
-import static org.mockito.Mockito.when;
-
-import android.annotation.NonNull;
-import android.annotation.UserIdInt;
-import android.content.pm.PackageManager;
-
-import org.mockito.Mock;
-import org.mockito.MockitoAnnotations;
-
-/** Mock to help test package name, UID, and certificate verification. */
-public class MockPackageManager {
-
-    @Mock private PackageManager mMockPackageManager;
-
-    public MockPackageManager() {
-        MockitoAnnotations.initMocks(this);
-    }
-
-    @NonNull
-    public PackageManager getMockPackageManager() {
-        return mMockPackageManager;
-    }
-
-    /** Mock a checkPermission call. */
-    public void mockCheckPermission(String permission, String packageName, int permissionResult) {
-        when(mMockPackageManager.checkPermission(permission, packageName))
-                .thenReturn(permissionResult);
-    }
-
-    /** Mock a NameNotFoundException if the package name isn't installed. */
-    public void mockThrowsNameNotFoundException(String packageName) {
-        try {
-            when(mMockPackageManager.getPackageUidAsUser(eq(packageName), /*userId=*/ anyInt()))
-                    .thenThrow(new PackageManager.NameNotFoundException());
-            when(mMockPackageManager.getPackageUidAsUser(
-                            eq(packageName), /*flags=*/ anyInt(), /*userId=*/ anyInt()))
-                    .thenThrow(new PackageManager.NameNotFoundException());
-        } catch (PackageManager.NameNotFoundException e) {
-            // Shouldn't ever happen since we're mocking the exception
-            e.printStackTrace();
-        }
-    }
-
-    /** Mocks that {@code uid} contains the {@code packageName} */
-    public void mockGetPackageUidAsUser(String packageName, @UserIdInt int callerUserId, int uid) {
-        try {
-            when(mMockPackageManager.getPackageUidAsUser(eq(packageName), eq(callerUserId)))
-                    .thenReturn(uid);
-            when(mMockPackageManager.getPackageUidAsUser(
-                            eq(packageName), /*flags=*/ anyInt(), eq(callerUserId)))
-                    .thenReturn(uid);
-        } catch (PackageManager.NameNotFoundException e) {
-            // Shouldn't ever happen since we're mocking the method.
-            e.printStackTrace();
-        }
-    }
-
-    /** Mocks that {@code packageName} has been signed with {@code sha256Cert}. */
-    public void mockAddSigningCertificate(String packageName, byte[] sha256Cert) {
-        when(mMockPackageManager.hasSigningCertificate(
-                        packageName, sha256Cert, PackageManager.CERT_INPUT_SHA256))
-                .thenReturn(true);
-    }
-
-    /** Mocks that {@code packageName} has NOT been signed with {@code sha256Cert}. */
-    public void mockRemoveSigningCertificate(String packageName, byte[] sha256Cert) {
-        when(mMockPackageManager.hasSigningCertificate(
-                        packageName, sha256Cert, PackageManager.CERT_INPUT_SHA256))
-                .thenReturn(false);
-    }
-}

services/tests/servicestests/src/com/android/server/appsearch/VisibilityStoreTest.java

@@ -25,10 +25,17 @@ import static android.content.pm.PackageManager.PERMISSION_GRANTED;
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.when;
+
+import android.annotation.NonNull;
 import android.app.appsearch.PackageIdentifier;
 import android.content.Context;
 import android.content.ContextWrapper;
 import android.content.pm.PackageManager;
+import android.os.UserHandle;
+import android.util.ArrayMap;
 
 import androidx.test.core.app.ApplicationProvider;
 
@@ -43,39 +50,47 @@ import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.TemporaryFolder;
+import org.mockito.Mockito;
 
 import java.util.Collections;
+import java.util.Map;
 
 public class VisibilityStoreTest {
-
     @Rule public TemporaryFolder mTemporaryFolder = new TemporaryFolder();
-    private MockPackageManager mMockPackageManager = new MockPackageManager();
+    private final Map<UserHandle, PackageManager> mMockPackageManagers = new ArrayMap<>();
     private Context mContext;
-    private AppSearchImpl mAppSearchImpl;
     private VisibilityStore mVisibilityStore;
     private int mUid;
 
     @Before
     public void setUp() throws Exception {
         Context context = ApplicationProvider.getApplicationContext();
-        mContext =
+        mContext = new ContextWrapper(context) {
-                new ContextWrapper(context) {
+            @Override
+            public Context createContextAsUser(UserHandle user, int flags) {
+                return new ContextWrapper(super.createContextAsUser(user, flags)) {
                     @Override
                     public PackageManager getPackageManager() {
-                        return mMockPackageManager.getMockPackageManager();
+                        return getMockPackageManager(user);
                     }
                 };
+            }
+
+            @Override
+            public PackageManager getPackageManager() {
+                return createContextAsUser(getUser(), /*flags=*/ 0).getPackageManager();
+            }
+        };
 
         // Give ourselves global query permissions
-        mAppSearchImpl =
+        AppSearchImpl appSearchImpl =
                 AppSearchImpl.create(
                         mTemporaryFolder.newFolder(),
                         mContext,
-                        mContext.getUserId(),
                         /*logger=*/ null);
 
         mUid = mContext.getPackageManager().getPackageUid(mContext.getPackageName(), /*flags=*/ 0);
-        mVisibilityStore = mAppSearchImpl.getVisibilityStoreLocked();
+        mVisibilityStore = appSearchImpl.getVisibilityStoreLocked();
     }
 
     /**
@@ -103,8 +118,10 @@ public class VisibilityStoreTest {
     @Test
     public void testSetVisibility_platformSurfaceable() throws Exception {
         // Make sure we have global query privileges
-        mMockPackageManager.mockCheckPermission(
+        PackageManager mockPackageManager = getMockPackageManager(mContext.getUser());
-                READ_GLOBAL_APP_SEARCH_DATA, mContext.getPackageName(), PERMISSION_GRANTED);
+        when(mockPackageManager
+                .checkPermission(READ_GLOBAL_APP_SEARCH_DATA, mContext.getPackageName()))
+                .thenReturn(PERMISSION_GRANTED);
 
         mVisibilityStore.setVisibility(
                 "package",
@@ -210,10 +227,13 @@ public class VisibilityStoreTest {
         int uidNotFooOrBar = 3;
 
         // Make sure none of them have global query privileges
-        mMockPackageManager.mockCheckPermission(
+        PackageManager mockPackageManager = getMockPackageManager(mContext.getUser());
-                READ_GLOBAL_APP_SEARCH_DATA, packageNameFoo, PERMISSION_DENIED);
+        when(mockPackageManager
-        mMockPackageManager.mockCheckPermission(
+                .checkPermission(READ_GLOBAL_APP_SEARCH_DATA, packageNameFoo))
-                READ_GLOBAL_APP_SEARCH_DATA, packageNameBar, PERMISSION_DENIED);
+                .thenReturn(PERMISSION_DENIED);
+        when(mockPackageManager
+                .checkPermission(READ_GLOBAL_APP_SEARCH_DATA, packageNameBar))
+                .thenReturn(PERMISSION_DENIED);
 
         // By default, a schema isn't package accessible.
         assertThat(
@@ -237,32 +257,43 @@ public class VisibilityStoreTest {
                         ImmutableList.of(new PackageIdentifier(packageNameBar, sha256CertBar))));
 
         // Should fail if PackageManager doesn't see that it has the proper certificate
-        mMockPackageManager.mockGetPackageUidAsUser(packageNameFoo, mContext.getUserId(), uidFoo);
+        when(mockPackageManager.getPackageUid(eq(packageNameFoo), /*flags=*/ anyInt()))
-        mMockPackageManager.mockRemoveSigningCertificate(packageNameFoo, sha256CertFoo);
+                .thenReturn(uidFoo);
+        when(mockPackageManager.hasSigningCertificate(
+                packageNameFoo, sha256CertFoo, PackageManager.CERT_INPUT_SHA256))
+                .thenReturn(false);
         assertThat(
                         mVisibilityStore.isSchemaSearchableByCaller(
                                 "package", "database", "prefix/schemaFoo", packageNameFoo, uidFoo))
                 .isFalse();
 
         // Should fail if PackageManager doesn't think the package belongs to the uid
-        mMockPackageManager.mockGetPackageUidAsUser(
+        when(mockPackageManager.getPackageUid(eq(packageNameFoo), /*flags=*/ anyInt()))
-                packageNameFoo, mContext.getUserId(), uidNotFooOrBar);
+                .thenReturn(uidNotFooOrBar);
-        mMockPackageManager.mockAddSigningCertificate(packageNameFoo, sha256CertFoo);
+        when(mockPackageManager.hasSigningCertificate(
+                packageNameFoo, sha256CertFoo, PackageManager.CERT_INPUT_SHA256))
+                .thenReturn(true);
         assertThat(
                         mVisibilityStore.isSchemaSearchableByCaller(
                                 "package", "database", "prefix/schemaFoo", packageNameFoo, uidFoo))
                 .isFalse();
 
         // But if uid and certificate match, then we should have access
-        mMockPackageManager.mockGetPackageUidAsUser(packageNameFoo, mContext.getUserId(), uidFoo);
+        when(mockPackageManager.getPackageUid(eq(packageNameFoo), /*flags=*/ anyInt()))
-        mMockPackageManager.mockAddSigningCertificate(packageNameFoo, sha256CertFoo);
+                .thenReturn(uidFoo);
+        when(mockPackageManager.hasSigningCertificate(
+                packageNameFoo, sha256CertFoo, PackageManager.CERT_INPUT_SHA256))
+                .thenReturn(true);
         assertThat(
                         mVisibilityStore.isSchemaSearchableByCaller(
                                 "package", "database", "prefix/schemaFoo", packageNameFoo, uidFoo))
                 .isTrue();
 
-        mMockPackageManager.mockGetPackageUidAsUser(packageNameBar, mContext.getUserId(), uidBar);
+        when(mockPackageManager.getPackageUid(eq(packageNameBar), /*flags=*/ anyInt()))
-        mMockPackageManager.mockAddSigningCertificate(packageNameBar, sha256CertBar);
+                .thenReturn(uidBar);
+        when(mockPackageManager.hasSigningCertificate(
+                packageNameBar, sha256CertBar, PackageManager.CERT_INPUT_SHA256))
+                .thenReturn(true);
         assertThat(
                         mVisibilityStore.isSchemaSearchableByCaller(
                                 "package", "database", "prefix/schemaBar", packageNameBar, uidBar))
@@ -278,15 +309,21 @@ public class VisibilityStoreTest {
                         "prefix/schemaFoo",
                         ImmutableList.of(new PackageIdentifier(packageNameFoo, sha256CertFoo))));
 
-        mMockPackageManager.mockGetPackageUidAsUser(packageNameFoo, mContext.getUserId(), uidFoo);
+        when(mockPackageManager.getPackageUid(eq(packageNameFoo), /*flags=*/ anyInt()))
-        mMockPackageManager.mockAddSigningCertificate(packageNameFoo, sha256CertFoo);
+                .thenReturn(uidFoo);
+        when(mockPackageManager.hasSigningCertificate(
+                packageNameFoo, sha256CertFoo, PackageManager.CERT_INPUT_SHA256))
+                .thenReturn(true);
         assertThat(
                         mVisibilityStore.isSchemaSearchableByCaller(
                                 "package", "database", "prefix/schemaFoo", packageNameFoo, uidFoo))
                 .isTrue();
 
-        mMockPackageManager.mockGetPackageUidAsUser(packageNameBar, mContext.getUserId(), uidBar);
+        when(mockPackageManager.getPackageUid(eq(packageNameBar), /*flags=*/ anyInt()))
-        mMockPackageManager.mockAddSigningCertificate(packageNameBar, sha256CertBar);
+                .thenReturn(uidBar);
+        when(mockPackageManager.hasSigningCertificate(
+                packageNameBar, sha256CertBar, PackageManager.CERT_INPUT_SHA256))
+                .thenReturn(true);
         assertThat(
                         mVisibilityStore.isSchemaSearchableByCaller(
                                 "package", "database", "prefix/schemaBar", packageNameBar, uidBar))
@@ -302,11 +339,13 @@ public class VisibilityStoreTest {
         int uidFoo = 1;
 
         // Pretend we can't find the Foo package.
-        mMockPackageManager.mockThrowsNameNotFoundException(packageNameFoo);
+        PackageManager mockPackageManager = getMockPackageManager(mContext.getUser());
+        when(mockPackageManager.getPackageUid(eq(packageNameFoo), /*flags=*/ anyInt()))
+                .thenThrow(new PackageManager.NameNotFoundException());
 
         // Make sure "foo" doesn't have global query privileges
-        mMockPackageManager.mockCheckPermission(
+        when(mockPackageManager.checkPermission(READ_GLOBAL_APP_SEARCH_DATA, packageNameFoo))
-                READ_GLOBAL_APP_SEARCH_DATA, packageNameFoo, PERMISSION_DENIED);
+                .thenReturn(PERMISSION_DENIED);
 
         // Grant package access
         mVisibilityStore.setVisibility(
@@ -332,10 +371,12 @@ public class VisibilityStoreTest {
         int uidFoo = 1;
 
         // Set it up such that the test package has global query privileges, but "foo" doesn't.
-        mMockPackageManager.mockCheckPermission(
+        PackageManager mockPackageManager = getMockPackageManager(mContext.getUser());
-                READ_GLOBAL_APP_SEARCH_DATA, mContext.getPackageName(), PERMISSION_GRANTED);
+        when(mockPackageManager.checkPermission(
-        mMockPackageManager.mockCheckPermission(
+                READ_GLOBAL_APP_SEARCH_DATA, mContext.getPackageName()))
-                READ_GLOBAL_APP_SEARCH_DATA, packageNameFoo, PERMISSION_DENIED);
+                .thenReturn(PERMISSION_GRANTED);
+        when(mockPackageManager.checkPermission(READ_GLOBAL_APP_SEARCH_DATA, packageNameFoo))
+                .thenReturn(PERMISSION_DENIED);
 
         mVisibilityStore.setVisibility(
                 /*packageName=*/ "",
@@ -354,8 +395,11 @@ public class VisibilityStoreTest {
                                 mUid))
                 .isTrue();
 
-        mMockPackageManager.mockGetPackageUidAsUser(packageNameFoo, mContext.getUserId(), uidFoo);
+        when(mockPackageManager.getPackageUid(eq(packageNameFoo), /*flags=*/ anyInt()))
-        mMockPackageManager.mockAddSigningCertificate(packageNameFoo, sha256CertFoo);
+                .thenReturn(uidFoo);
+        when(mockPackageManager.hasSigningCertificate(
+                packageNameFoo, sha256CertFoo, PackageManager.CERT_INPUT_SHA256))
+                .thenReturn(true);
         assertThat(
                         mVisibilityStore.isSchemaSearchableByCaller(
                                 /*packageName=*/ "",
@@ -365,4 +409,14 @@ public class VisibilityStoreTest {
                                 uidFoo))
                 .isTrue();
     }
+
+    @NonNull
+    private PackageManager getMockPackageManager(@NonNull UserHandle user) {
+        PackageManager pm = mMockPackageManagers.get(user);
+        if (pm == null) {
+            pm = Mockito.mock(PackageManager.class);
+            mMockPackageManagers.put(user, pm);
+        }
+        return pm;
+    }
 }

services/tests/servicestests/src/com/android/server/appsearch/external/localstorage/AppSearchImplTest.java

@@ -86,7 +86,6 @@ public class AppSearchImplTest {
                 AppSearchImpl.create(
                         mTemporaryFolder.newFolder(),
                         context,
-                        VisibilityStore.NO_OP_USER_ID,
                         /*logger=*/ null);
     }
 
@@ -494,9 +493,7 @@ public class AppSearchImplTest {
         // Setup the index
         Context context = ApplicationProvider.getApplicationContext();
         File appsearchDir = mTemporaryFolder.newFolder();
-        AppSearchImpl appSearchImpl =
+        AppSearchImpl appSearchImpl = AppSearchImpl.create(appsearchDir, context, /*logger=*/ null);
-                AppSearchImpl.create(
-                        appsearchDir, context, VisibilityStore.NO_OP_USER_ID, /*logger=*/ null);
 
         // Insert schema
         List<AppSearchSchema> schemas =
@@ -557,9 +554,7 @@ public class AppSearchImplTest {
 
         // Initialize AppSearchImpl. This should cause a reset.
         appSearchImpl.close();
-        appSearchImpl =
+        appSearchImpl = AppSearchImpl.create(appsearchDir, context, testLogger);
-                AppSearchImpl.create(
-                        appsearchDir, context, VisibilityStore.NO_OP_USER_ID, testLogger);
 
         // Check recovery state
         InitializeStats initStats = testLogger.mInitializeStats;
@@ -1688,7 +1683,6 @@ public class AppSearchImplTest {
                 AppSearchImpl.create(
                         mTemporaryFolder.newFolder(),
                         context,
-                        VisibilityStore.NO_OP_USER_ID,
                         /*logger=*/ null);
 
         // Initial check that we could do something at first.
@@ -1836,9 +1830,7 @@ public class AppSearchImplTest {
         // Setup the index
         Context context = ApplicationProvider.getApplicationContext();
         File appsearchDir = mTemporaryFolder.newFolder();
-        AppSearchImpl appSearchImpl =
+        AppSearchImpl appSearchImpl = AppSearchImpl.create(appsearchDir, context, /*logger=*/ null);
-                AppSearchImpl.create(
-                        appsearchDir, context, VisibilityStore.NO_OP_USER_ID, /*logger=*/ null);
 
         List<AppSearchSchema> schemas =
                 Collections.singletonList(new AppSearchSchema.Builder("type").build());
@@ -1864,8 +1856,7 @@ public class AppSearchImplTest {
 
         // That document should be visible even from another instance.
         AppSearchImpl appSearchImpl2 =
-                AppSearchImpl.create(
+                AppSearchImpl.create(appsearchDir, context, /*logger=*/ null);
-                        appsearchDir, context, VisibilityStore.NO_OP_USER_ID, /*logger=*/ null);
         getResult =
                 appSearchImpl2.getDocument(
                         "package", "database", "namespace1", "id1", Collections.emptyMap());
@@ -1877,9 +1868,7 @@ public class AppSearchImplTest {
         // Setup the index
         Context context = ApplicationProvider.getApplicationContext();
         File appsearchDir = mTemporaryFolder.newFolder();
-        AppSearchImpl appSearchImpl =
+        AppSearchImpl appSearchImpl = AppSearchImpl.create(appsearchDir, context, /*logger=*/ null);
-                AppSearchImpl.create(
-                        appsearchDir, context, VisibilityStore.NO_OP_USER_ID, /*logger=*/ null);
 
         List<AppSearchSchema> schemas =
                 Collections.singletonList(new AppSearchSchema.Builder("type").build());
@@ -1929,8 +1918,7 @@ public class AppSearchImplTest {
 
         // Only the second document should be retrievable from another instance.
         AppSearchImpl appSearchImpl2 =
-                AppSearchImpl.create(
+                AppSearchImpl.create(appsearchDir, context, /*logger=*/ null);
-                        appsearchDir, context, VisibilityStore.NO_OP_USER_ID, /*logger=*/ null);
         expectThrows(
                 AppSearchException.class,
                 () ->
@@ -1951,9 +1939,7 @@ public class AppSearchImplTest {
         // Setup the index
         Context context = ApplicationProvider.getApplicationContext();
         File appsearchDir = mTemporaryFolder.newFolder();
-        AppSearchImpl appSearchImpl =
+        AppSearchImpl appSearchImpl = AppSearchImpl.create(appsearchDir, context, /*logger=*/ null);
-                AppSearchImpl.create(
-                        appsearchDir, context, VisibilityStore.NO_OP_USER_ID, /*logger=*/ null);
 
         List<AppSearchSchema> schemas =
                 Collections.singletonList(new AppSearchSchema.Builder("type").build());
@@ -2011,8 +1997,7 @@ public class AppSearchImplTest {
 
         // Only the second document should be retrievable from another instance.
         AppSearchImpl appSearchImpl2 =
-                AppSearchImpl.create(
+                AppSearchImpl.create(appsearchDir, context, /*logger=*/ null);
-                        appsearchDir, context, VisibilityStore.NO_OP_USER_ID, /*logger=*/ null);
         expectThrows(
                 AppSearchException.class,
                 () ->

services/tests/servicestests/src/com/android/server/appsearch/external/localstorage/AppSearchLoggerTest.java

@@ -40,7 +40,6 @@ import com.android.server.appsearch.proto.PutDocumentStatsProto;
 import com.android.server.appsearch.proto.QueryStatsProto;
 import com.android.server.appsearch.proto.ScoringSpecProto;
 import com.android.server.appsearch.proto.TermMatchType;
-import com.android.server.appsearch.visibilitystore.VisibilityStore;
 
 import org.junit.Before;
 import org.junit.Rule;
@@ -64,7 +63,6 @@ public class AppSearchLoggerTest {
                 AppSearchImpl.create(
                         mTemporaryFolder.newFolder(),
                         context,
-                        VisibilityStore.NO_OP_USER_ID,
                         /*logger=*/ null);
         mLogger = new TestLogger();
     }
@@ -292,7 +290,6 @@ public class AppSearchLoggerTest {
                 AppSearchImpl.create(
                         mTemporaryFolder.newFolder(),
                         context,
-                        VisibilityStore.NO_OP_USER_ID,
                         mLogger);
 
         InitializeStats iStats = mLogger.mInitializeStats;

services/tests/servicestests/src/com/android/server/appsearch/stats/PlatformLoggerTest.java

@@ -22,6 +22,7 @@ import static org.mockito.Mockito.anyInt;
 import static org.mockito.Mockito.eq;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 import android.annotation.NonNull;
 import android.content.Context;
@@ -29,46 +30,46 @@ import android.content.ContextWrapper;
 import android.content.pm.PackageManager;
 import android.os.SystemClock;
 import android.os.UserHandle;
+import android.util.ArrayMap;
 import android.util.SparseIntArray;
 
 import androidx.test.core.app.ApplicationProvider;
 
-import com.android.server.appsearch.external.localstorage.MockPackageManager;
 import com.android.server.appsearch.external.localstorage.stats.CallStats;
 
 import org.junit.Before;
 import org.junit.Test;
+import org.mockito.Mockito;
 
 import java.io.UnsupportedEncodingException;
 import java.math.BigInteger;
+import java.nio.charset.StandardCharsets;
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
+import java.util.Map;
 
 public class PlatformLoggerTest {
     private static final int TEST_MIN_TIME_INTERVAL_BETWEEN_SAMPLES_MILLIS = 100;
     private static final int TEST_DEFAULT_SAMPLING_INTERVAL = 10;
     private static final String TEST_PACKAGE_NAME = "packageName";
-    private MockPackageManager mMockPackageManager = new MockPackageManager();
+
+    private final Map<UserHandle, PackageManager> mMockPackageManagers = new ArrayMap<>();
     private Context mContext;
 
     @Before
     public void setUp() throws Exception {
         Context context = ApplicationProvider.getApplicationContext();
-        mContext =
+        mContext = new ContextWrapper(context) {
-                new ContextWrapper(context) {
+            @Override
+            public Context createContextAsUser(UserHandle user, int flags) {
+                return new ContextWrapper(super.createContextAsUser(user, flags)) {
                     @Override
                     public PackageManager getPackageManager() {
-                        return mMockPackageManager.getMockPackageManager();
+                        return getMockPackageManager(user);
                     }
                 };
-    }
+            }
-
+        };
-    static int calculateHashCodeMd5withBigInteger(@NonNull String str) throws
-            NoSuchAlgorithmException, UnsupportedEncodingException {
-        MessageDigest md = MessageDigest.getInstance("MD5");
-        md.update(str.getBytes(/*charsetName=*/ "UTF-8"));
-        byte[] digest = md.digest();
-        return new BigInteger(digest).intValue();
     }
 
     @Test
@@ -294,39 +295,51 @@ public class PlatformLoggerTest {
                         TEST_MIN_TIME_INTERVAL_BETWEEN_SAMPLES_MILLIS,
                         TEST_DEFAULT_SAMPLING_INTERVAL,
                         /*samplingIntervals=*/ new SparseIntArray()));
-        mMockPackageManager.mockGetPackageUidAsUser(testPackageName, mContext.getUserId(), testUid);
+        PackageManager mockPackageManager = getMockPackageManager(mContext.getUser());
+        when(mockPackageManager.getPackageUid(testPackageName, /*flags=*/0)).thenReturn(testUid);
 
-        //
         // First time, no cache
-        //
         PlatformLogger.ExtraStats extraStats = logger.createExtraStatsLocked(testPackageName,
                 CallStats.CALL_TYPE_PUT_DOCUMENT);
-
+        verify(mockPackageManager, times(1))
-        verify(mMockPackageManager.getMockPackageManager(), times(1)).getPackageUidAsUser(
+                .getPackageUid(eq(testPackageName), /*flags=*/ anyInt());
-                eq(testPackageName), /*userId=*/ anyInt());
         assertThat(extraStats.mPackageUid).isEqualTo(testUid);
 
-        //
         // Second time, we have cache
-        //
         extraStats = logger.createExtraStatsLocked(testPackageName,
                 CallStats.CALL_TYPE_PUT_DOCUMENT);
 
         // Count is still one since we will use the cache
-        verify(mMockPackageManager.getMockPackageManager(), times(1)).getPackageUidAsUser(
+        verify(mockPackageManager, times(1))
-                eq(testPackageName), /*userId=*/ anyInt());
+                .getPackageUid(eq(testPackageName), /*flags=*/ anyInt());
         assertThat(extraStats.mPackageUid).isEqualTo(testUid);
 
-        //
         // Remove the cache and try again
-        //
         assertThat(logger.removeCachedUidForPackage(testPackageName)).isEqualTo(testUid);
         extraStats = logger.createExtraStatsLocked(testPackageName,
                 CallStats.CALL_TYPE_PUT_DOCUMENT);
 
         // count increased by 1 since cache is cleared
-        verify(mMockPackageManager.getMockPackageManager(), times(2)).getPackageUidAsUser(
+        verify(mockPackageManager, times(2))
-                eq(testPackageName), /*userId=*/ anyInt());
+                .getPackageUid(eq(testPackageName), /*flags=*/ anyInt());
         assertThat(extraStats.mPackageUid).isEqualTo(testUid);
     }
+
+    private static int calculateHashCodeMd5withBigInteger(@NonNull String str)
+            throws NoSuchAlgorithmException {
+        MessageDigest md = MessageDigest.getInstance("MD5");
+        md.update(str.getBytes(StandardCharsets.UTF_8));
+        byte[] digest = md.digest();
+        return new BigInteger(digest).intValue();
+    }
+
+    @NonNull
+    private PackageManager getMockPackageManager(@NonNull UserHandle user) {
+        PackageManager pm = mMockPackageManagers.get(user);
+        if (pm == null) {
+            pm = Mockito.mock(PackageManager.class);
+            mMockPackageManagers.put(user, pm);
+        }
+        return pm;
+    }
 }