Make device wide DO features available if all users affiliated

Currently, those features are available on single user devices only
(since they collect privacy sensitive data device wide). Now making
them available as long as all users are affiliated.

It'll take a certain amount of time between user creation and the DPC
of that new user setting the appropriate affiliation ids. The DO won't
be able to access the logs during that time (and won't get any "logs
ready" callback). Once the affiliation ids are set, if they match,
logs become available again - this includes logs collected while the
user was being setup. Some logs might be lost though if the amount of
data exceeds the internal limit.

Test: runtest -c com.android.server.devicepolicy.DevicePolicyManagerTest frameworks-services

Test: cts-tradefed run cts -a armeabi-v7a --module CtsDevicePolicyManagerTestCases --test com.android.cts.devicepolicy.DeviceOwnerTest

Bug: 32326223

Change-Id: Idfe881dd6497d3ad2bead10addfd37b98b8a6e2b
diff --git a/core/java/android/app/admin/DeviceAdminReceiver.java b/core/java/android/app/admin/DeviceAdminReceiver.java
index aae80ed..c82bb57 100644
--- a/core/java/android/app/admin/DeviceAdminReceiver.java
+++ b/core/java/android/app/admin/DeviceAdminReceiver.java
@@ -681,6 +681,10 @@
     /**
      * Called when a new batch of security logs can be retrieved.
      *
+     * <p>If a secondary user or profile is created, this callback won't be received until all users
+     * become affiliated again (even if security logging is enabled).
+     * See {@link DevicePolicyManager#setAffiliationIds}
+     *
      * <p>This callback is only applicable to device owners.
      *
      * @param context The running context as per {@link #onReceive}.
@@ -695,13 +699,18 @@
      * ever be called when network logging is enabled. The logs can only be retrieved while network
      * logging is enabled.
      *
+     * <p>If a secondary user or profile is created, this callback won't be received until all users
+     * become affiliated again (even if network logging is enabled). It will also no longer be
+     * possible to retrieve the network logs batch with the most recent {@code batchToken} provided
+     * by this callback. See {@link DevicePolicyManager#setAffiliationIds}.
+     *
      * <p>This callback is only applicable to device owners.
      *
      * @param context The running context as per {@link #onReceive}.
      * @param intent The received intent as per {@link #onReceive}.
      * @param batchToken The token representing the current batch of network logs.
      * @param networkLogsCount The total count of events in the current batch of network logs.
-     * @see DevicePolicyManager#retrieveNetworkLogs(ComponentName)
+     * @see DevicePolicyManager#retrieveNetworkLogs
      */
     public void onNetworkLogsAvailable(Context context, Intent intent, long batchToken,
             int networkLogsCount) {
diff --git a/core/java/android/app/admin/DevicePolicyManager.java b/core/java/android/app/admin/DevicePolicyManager.java
index 074326f..c95e011 100644
--- a/core/java/android/app/admin/DevicePolicyManager.java
+++ b/core/java/android/app/admin/DevicePolicyManager.java
@@ -3651,15 +3651,16 @@
     /**
      * Called by a device owner to request a bugreport.
      * <p>
-     * There must be only one user on the device, managed by the device owner. Otherwise a
-     * {@link SecurityException} will be thrown.
+     * If the device contains secondary users or profiles, they must be affiliated with the device
+     * owner user. Otherwise a {@link SecurityException} will be thrown. See
+     * {@link #setAffiliationIds}.
      *
      * @param admin Which {@link DeviceAdminReceiver} this request is associated with.
      * @return {@code true} if the bugreport collection started successfully, or {@code false} if it
      *         wasn't triggered because a previous bugreport operation is still active (either the
      *         bugreport is still running or waiting for the user to share or decline)
-     * @throws SecurityException if {@code admin} is not a device owner, or if there are users other
-     *             than the one managed by the device owner.
+     * @throws SecurityException if {@code admin} is not a device owner, or there is at least one
+     *         profile or secondary user that is not affiliated with the device owner user.
      */
     public boolean requestBugreport(@NonNull ComponentName admin) {
         throwIfParentInstance("requestBugreport");
@@ -6631,14 +6632,16 @@
     }
 
     /**
-     * Called by device owner to control the security logging feature. Logging can only be
-     * enabled on single user devices where the sole user is managed by the device owner.
+     * Called by device owner to control the security logging feature.
      *
      * <p> Security logs contain various information intended for security auditing purposes.
      * See {@link SecurityEvent} for details.
      *
-     * <p>There must be only one user on the device, managed by the device owner.
-     * Otherwise a {@link SecurityException} will be thrown.
+     * <p><strong>Note:</strong> The device owner won't be able to retrieve security logs if there
+     * are unaffiliated secondary users or profiles on the device, regardless of whether the
+     * feature is enabled. Logs will be discarded if the internal buffer fills up while waiting for
+     * all users to become affiliated. Therefore it's recommended that affiliation ids are set for
+     * new users as soon as possible after provisioning via {@link #setAffiliationIds}.
      *
      * @param admin Which device owner this request is associated with.
      * @param enabled whether security logging should be enabled or not.
@@ -6680,13 +6683,16 @@
      * <p> Access to the logs is rate limited and it will only return new logs after the device
      * owner has been notified via {@link DeviceAdminReceiver#onSecurityLogsAvailable}.
      *
-     * <p>There must be only one user on the device, managed by the device owner.
-     * Otherwise a {@link SecurityException} will be thrown.
+     * <p>If there is any other user or profile on the device, it must be affiliated with the
+     * device owner. Otherwise a {@link SecurityException} will be thrown. See
+     * {@link #setAffiliationIds}
      *
      * @param admin Which device owner this request is associated with.
      * @return the new batch of security logs which is a list of {@link SecurityEvent},
      * or {@code null} if rate limitation is exceeded or if logging is currently disabled.
-     * @throws SecurityException if {@code admin} is not a device owner.
+     * @throws SecurityException if {@code admin} is not a device owner, or there is at least one
+     * profile or secondary user that is not affiliated with the device owner user.
+     * @see DeviceAdminReceiver#onSecurityLogsAvailable
      */
     public @Nullable List<SecurityEvent> retrieveSecurityLogs(@NonNull ComponentName admin) {
         throwIfParentInstance("retrieveSecurityLogs");
@@ -6726,14 +6732,17 @@
      * will result in {@code null} being returned. The device logs are retrieved from a RAM region
      * which is not guaranteed to be corruption-free during power cycles, as a result be cautious
      * about data corruption when parsing. </strong>
-     * <p>
-     * There must be only one user on the device, managed by the device owner. Otherwise a
-     * {@link SecurityException} will be thrown.
+     *
+     * <p>If there is any other user or profile on the device, it must be affiliated with the
+     * device owner. Otherwise a {@link SecurityException} will be thrown. See
+     * {@link #setAffiliationIds}
      *
      * @param admin Which device owner this request is associated with.
      * @return Device logs from before the latest reboot of the system, or {@code null} if this API
      *         is not supported on the device.
-     * @throws SecurityException if {@code admin} is not a device owner.
+     * @throws SecurityException if {@code admin} is not a device owner, or there is at least one
+     * profile or secondary user that is not affiliated with the device owner user.
+     * @see #retrieveSecurityLogs
      */
     public @Nullable List<SecurityEvent> retrievePreRebootSecurityLogs(
             @NonNull ComponentName admin) {
@@ -6939,6 +6948,12 @@
      * Indicates the entity that controls the device or profile owner. Two users/profiles are
      * affiliated if the set of ids set by their device or profile owners intersect.
      *
+     * <p><strong>Note:</strong> Features that depend on user affiliation (such as security logging
+     * or {@link #bindDeviceAdminServiceAsUser}) won't be available when a secondary user or profile
+     * is created, until it becomes affiliated. Therefore it is recommended that the appropriate
+     * affiliation ids are set by its profile owner as soon as possible after the user/profile is
+     * created.
+     *
      * @param admin Which profile or device owner this request is associated with.
      * @param ids A list of opaque non-empty affiliation ids. Duplicate elements will be ignored.
      *
@@ -7138,15 +7153,19 @@
     }
 
     /**
-     * Called by a device owner to control the network logging feature. Logging can only be
-     * enabled on single user devices where the sole user is managed by the device owner. If a new
-     * user is added on the device, logging is disabled.
+     * Called by a device owner to control the network logging feature.
      *
      * <p> Network logs contain DNS lookup and connect() library call events.
      *
+     * <p><strong>Note:</strong> The device owner won't be able to retrieve network logs if there
+     * are unaffiliated secondary users or profiles on the device, regardless of whether the
+     * feature is enabled. Logs will be discarded if the internal buffer fills up while waiting for
+     * all users to become affiliated. Therefore it's recommended that affiliation ids are set for
+     * new users as soon as possible after provisioning via {@link #setAffiliationIds}.
+     *
      * @param admin Which {@link DeviceAdminReceiver} this request is associated with.
      * @param enabled whether network logging should be enabled or not.
-     * @throws {@link SecurityException} if {@code admin} is not a device owner.
+     * @throws SecurityException if {@code admin} is not a device owner.
      * @see #retrieveNetworkLogs
      */
     public void setNetworkLoggingEnabled(@NonNull ComponentName admin, boolean enabled) {
@@ -7164,7 +7183,7 @@
      * @param admin Which {@link DeviceAdminReceiver} this request is associated with. Can only
      * be {@code null} if the caller has MANAGE_USERS permission.
      * @return {@code true} if network logging is enabled by device owner, {@code false} otherwise.
-     * @throws {@link SecurityException} if {@code admin} is not a device owner and caller has
+     * @throws SecurityException if {@code admin} is not a device owner and caller has
      * no MANAGE_USERS permission
      */
     public boolean isNetworkLoggingEnabled(@Nullable ComponentName admin) {
@@ -7190,12 +7209,19 @@
      * after the device device owner has been notified via
      * {@link DeviceAdminReceiver#onNetworkLogsAvailable}.
      *
+     * <p>If a secondary user or profile is created, calling this method will throw a
+     * {@link SecurityException} until all users become affiliated again. It will also no longer be
+     * possible to retrieve the network logs batch with the most recent batchToken provided
+     * by {@link DeviceAdminReceiver#onNetworkLogsAvailable}. See
+     * {@link DevicePolicyManager#setAffiliationIds}.
+     *
      * @param admin Which {@link DeviceAdminReceiver} this request is associated with.
      * @param batchToken A token of the batch to retrieve
      * @return A new batch of network logs which is a list of {@link NetworkEvent}. Returns
      *        {@code null} if the batch represented by batchToken is no longer available or if
      *        logging is disabled.
-     * @throws {@link SecurityException} if {@code admin} is not a device owner.
+     * @throws SecurityException if {@code admin} is not a device owner, or there is at least one
+     * profile or secondary user that is not affiliated with the device owner user.
      * @see DeviceAdminReceiver#onNetworkLogsAvailable
      */
     public @Nullable List<NetworkEvent> retrieveNetworkLogs(@NonNull ComponentName admin,
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
index cfd2bed..6e4c425 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
@@ -17,7 +17,6 @@
 package com.android.server.devicepolicy;
 
 import static android.Manifest.permission.MANAGE_CA_CERTIFICATES;
-import static android.app.admin.DevicePolicyManager.ACTION_SHOW_DEVICE_MONITORING_DIALOG;
 import static android.app.admin.DevicePolicyManager.CODE_ACCOUNTS_NOT_EMPTY;
 import static android.app.admin.DevicePolicyManager.CODE_ADD_MANAGED_PROFILE_DISALLOWED;
 import static android.app.admin.DevicePolicyManager.CODE_CANNOT_ADD_MANAGED_PROFILE;
@@ -49,7 +48,6 @@
 import android.accessibilityservice.AccessibilityServiceInfo;
 import android.accounts.Account;
 import android.accounts.AccountManager;
-import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.UserIdInt;
@@ -180,8 +178,6 @@
 import java.io.FileOutputStream;
 import java.io.IOException;
 import java.io.PrintWriter;
-import java.lang.annotation.Retention;
-import java.lang.annotation.RetentionPolicy;
 import java.nio.charset.StandardCharsets;
 import java.security.cert.CertificateException;
 import java.security.cert.CertificateFactory;
@@ -337,7 +333,8 @@
     private static final long MINIMUM_STRONG_AUTH_TIMEOUT_MS = 1 * 60 * 60 * 1000; // 1h
 
     /**
-     * Strings logged with {@link #PROVISIONING_ENTRY_POINT_ADB}.
+     * Strings logged with {@link
+     * com.android.internal.logging.nano.MetricsProto.MetricsEvent#PROVISIONING_ENTRY_POINT_ADB}.
      */
     private static final String LOG_TAG_PROFILE_OWNER = "profile-owner";
     private static final String LOG_TAG_DEVICE_OWNER = "device-owner";
@@ -552,11 +549,25 @@
             }
             if (Intent.ACTION_USER_ADDED.equals(action)) {
                 sendUserAddedOrRemovedCommand(DeviceAdminReceiver.ACTION_USER_ADDED, userHandle);
-                disableDeviceOwnerManagedSingleUserFeaturesIfNeeded();
+                synchronized (DevicePolicyManagerService.this) {
+                    // It might take a while for the user to become affiliated. Make security
+                    // and network logging unavailable in the meantime.
+                    maybePauseDeviceWideLoggingLocked();
+                }
             } else if (Intent.ACTION_USER_REMOVED.equals(action)) {
                 sendUserAddedOrRemovedCommand(DeviceAdminReceiver.ACTION_USER_REMOVED, userHandle);
-                disableDeviceOwnerManagedSingleUserFeaturesIfNeeded();
-                removeUserData(userHandle);
+                synchronized (DevicePolicyManagerService.this) {
+                    // Check whether the user is affiliated, *before* removing its data.
+                    boolean isRemovedUserAffiliated = isUserAffiliatedWithDeviceLocked(userHandle);
+                    removeUserData(userHandle);
+                    if (!isRemovedUserAffiliated) {
+                        // We discard the logs when unaffiliated users are deleted (so that the
+                        // device owner cannot retrieve data about that user after it's gone).
+                        discardDeviceWideLogsLocked();
+                        // Resume logging if all remaining users are affiliated.
+                        maybeResumeDeviceWideLoggingLocked();
+                    }
+                }
             } else if (Intent.ACTION_USER_STARTED.equals(action)) {
                 synchronized (DevicePolicyManagerService.this) {
                     // Reset the policy data
@@ -1858,9 +1869,10 @@
             if (mOwners.hasDeviceOwner()) {
                 mInjector.systemPropertiesSet(PROPERTY_DEVICE_OWNER_PRESENT, "true");
                 Slog.i(LOG_TAG, "Set ro.device_owner property to true");
-                disableDeviceOwnerManagedSingleUserFeaturesIfNeeded();
+
                 if (mInjector.securityLogGetLoggingEnabledProperty()) {
                     mSecurityLogMonitor.start();
+                    maybePauseDeviceWideLoggingLocked();
                 }
             } else {
                 mInjector.systemPropertiesSet(PROPERTY_DEVICE_OWNER_PRESENT, "false");
@@ -3124,7 +3136,7 @@
     }
 
     // It's temporary solution to clear DISALLOW_ADD_USER after CTS
-    // TODO: b/31952368 when the restriction is moved from system to the device owner,
+    // STOPSHIP(b/31952368) when the restriction is moved from system to the device owner,
     // it can be removed.
     private void clearDeviceOwnerUserRestrictionLocked(UserHandle userHandle) {
         if (mUserManager.hasUserRestriction(UserManager.DISALLOW_ADD_USER, userHandle)) {
@@ -5619,34 +5631,12 @@
         }
     }
 
-    private boolean isDeviceOwnerManagedSingleUserDevice() {
-        synchronized (this) {
-            if (!mOwners.hasDeviceOwner()) {
-                return false;
-            }
-        }
-        final long callingIdentity = mInjector.binderClearCallingIdentity();
-        try {
-            if (mInjector.userManagerIsSplitSystemUser()) {
-                // In split system user mode, only allow the case where the device owner is managing
-                // the only non-system user of the device
-                return (mUserManager.getUserCount() == 2
-                        && mOwners.getDeviceOwnerUserId() != UserHandle.USER_SYSTEM);
-            } else  {
-                return mUserManager.getUserCount() == 1;
-            }
-        } finally {
-            mInjector.binderRestoreCallingIdentity(callingIdentity);
-        }
-    }
-
-    private void ensureDeviceOwnerManagingSingleUser(ComponentName who) throws SecurityException {
+    private void ensureDeviceOwnerAndAllUsersAffiliated(ComponentName who) throws SecurityException {
         synchronized (this) {
             getActiveAdminForCallerLocked(who, DeviceAdminInfo.USES_POLICY_DEVICE_OWNER);
-        }
-        if (!isDeviceOwnerManagedSingleUserDevice()) {
-            throw new SecurityException(
-                    "There should only be one user, managed by Device Owner");
+            if (!areAllUsersAffiliatedWithDeviceLocked()) {
+                throw new SecurityException("Not all users are affiliated.");
+            }
         }
     }
 
@@ -5656,7 +5646,11 @@
             return false;
         }
         Preconditions.checkNotNull(who, "ComponentName is null");
-        ensureDeviceOwnerManagingSingleUser(who);
+
+        // TODO: If an unaffiliated user is removed, the admin will be able to request a bugreport
+        // which could still contain data related to that user. Should we disallow that, e.g. until
+        // next boot? Might not be needed given that this still requires user consent.
+        ensureDeviceOwnerAndAllUsersAffiliated(who);
 
         if (mRemoteBugreportServiceIsActive.get()
                 || (getDeviceOwnerRemoteBugreportUri() != null)) {
@@ -5681,7 +5675,8 @@
             mRemoteBugreportServiceIsActive.set(true);
             mRemoteBugreportSharingAccepted.set(false);
             registerRemoteBugreportReceivers();
-            mInjector.getNotificationManager().notifyAsUser(LOG_TAG, RemoteBugreportUtils.NOTIFICATION_ID,
+            mInjector.getNotificationManager().notifyAsUser(LOG_TAG,
+                    RemoteBugreportUtils.NOTIFICATION_ID,
                     RemoteBugreportUtils.buildNotification(mContext,
                             DevicePolicyManager.NOTIFICATION_BUGREPORT_STARTED), UserHandle.ALL);
             mHandler.postDelayed(mRemoteBugreportTimeoutRunnable,
@@ -6228,6 +6223,7 @@
             admin.userRestrictions = null;
             admin.defaultEnabledRestrictionsAlreadySet.clear();
             admin.forceEphemeralUsers = false;
+            admin.isNetworkLoggingEnabled = false;
             mUserManagerInternal.setForceEphemeralUsers(admin.forceEphemeralUsers);
             final DevicePolicyData policyData = getUserData(UserHandle.USER_SYSTEM);
             policyData.mLastSecurityLogRetrievalTime = -1;
@@ -6240,7 +6236,11 @@
         mOwners.clearDeviceOwner();
         mOwners.writeDeviceOwner();
         updateDeviceOwnerLocked();
-        disableDeviceOwnerManagedSingleUserFeaturesIfNeeded();
+
+        mInjector.securityLogSetLoggingEnabledProperty(false);
+        mSecurityLogMonitor.stop();
+        setNetworkLoggingActiveInternal(false);
+
         try {
             if (mInjector.getIBackupManager() != null) {
                 // Reactivate backup service.
@@ -8227,7 +8227,7 @@
         synchronized (this) {
             getActiveAdminForCallerLocked(who, DeviceAdminInfo.USES_POLICY_PROFILE_OWNER);
             final int userHandle = mInjector.userHandleGetCallingUserId();
-            if (isUserAffiliatedWithDevice(userHandle)) {
+            if (isUserAffiliatedWithDeviceLocked(userHandle)) {
                 setLockTaskPackagesLocked(userHandle, new ArrayList<>(Arrays.asList(packages)));
             } else {
                 throw new SecurityException("Admin " + who +
@@ -9408,6 +9408,12 @@
                 getUserData(UserHandle.USER_SYSTEM).mAffiliationIds = affiliationIds;
                 saveSettingsLocked(UserHandle.USER_SYSTEM);
             }
+
+            // Affiliation status for any user, not just the calling user, might have changed.
+            // The device owner user will still be affiliated after changing its affiliation ids,
+            // but as a result of that other users might become affiliated or un-affiliated.
+            maybePauseDeviceWideLoggingLocked();
+            maybeResumeDeviceWideLoggingLocked();
         }
     }
 
@@ -9427,84 +9433,78 @@
 
     @Override
     public boolean isAffiliatedUser() {
-        return isUserAffiliatedWithDevice(mInjector.userHandleGetCallingUserId());
+        if (!mHasFeature) {
+            return false;
+        }
+
+        synchronized (this) {
+            return isUserAffiliatedWithDeviceLocked(mInjector.userHandleGetCallingUserId());
+        }
     }
 
-    private boolean isUserAffiliatedWithDevice(int userId) {
-        synchronized (this) {
-            if (!mOwners.hasDeviceOwner()) {
-                return false;
-            }
-            if (userId == mOwners.getDeviceOwnerUserId()) {
-                // The user that the DO is installed on is always affiliated with the device.
+    private boolean isUserAffiliatedWithDeviceLocked(int userId) {
+        if (!mOwners.hasDeviceOwner()) {
+            return false;
+        }
+        if (userId == mOwners.getDeviceOwnerUserId()) {
+            // The user that the DO is installed on is always affiliated with the device.
+            return true;
+        }
+        if (userId == UserHandle.USER_SYSTEM) {
+            // The system user is always affiliated in a DO device, even if the DO is set on a
+            // different user. This could be the case if the DO is set in the primary user
+            // of a split user device.
+            return true;
+        }
+        final ComponentName profileOwner = getProfileOwner(userId);
+        if (profileOwner == null) {
+            return false;
+        }
+        final Set<String> userAffiliationIds = getUserData(userId).mAffiliationIds;
+        final Set<String> deviceAffiliationIds =
+                getUserData(UserHandle.USER_SYSTEM).mAffiliationIds;
+        for (String id : userAffiliationIds) {
+            if (deviceAffiliationIds.contains(id)) {
                 return true;
             }
-            if (userId == UserHandle.USER_SYSTEM) {
-                // The system user is always affiliated in a DO device, even if the DO is set on a
-                // different user. This could be the case if the DO is set in the primary user
-                // of a split user device.
-                return true;
-            }
-            final ComponentName profileOwner = getProfileOwner(userId);
-            if (profileOwner == null) {
-                return false;
-            }
-            final Set<String> userAffiliationIds = getUserData(userId).mAffiliationIds;
-            final Set<String> deviceAffiliationIds =
-                    getUserData(UserHandle.USER_SYSTEM).mAffiliationIds;
-            for (String id : userAffiliationIds) {
-                if (deviceAffiliationIds.contains(id)) {
-                    return true;
-                }
-            }
         }
         return false;
     }
 
-    private synchronized void disableDeviceOwnerManagedSingleUserFeaturesIfNeeded() {
-        final boolean isSingleUserManagedDevice = isDeviceOwnerManagedSingleUserDevice();
-
-        // disable security logging if needed
-        if (!isSingleUserManagedDevice) {
-            mInjector.securityLogSetLoggingEnabledProperty(false);
-            Slog.w(LOG_TAG, "Security logging turned off as it's no longer a single user managed"
-                    + " device.");
-        }
-
-        // disable backup service if needed
-        // note: when clearing DO, the backup service shouldn't be disabled if it was enabled by
-        // the device owner
-        if (mOwners.hasDeviceOwner() && !isSingleUserManagedDevice) {
-            setBackupServiceEnabledInternal(false);
-            Slog.w(LOG_TAG, "Backup is off as it's a managed device that has more that one user.");
-        }
-
-        // disable network logging if needed
-        if (!isSingleUserManagedDevice) {
-            setNetworkLoggingActiveInternal(false);
-            Slog.w(LOG_TAG, "Network logging turned off as it's no longer a single user managed"
-                    + " device.");
-            // if there still is a device owner, disable logging policy, otherwise the admin
-            // has been nuked
-            if (mOwners.hasDeviceOwner()) {
-                getDeviceOwnerAdminLocked().isNetworkLoggingEnabled = false;
-                saveSettingsLocked(mOwners.getDeviceOwnerUserId());
+    private boolean areAllUsersAffiliatedWithDeviceLocked() {
+        final long ident = mInjector.binderClearCallingIdentity();
+        try {
+            final List<UserInfo> userInfos = mUserManager.getUsers();
+            for (int i = 0; i < userInfos.size(); i++) {
+                int userId = userInfos.get(i).id;
+                if (!isUserAffiliatedWithDeviceLocked(userId)) {
+                    Slog.d(LOG_TAG, "User id " + userId + " not affiliated.");
+                    return false;
+                }
             }
+        } finally {
+            mInjector.binderRestoreCallingIdentity(ident);
         }
+
+        return true;
     }
 
     @Override
     public void setSecurityLoggingEnabled(ComponentName admin, boolean enabled) {
+        if (!mHasFeature) {
+            return;
+        }
         Preconditions.checkNotNull(admin);
-        ensureDeviceOwnerManagingSingleUser(admin);
 
         synchronized (this) {
+            getActiveAdminForCallerLocked(admin, DeviceAdminInfo.USES_POLICY_DEVICE_OWNER);
             if (enabled == mInjector.securityLogGetLoggingEnabledProperty()) {
                 return;
             }
             mInjector.securityLogSetLoggingEnabledProperty(enabled);
             if (enabled) {
                 mSecurityLogMonitor.start();
+                maybePauseDeviceWideLoggingLocked();
             } else {
                 mSecurityLogMonitor.stop();
             }
@@ -9513,6 +9513,10 @@
 
     @Override
     public boolean isSecurityLoggingEnabled(ComponentName admin) {
+        if (!mHasFeature) {
+            return false;
+        }
+
         Preconditions.checkNotNull(admin);
         synchronized (this) {
             getActiveAdminForCallerLocked(admin, DeviceAdminInfo.USES_POLICY_DEVICE_OWNER);
@@ -9531,10 +9535,15 @@
 
     @Override
     public ParceledListSlice<SecurityEvent> retrievePreRebootSecurityLogs(ComponentName admin) {
-        Preconditions.checkNotNull(admin);
-        ensureDeviceOwnerManagingSingleUser(admin);
+        if (!mHasFeature) {
+            return null;
+        }
 
-        if (!mContext.getResources().getBoolean(R.bool.config_supportPreRebootSecurityLogs)) {
+        Preconditions.checkNotNull(admin);
+        ensureDeviceOwnerAndAllUsersAffiliated(admin);
+
+        if (!mContext.getResources().getBoolean(R.bool.config_supportPreRebootSecurityLogs)
+                || !mInjector.securityLogGetLoggingEnabledProperty()) {
             return null;
         }
 
@@ -9552,8 +9561,16 @@
 
     @Override
     public ParceledListSlice<SecurityEvent> retrieveSecurityLogs(ComponentName admin) {
+        if (!mHasFeature) {
+            return null;
+        }
+
         Preconditions.checkNotNull(admin);
-        ensureDeviceOwnerManagingSingleUser(admin);
+        ensureDeviceOwnerAndAllUsersAffiliated(admin);
+
+        if (!mInjector.securityLogGetLoggingEnabledProperty()) {
+            return null;
+        }
 
         recordSecurityLogRetrievalTime();
 
@@ -9760,18 +9777,21 @@
         }
     }
 
+    // TODO(b/22388012): When backup is available for secondary users and profiles, consider
+    // whether there are any privacy/security implications of enabling the backup service here
+    // if there are other users or profiles unmanaged or managed by a different entity (i.e. not
+    // affiliated).
     @Override
     public void setBackupServiceEnabled(ComponentName admin, boolean enabled) {
-        Preconditions.checkNotNull(admin);
         if (!mHasFeature) {
             return;
         }
-        ensureDeviceOwnerManagingSingleUser(admin);
-        setBackupServiceEnabledInternal(enabled);
-    }
+        Preconditions.checkNotNull(admin);
+        synchronized (this) {
+            getActiveAdminForCallerLocked(admin, DeviceAdminInfo.USES_POLICY_DEVICE_OWNER);
+        }
 
-    private synchronized void setBackupServiceEnabledInternal(boolean enabled) {
-        long ident = mInjector.binderClearCallingIdentity();
+        final long ident = mInjector.binderClearCallingIdentity();
         try {
             IBackupManager ibm = mInjector.getIBackupManager();
             if (ibm != null) {
@@ -9872,7 +9892,7 @@
             final boolean isCallerDeviceOwner = isDeviceOwner(callingOwner);
             final boolean isCallerManagedProfile = isManagedProfile(callingUserId);
             if ((!isCallerDeviceOwner && !isCallerManagedProfile)
-                    || !isUserAffiliatedWithDevice(callingUserId)) {
+                    || !isUserAffiliatedWithDeviceLocked(callingUserId)) {
                 return targetUsers;
             }
 
@@ -9892,7 +9912,7 @@
 
                         // Both must be the same package and be affiliated in order to bind.
                         if (callingOwnerPackage.equals(targetOwnerPackage)
-                               && isUserAffiliatedWithDevice(userId)) {
+                               && isUserAffiliatedWithDeviceLocked(userId)) {
                             targetUsers.add(UserHandle.of(userId));
                         }
                     }
@@ -9990,7 +10010,7 @@
             return;
         }
         Preconditions.checkNotNull(admin);
-        ensureDeviceOwnerManagingSingleUser(admin);
+        getActiveAdminForCallerLocked(admin, DeviceAdminInfo.USES_POLICY_DEVICE_OWNER);
 
         if (enabled == isNetworkLoggingEnabledInternalLocked()) {
             // already in the requested state
@@ -10017,10 +10037,10 @@
                     Slog.wtf(LOG_TAG, "Network logging could not be started due to the logging"
                             + " service not being available yet.");
                 }
+                maybePauseDeviceWideLoggingLocked();
                 sendNetworkLoggingNotificationLocked();
             } else {
                 if (mNetworkLogger != null && !mNetworkLogger.stopNetworkLogging()) {
-                    mNetworkLogger = null;
                     Slog.wtf(LOG_TAG, "Network logging could not be stopped due to the logging"
                             + " service not being available yet.");
                 }
@@ -10032,6 +10052,44 @@
         }
     }
 
+    /** Pauses security and network logging if there are unaffiliated users on the device */
+    private void maybePauseDeviceWideLoggingLocked() {
+        if (!areAllUsersAffiliatedWithDeviceLocked()) {
+            Slog.i(LOG_TAG, "There are unaffiliated users, security and network logging will be "
+                    + "paused if enabled.");
+            mSecurityLogMonitor.pause();
+            if (mNetworkLogger != null) {
+                mNetworkLogger.pause();
+            }
+        }
+    }
+
+    /** Resumes security and network logging (if they are enabled) if all users are affiliated */
+    private void maybeResumeDeviceWideLoggingLocked() {
+        if (areAllUsersAffiliatedWithDeviceLocked()) {
+            final long ident = mInjector.binderClearCallingIdentity();
+            try {
+                mSecurityLogMonitor.resume();
+                if (mNetworkLogger != null) {
+                    mNetworkLogger.resume();
+                }
+            } finally {
+                mInjector.binderRestoreCallingIdentity(ident);
+            }
+        }
+    }
+
+    /** Deletes any security and network logs that might have been collected so far */
+    private void discardDeviceWideLogsLocked() {
+        mSecurityLogMonitor.discardLogs();
+        if (mNetworkLogger != null) {
+            mNetworkLogger.discardLogs();
+        }
+        // TODO: We should discard pre-boot security logs here too, as otherwise those
+        // logs (which might contain data from the user just removed) will be
+        // available after next boot.
+    }
+
     @Override
     public boolean isNetworkLoggingEnabled(ComponentName admin) {
         if (!mHasFeature) {
@@ -10056,32 +10114,27 @@
      * @see NetworkLoggingHandler#MAX_EVENTS_PER_BATCH
      */
     @Override
-    public synchronized List<NetworkEvent> retrieveNetworkLogs(ComponentName admin,
-            long batchToken) {
+    public List<NetworkEvent> retrieveNetworkLogs(ComponentName admin, long batchToken) {
         if (!mHasFeature) {
             return null;
         }
         Preconditions.checkNotNull(admin);
-        ensureDeviceOwnerManagingSingleUser(admin);
+        ensureDeviceOwnerAndAllUsersAffiliated(admin);
 
-        if (mNetworkLogger == null) {
-            return null;
-        }
-
-        if (!isNetworkLoggingEnabledInternalLocked()) {
-            return null;
-        }
-
-        final long currentTime = System.currentTimeMillis();
         synchronized (this) {
+            if (mNetworkLogger == null
+                    || !isNetworkLoggingEnabledInternalLocked()) {
+                return null;
+            }
+
+            final long currentTime = System.currentTimeMillis();
             DevicePolicyData policyData = getUserData(UserHandle.USER_SYSTEM);
             if (currentTime > policyData.mLastNetworkLogsRetrievalTime) {
                 policyData.mLastNetworkLogsRetrievalTime = currentTime;
                 saveSettingsLocked(UserHandle.USER_SYSTEM);
             }
+            return mNetworkLogger.retrieveLogs(batchToken);
         }
-
-        return mNetworkLogger.retrieveLogs(batchToken);
     }
 
     private void sendNetworkLoggingNotificationLocked() {
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/NetworkLogger.java b/services/devicepolicy/java/com/android/server/devicepolicy/NetworkLogger.java
index b82cb3c..0085931 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/NetworkLogger.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/NetworkLogger.java
@@ -31,7 +31,6 @@
 
 import com.android.server.ServiceThread;
 
-import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.atomic.AtomicBoolean;
 
@@ -130,6 +129,8 @@
         Log.d(TAG, "Stopping network logging");
         // stop the logging regardless of whether we fail to unregister listener
         mIsLoggingEnabled.set(false);
+        discardLogs();
+
         try {
             if (!checkIpConnectivityMetricsService()) {
                 // the IIpConnectivityMetrics service should have been present at this point
@@ -140,9 +141,43 @@
             return mIpConnectivityMetrics.unregisterNetdEventCallback();
         } catch (RemoteException re) {
             Slog.wtf(TAG, "Failed to make remote calls to unregister the callback", re);
-        } finally {
-            mHandlerThread.quitSafely();
             return true;
+        } finally {
+            if (mHandlerThread != null) {
+                mHandlerThread.quitSafely();
+            }
+        }
+    }
+
+    /**
+     * If logs are being collected, keep collecting them but stop notifying the device owner that
+     * new logs are available (since they cannot be retrieved)
+     */
+    void pause() {
+        if (mNetworkLoggingHandler != null) {
+            mNetworkLoggingHandler.pause();
+        }
+    }
+
+    /**
+     * If logs are being collected, start notifying the device owner when logs are ready to be
+     * collected again (if it was paused).
+     * <p>If logging is enabled and there are logs ready to be retrieved, this method will attempt
+     * to notify the device owner. Therefore calling identity should be cleared before calling it
+     * (in case the method is called from a user other than the DO's user).
+     */
+    void resume() {
+        if (mNetworkLoggingHandler != null) {
+            mNetworkLoggingHandler.resume();
+        }
+    }
+
+    /**
+     * Discard all collected logs.
+     */
+    void discardLogs() {
+        if (mNetworkLoggingHandler != null) {
+            mNetworkLoggingHandler.discardLogs();
         }
     }
 
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/NetworkLoggingHandler.java b/services/devicepolicy/java/com/android/server/devicepolicy/NetworkLoggingHandler.java
index baa4c13..7d68412 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/NetworkLoggingHandler.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/NetworkLoggingHandler.java
@@ -55,10 +55,16 @@
     @GuardedBy("this")
     private ArrayList<NetworkEvent> mFullBatch;
 
-    // each full batch is represented by its token, which the DPC has to provide back to revieve it
+    @GuardedBy("this")
+    private boolean mPaused = false;
+
+    // each full batch is represented by its token, which the DPC has to provide back to retrieve it
     @GuardedBy("this")
     private long mCurrentFullBatchToken;
 
+    @GuardedBy("this")
+    private long mLastRetrievedFullBatchToken;
+
     NetworkLoggingHandler(Looper looper, DevicePolicyManagerService dpm) {
         super(looper);
         mDpm = dpm;
@@ -70,15 +76,19 @@
             case LOG_NETWORK_EVENT_MSG: {
                 NetworkEvent networkEvent = msg.getData().getParcelable(NETWORK_EVENT_KEY);
                 if (networkEvent != null) {
-                    mNetworkEvents.add(networkEvent);
-                    if (mNetworkEvents.size() >= MAX_EVENTS_PER_BATCH) {
-                        finalizeBatchAndNotifyDeviceOwnerIfNotEmpty();
+                    synchronized (NetworkLoggingHandler.this) {
+                        mNetworkEvents.add(networkEvent);
+                        if (mNetworkEvents.size() >= MAX_EVENTS_PER_BATCH) {
+                            finalizeBatchAndNotifyDeviceOwnerLocked();
+                        }
                     }
                 }
                 break;
             }
             case FINALIZE_BATCH_MSG: {
-                finalizeBatchAndNotifyDeviceOwnerIfNotEmpty();
+                synchronized (NetworkLoggingHandler.this) {
+                    finalizeBatchAndNotifyDeviceOwnerLocked();
+                }
                 break;
             }
         }
@@ -91,22 +101,49 @@
                 + "ms from now.");
     }
 
-    private synchronized void finalizeBatchAndNotifyDeviceOwnerIfNotEmpty() {
+    synchronized void pause() {
+        Log.d(TAG, "Paused network logging");
+        mPaused = true;
+    }
+
+    synchronized void resume() {
+        if (!mPaused) {
+            Log.d(TAG, "Attempted to resume network logging, but logging is not paused.");
+            return;
+        }
+
+        Log.d(TAG, "Resumed network logging. Current batch="
+                + mCurrentFullBatchToken + ", LastRetrievedBatch=" + mLastRetrievedFullBatchToken);
+        mPaused = false;
+
+        // If there is a full batch ready that the device owner hasn't been notified about, do it
+        // now.
+        if (mFullBatch != null && mFullBatch.size() > 0
+                && mLastRetrievedFullBatchToken != mCurrentFullBatchToken) {
+            scheduleBatchFinalization();
+            notifyDeviceOwnerLocked();
+        }
+    }
+
+    synchronized void discardLogs() {
+        mFullBatch = null;
+        mNetworkEvents = new ArrayList<NetworkEvent>();
+        Log.d(TAG, "Discarded all network logs");
+    }
+
+    @GuardedBy("this")
+    private void finalizeBatchAndNotifyDeviceOwnerLocked() {
         if (mNetworkEvents.size() > 0) {
             // finalize the batch and start a new one from scratch
             mFullBatch = mNetworkEvents;
             mCurrentFullBatchToken++;
             mNetworkEvents = new ArrayList<NetworkEvent>();
-            // notify DO that there's a new non-empty batch waiting
-            Bundle extras = new Bundle();
-            extras.putLong(DeviceAdminReceiver.EXTRA_NETWORK_LOGS_TOKEN, mCurrentFullBatchToken);
-            extras.putInt(DeviceAdminReceiver.EXTRA_NETWORK_LOGS_COUNT, mFullBatch.size());
-            Log.d(TAG, "Sending network logging batch broadcast to device owner, batchToken: "
-                    + mCurrentFullBatchToken);
-            mDpm.sendDeviceOwnerCommand(DeviceAdminReceiver.ACTION_NETWORK_LOGS_AVAILABLE, extras);
+            if (!mPaused) {
+                notifyDeviceOwnerLocked();
+            }
         } else {
             // don't notify the DO, since there are no events; DPC can still retrieve
-            // the last full batch
+            // the last full batch if not paused.
             Log.d(TAG, "Was about to finalize the batch, but there were no events to send to"
                     + " the DPC, the batchToken of last available batch: "
                     + mCurrentFullBatchToken);
@@ -115,10 +152,21 @@
         scheduleBatchFinalization();
     }
 
+    @GuardedBy("this")
+    private void notifyDeviceOwnerLocked() {
+        Bundle extras = new Bundle();
+        extras.putLong(DeviceAdminReceiver.EXTRA_NETWORK_LOGS_TOKEN, mCurrentFullBatchToken);
+        extras.putInt(DeviceAdminReceiver.EXTRA_NETWORK_LOGS_COUNT, mFullBatch.size());
+        Log.d(TAG, "Sending network logging batch broadcast to device owner, batchToken: "
+                + mCurrentFullBatchToken);
+        mDpm.sendDeviceOwnerCommand(DeviceAdminReceiver.ACTION_NETWORK_LOGS_AVAILABLE, extras);
+    }
+
     synchronized List<NetworkEvent> retrieveFullLogBatch(long batchToken) {
         if (batchToken != mCurrentFullBatchToken) {
             return null;
         }
+        mLastRetrievedFullBatchToken = mCurrentFullBatchToken;
         return mFullBatch;
     }
 }
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/SecurityLogMonitor.java b/services/devicepolicy/java/com/android/server/devicepolicy/SecurityLogMonitor.java
index 79702a8..18f06be 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/SecurityLogMonitor.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/SecurityLogMonitor.java
@@ -19,6 +19,7 @@
 import android.app.admin.DeviceAdminReceiver;
 import android.app.admin.SecurityLog;
 import android.app.admin.SecurityLog.SecurityEvent;
+import android.os.SystemClock;
 import android.util.Log;
 import android.util.Slog;
 
@@ -50,7 +51,7 @@
         mService = service;
     }
 
-    private static final boolean DEBUG = false;
+    private static final boolean DEBUG = false;  // STOPSHIP if true.
     private static final String TAG = "SecurityLogMonitor";
     /**
      * Each log entry can hold up to 4K bytes (but as of {@link android.os.Build.VERSION_CODES#N}
@@ -78,17 +79,25 @@
     private ArrayList<SecurityEvent> mPendingLogs = new ArrayList<SecurityEvent>();
     @GuardedBy("mLock")
     private boolean mAllowedToRetrieve = false;
-    // When DO will be allowed to retrieves the log, in milliseconds.
+
+    /**
+     * When DO will be allowed to retrieve the log, in milliseconds since boot (as per
+     * {@link SystemClock#elapsedRealtime()})
+     */
     @GuardedBy("mLock")
-    private long mNextAllowedRetrivalTimeMillis = -1;
+    private long mNextAllowedRetrievalTimeMillis = -1;
+    @GuardedBy("mLock")
+    private boolean mPaused = false;
 
     void start() {
+        Slog.i(TAG, "Starting security logging.");
         mLock.lock();
         try {
             if (mMonitorThread == null) {
                 mPendingLogs = new ArrayList<SecurityEvent>();
                 mAllowedToRetrieve = false;
-                mNextAllowedRetrivalTimeMillis = -1;
+                mNextAllowedRetrievalTimeMillis = -1;
+                mPaused = false;
 
                 mMonitorThread = new Thread(this);
                 mMonitorThread.start();
@@ -99,6 +108,7 @@
     }
 
     void stop() {
+        Slog.i(TAG, "Stopping security logging.");
         mLock.lock();
         try {
             if (mMonitorThread != null) {
@@ -111,7 +121,8 @@
                 // Reset state and clear buffer
                 mPendingLogs = new ArrayList<SecurityEvent>();
                 mAllowedToRetrieve = false;
-                mNextAllowedRetrivalTimeMillis = -1;
+                mNextAllowedRetrievalTimeMillis = -1;
+                mPaused = false;
                 mMonitorThread = null;
             }
         } finally {
@@ -120,6 +131,58 @@
     }
 
     /**
+     * If logs are being collected, keep collecting them but stop notifying the device owner that
+     * new logs are available (since they cannot be retrieved).
+     */
+    void pause() {
+        Slog.i(TAG, "Paused.");
+
+        mLock.lock();
+        mPaused = true;
+        mAllowedToRetrieve = false;
+        mLock.unlock();
+    }
+
+    /**
+     * If logs are being collected, start notifying the device owner when logs are ready to be
+     * retrieved again (if it was paused).
+     * <p>If logging is enabled and there are logs ready to be retrieved, this method will attempt
+     * to notify the device owner. Therefore calling identity should be cleared before calling it
+     * (in case the method is called from a user other than the DO's user).
+     */
+    void resume() {
+        mLock.lock();
+        try {
+            if (!mPaused) {
+                Log.d(TAG, "Attempted to resume, but logging is not paused.");
+                return;
+            }
+            mPaused = false;
+            mAllowedToRetrieve = false;
+        } finally {
+            mLock.unlock();
+        }
+
+        Slog.i(TAG, "Resumed.");
+        try {
+            notifyDeviceOwnerIfNeeded();
+        } catch (InterruptedException e) {
+            Log.w(TAG, "Thread interrupted.", e);
+        }
+    }
+
+    /**
+     * Discard all collected logs.
+     */
+    void discardLogs() {
+        mLock.lock();
+        mAllowedToRetrieve = false;
+        mPendingLogs = new ArrayList<SecurityEvent>();
+        mLock.unlock();
+        Slog.i(TAG, "Discarded all logs.");
+    }
+
+    /**
      * Returns the new batch of logs since the last call to this method. Returns null if
      * rate limit is exceeded.
      */
@@ -128,7 +191,7 @@
         try {
             if (mAllowedToRetrieve) {
                 mAllowedToRetrieve = false;
-                mNextAllowedRetrivalTimeMillis = System.currentTimeMillis()
+                mNextAllowedRetrievalTimeMillis = SystemClock.elapsedRealtime()
                         + RATE_LIMIT_INTERVAL_MILLISECONDS;
                 List<SecurityEvent> result = mPendingLogs;
                 mPendingLogs = new ArrayList<SecurityEvent>();
@@ -163,7 +226,7 @@
                     SecurityLog.readEventsSince(lastLogTimestampNanos + 1, logs);
                 }
                 if (!logs.isEmpty()) {
-                    if (DEBUG) Slog.d(TAG, "processing new logs");
+                    if (DEBUG) Slog.d(TAG, "processing new logs. Events: " + logs.size());
                     mLock.lockInterruptibly();
                     try {
                         mPendingLogs.addAll(logs);
@@ -172,6 +235,7 @@
                             mPendingLogs = new ArrayList<SecurityEvent>(mPendingLogs.subList(
                                     mPendingLogs.size() - (BUFFER_ENTRIES_MAXIMUM_LEVEL / 2),
                                     mPendingLogs.size()));
+                            Slog.i(TAG, "Pending logs buffer full. Discarding old logs.");
                         }
                     } finally {
                         mLock.unlock();
@@ -188,7 +252,7 @@
                 break;
             }
         }
-        if (DEBUG) Slog.d(TAG, "MonitorThread exit.");
+        Slog.i(TAG, "MonitorThread exit.");
     }
 
     private void notifyDeviceOwnerIfNeeded() throws InterruptedException {
@@ -196,15 +260,24 @@
         boolean allowToRetrieveNow = false;
         mLock.lockInterruptibly();
         try {
+            if (mPaused) {
+                return;
+            }
+
+            // STOPSHIP(b/34186771): If the previous notification didn't reach the DO and logs were
+            // not retrieved (e.g. the broadcast was sent before the user was unlocked), no more
+            // subsequent callbacks will be sent. We should make sure that the DO gets notified
+            // before logs are lost.
             int logSize = mPendingLogs.size();
             if (logSize >= BUFFER_ENTRIES_NOTIFICATION_LEVEL) {
                 // Allow DO to retrieve logs if too many pending logs
                 allowToRetrieveNow = true;
+                if (DEBUG) Slog.d(TAG, "Number of log entries over threshold: " + logSize);
             } else if (logSize > 0) {
-                if (mNextAllowedRetrivalTimeMillis == -1 ||
-                        System.currentTimeMillis() >= mNextAllowedRetrivalTimeMillis) {
+                if (SystemClock.elapsedRealtime() >= mNextAllowedRetrievalTimeMillis) {
                     // Rate limit reset
                     allowToRetrieveNow = true;
+                    if (DEBUG) Slog.d(TAG, "Timeout reached");
                 }
             }
             shouldNotifyDO = (!mAllowedToRetrieve) && allowToRetrieveNow;
@@ -213,7 +286,7 @@
             mLock.unlock();
         }
         if (shouldNotifyDO) {
-            if (DEBUG) Slog.d(TAG, "notify DO");
+            Slog.i(TAG, "notify DO");
             mService.sendDeviceOwnerCommand(DeviceAdminReceiver.ACTION_SECURITY_LOGS_AVAILABLE,
                     null);
         }
diff --git a/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java b/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
index 182f045..33c8a12 100644
--- a/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
+++ b/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
@@ -2767,7 +2767,10 @@
     public void testGetLastSecurityLogRetrievalTime() throws Exception {
         mContext.binder.callingUid = DpmMockContext.CALLER_SYSTEM_USER_UID;
         setupDeviceOwner();
-        when(mContext.userManager.getUserCount()).thenReturn(1);
+
+        // setUp() adds a secondary user for CALLER_USER_HANDLE. Remove it as otherwise the
+        // feature is disabled because there are non-affiliated secondary users.
+        mContext.removeUser(DpmMockContext.CALLER_USER_HANDLE);
         when(mContext.resources.getBoolean(R.bool.config_supportPreRebootSecurityLogs))
                 .thenReturn(true);
 
@@ -2776,6 +2779,10 @@
 
         // Enabling logging should not change the timestamp.
         dpm.setSecurityLoggingEnabled(admin1, true);
+        verify(mContext.settings)
+                .securityLogSetLoggingEnabledProperty(true);
+        when(mContext.settings.securityLogGetLoggingEnabledProperty())
+                .thenReturn(true);
         assertEquals(-1, dpm.getLastSecurityLogRetrievalTime());
 
         // Retrieving the logs should update the timestamp.
@@ -2828,7 +2835,7 @@
     public void testGetLastBugReportRequestTime() throws Exception {
         mContext.binder.callingUid = DpmMockContext.CALLER_SYSTEM_USER_UID;
         setupDeviceOwner();
-        when(mContext.userManager.getUserCount()).thenReturn(1);
+
         mContext.packageName = admin1.getPackageName();
         mContext.applicationInfo = new ApplicationInfo();
         when(mContext.resources.getColor(eq(R.color.notification_action_list), anyObject()))
@@ -2836,6 +2843,10 @@
         when(mContext.resources.getColor(eq(R.color.notification_material_background_color),
                 anyObject())).thenReturn(Color.WHITE);
 
+        // setUp() adds a secondary user for CALLER_USER_HANDLE. Remove it as otherwise the
+        // feature is disabled because there are non-affiliated secondary users.
+        mContext.removeUser(DpmMockContext.CALLER_USER_HANDLE);
+
         // No bug reports were requested so far.
         assertEquals(-1, dpm.getLastBugReportRequestTime());
 
@@ -2873,7 +2884,16 @@
     public void testGetLastNetworkLogRetrievalTime() throws Exception {
         mContext.binder.callingUid = DpmMockContext.CALLER_SYSTEM_USER_UID;
         setupDeviceOwner();
-        when(mContext.userManager.getUserCount()).thenReturn(1);
+        mContext.packageName = admin1.getPackageName();
+        mContext.applicationInfo = new ApplicationInfo();
+        when(mContext.resources.getColor(eq(R.color.notification_action_list), anyObject()))
+                .thenReturn(Color.WHITE);
+        when(mContext.resources.getColor(eq(R.color.notification_material_background_color),
+                anyObject())).thenReturn(Color.WHITE);
+
+        // setUp() adds a secondary user for CALLER_USER_HANDLE. Remove it as otherwise the
+        // feature is disabled because there are non-affiliated secondary users.
+        mContext.removeUser(DpmMockContext.CALLER_USER_HANDLE);
         when(mContext.iipConnectivityMetrics.registerNetdEventCallback(anyObject()))
                 .thenReturn(true);
 
diff --git a/services/tests/servicestests/src/com/android/server/devicepolicy/DpmMockContext.java b/services/tests/servicestests/src/com/android/server/devicepolicy/DpmMockContext.java
index 65255d9..346af62 100644
--- a/services/tests/servicestests/src/com/android/server/devicepolicy/DpmMockContext.java
+++ b/services/tests/servicestests/src/com/android/server/devicepolicy/DpmMockContext.java
@@ -389,6 +389,18 @@
         return dir;
     }
 
+    public void removeUser(int userId) {
+        for (int i = 0; i < mUserInfos.size(); i++) {
+            if (mUserInfos.get(i).id == userId) {
+                mUserInfos.remove(i);
+                break;
+            }
+        }
+        when(userManager.getUserInfo(eq(userId))).thenReturn(null);
+
+        when(userManager.isUserRunning(eq(new UserHandle(userId)))).thenReturn(false);
+    }
+
     private UserInfo getUserInfo(int userId) {
         for (UserInfo ui : mUserInfos) {
             if (ui.id == userId) {