Profile lock timeout.

When a managed profile has separate lock (a.k.a work challenge)
and maximum screen off timeout set by admin, this timeout will
only cause the profile to get locked, not the whole device.

PowerManagerService now tracks some per-profile state for profiles
that have lock timeout set by admin and have separate lock:
1. timeout set by admin.
2. wake lock summary
3. last activity time
4. whether the profile is considered active
5. whether the profile got locked last time it went inactive.

Wake lock summaries and last activity times are updated at the
same time as global wake lock summary/user activity is updated.

Test: cts-tradefed run commandAndExit cts-dev -a armeabi-v7a -m CtsDevicePolicyManagerTestCases -t com.android.cts.devicepolicy.ManagedProfileTest
Test: runtest -c com.android.server.devicepolicy.DevicePolicyManagerTest frameworks-services
Bug: 63908311
Change-Id: Ief7be4e0bf12bdbedef94a129a13d07a9f98e75c
diff --git a/services/core/java/com/android/server/locksettings/LockSettingsService.java b/services/core/java/com/android/server/locksettings/LockSettingsService.java
index 60f451a..c883262 100644
--- a/services/core/java/com/android/server/locksettings/LockSettingsService.java
+++ b/services/core/java/com/android/server/locksettings/LockSettingsService.java
@@ -36,6 +36,7 @@
 import android.app.NotificationManager;
 import android.app.PendingIntent;
 import android.app.admin.DevicePolicyManager;
+import android.app.admin.DevicePolicyManagerInternal;
 import android.app.admin.PasswordMetrics;
 import android.app.backup.BackupManager;
 import android.app.trust.IStrongAuthTracker;
@@ -83,6 +84,7 @@
 import android.util.Log;
 import android.util.Slog;
 
+import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.messages.nano.SystemMessageProto.SystemMessage;
 import com.android.internal.notification.SystemNotificationChannels;
@@ -93,6 +95,7 @@
 import com.android.internal.widget.ILockSettings;
 import com.android.internal.widget.LockPatternUtils;
 import com.android.internal.widget.VerifyCredentialResponse;
+import com.android.server.LocalServices;
 import com.android.server.SystemService;
 import com.android.server.locksettings.LockSettingsStorage.CredentialHash;
 import com.android.server.locksettings.SyntheticPasswordManager.AuthenticationResult;
@@ -876,14 +879,26 @@
             String managedUserPassword) {
         checkWritePermission(userId);
         synchronized (mSeparateChallengeLock) {
-            setBoolean(SEPARATE_PROFILE_CHALLENGE_KEY, enabled, userId);
-            if (enabled) {
-                mStorage.removeChildProfileLock(userId);
-                removeKeystoreProfileKey(userId);
-            } else {
-                tieManagedProfileLockIfNecessary(userId, managedUserPassword);
-            }
+            setSeparateProfileChallengeEnabledLocked(userId, enabled, managedUserPassword);
         }
+        notifySeparateProfileChallengeChanged(userId);
+    }
+
+    @GuardedBy("mSeparateChallengeLock")
+    private void setSeparateProfileChallengeEnabledLocked(@UserIdInt int userId, boolean enabled,
+            String managedUserPassword) {
+        setBoolean(SEPARATE_PROFILE_CHALLENGE_KEY, enabled, userId);
+        if (enabled) {
+            mStorage.removeChildProfileLock(userId);
+            removeKeystoreProfileKey(userId);
+        } else {
+            tieManagedProfileLockIfNecessary(userId, managedUserPassword);
+        }
+    }
+
+    private void notifySeparateProfileChallengeChanged(int userId) {
+        LocalServices.getService(DevicePolicyManagerInternal.class)
+                .reportSeparateProfileChallengeChanged(userId);
     }
 
     @Override
@@ -1219,9 +1234,10 @@
         checkWritePermission(userId);
         synchronized (mSeparateChallengeLock) {
             setLockCredentialInternal(credential, type, savedCredential, requestedQuality, userId);
-            setSeparateProfileChallengeEnabled(userId, true, null);
+            setSeparateProfileChallengeEnabledLocked(userId, true, null);
             notifyPasswordChanged(userId);
         }
+        notifySeparateProfileChallengeChanged(userId);
     }
 
     private void setLockCredentialInternal(String credential, int credentialType,
@@ -2344,9 +2360,10 @@
         }
         if (result) {
             synchronized (mSeparateChallengeLock) {
-                setSeparateProfileChallengeEnabled(userId, true, null);
+                setSeparateProfileChallengeEnabledLocked(userId, true, null);
             }
             notifyPasswordChanged(userId);
+            notifySeparateProfileChallengeChanged(userId);
         }
         return result;
     }
diff --git a/services/core/java/com/android/server/power/Notifier.java b/services/core/java/com/android/server/power/Notifier.java
index 8ee26f29..e5a23ea 100644
--- a/services/core/java/com/android/server/power/Notifier.java
+++ b/services/core/java/com/android/server/power/Notifier.java
@@ -16,6 +16,7 @@
 
 package com.android.server.power;
 
+import android.annotation.UserIdInt;
 import android.app.ActivityManagerInternal;
 import android.app.AppOpsManager;
 
@@ -27,6 +28,7 @@
 import com.android.server.LocalServices;
 import com.android.server.policy.WindowManagerPolicy;
 
+import android.app.trust.TrustManager;
 import android.content.BroadcastReceiver;
 import android.content.Context;
 import android.content.Intent;
@@ -82,6 +84,7 @@
     private static final int MSG_BROADCAST = 2;
     private static final int MSG_WIRELESS_CHARGING_STARTED = 3;
     private static final int MSG_SCREEN_BRIGHTNESS_BOOST_CHANGED = 4;
+    private static final int MSG_PROFILE_TIMED_OUT = 5;
 
     private final Object mLock = new Object();
 
@@ -93,6 +96,7 @@
     private final ActivityManagerInternal mActivityManagerInternal;
     private final InputManagerInternal mInputManagerInternal;
     private final InputMethodManagerInternal mInputMethodManagerInternal;
+    private final TrustManager mTrustManager;
 
     private final NotifierHandler mHandler;
     private final Intent mScreenOnIntent;
@@ -138,6 +142,7 @@
         mActivityManagerInternal = LocalServices.getService(ActivityManagerInternal.class);
         mInputManagerInternal = LocalServices.getService(InputManagerInternal.class);
         mInputMethodManagerInternal = LocalServices.getService(InputMethodManagerInternal.class);
+        mTrustManager = mContext.getSystemService(TrustManager.class);
 
         mHandler = new NotifierHandler(looper);
         mScreenOnIntent = new Intent(Intent.ACTION_SCREEN_ON);
@@ -559,6 +564,16 @@
         mHandler.sendMessage(msg);
     }
 
+    /**
+     * Called when profile screen lock timeout has expired.
+     */
+    public void onProfileTimeout(@UserIdInt int userId) {
+        final Message msg = mHandler.obtainMessage(MSG_PROFILE_TIMED_OUT);
+        msg.setAsynchronous(true);
+        msg.arg1 = userId;
+        mHandler.sendMessage(msg);
+    }
+
     private void updatePendingBroadcastLocked() {
         if (!mBroadcastInProgress
                 && mPendingInteractiveState != INTERACTIVE_STATE_UNKNOWN
@@ -710,28 +725,33 @@
         mSuspendBlocker.release();
     }
 
+    private void lockProfile(@UserIdInt int userId) {
+        mTrustManager.setDeviceLockedForUser(userId, true /*locked*/);
+    }
+
     private final class NotifierHandler extends Handler {
+
         public NotifierHandler(Looper looper) {
             super(looper, null, true /*async*/);
         }
-
         @Override
         public void handleMessage(Message msg) {
             switch (msg.what) {
                 case MSG_USER_ACTIVITY:
                     sendUserActivity();
                     break;
-
                 case MSG_BROADCAST:
                     sendNextBroadcast();
                     break;
-
                 case MSG_WIRELESS_CHARGING_STARTED:
                     playWirelessChargingStartedSound();
                     break;
                 case MSG_SCREEN_BRIGHTNESS_BOOST_CHANGED:
                     sendBrightnessBoostChangedBroadcast();
                     break;
+                case MSG_PROFILE_TIMED_OUT:
+                    lockProfile(msg.arg1);
+                    break;
             }
         }
     }
diff --git a/services/core/java/com/android/server/power/PowerManagerService.java b/services/core/java/com/android/server/power/PowerManagerService.java
index 7f1a534..0b590bc 100644
--- a/services/core/java/com/android/server/power/PowerManagerService.java
+++ b/services/core/java/com/android/server/power/PowerManagerService.java
@@ -16,9 +16,10 @@
 
 package com.android.server.power;
 
-import android.Manifest;
 import android.annotation.IntDef;
+import android.annotation.UserIdInt;
 import android.app.ActivityManager;
+import android.app.SynchronousUserSwitchObserver;
 import android.content.BroadcastReceiver;
 import android.content.ContentResolver;
 import android.content.Context;
@@ -410,12 +411,12 @@
     private boolean mDozeAfterScreenOffConfig;
 
     // The minimum screen off timeout, in milliseconds.
-    private int mMinimumScreenOffTimeoutConfig;
+    private long mMinimumScreenOffTimeoutConfig;
 
     // The screen dim duration, in milliseconds.
     // This is subtracted from the end of the screen off timeout so the
     // minimum screen off timeout should be longer than this.
-    private int mMaximumScreenDimDurationConfig;
+    private long mMaximumScreenDimDurationConfig;
 
     // The maximum screen dim time expressed as a ratio relative to the screen
     // off timeout.  If the screen off timeout is very short then we want the
@@ -427,14 +428,14 @@
     private boolean mSupportsDoubleTapWakeConfig;
 
     // The screen off timeout setting value in milliseconds.
-    private int mScreenOffTimeoutSetting;
+    private long mScreenOffTimeoutSetting;
 
     // The sleep timeout setting value in milliseconds.
-    private int mSleepTimeoutSetting;
+    private long mSleepTimeoutSetting;
 
     // The maximum allowable screen off timeout according to the device
     // administration policy.  Overrides other settings.
-    private int mMaximumScreenOffTimeoutFromDeviceAdmin = Integer.MAX_VALUE;
+    private long mMaximumScreenOffTimeoutFromDeviceAdmin = Long.MAX_VALUE;
 
     // The stay on while plugged in setting.
     // A bitfield of battery conditions under which to make the screen stay on.
@@ -555,6 +556,46 @@
     // True if we are currently in VR Mode.
     private boolean mIsVrModeEnabled;
 
+    private final class ForegroundProfileObserver extends SynchronousUserSwitchObserver {
+        @Override
+        public void onUserSwitching(int newUserId) throws RemoteException {}
+
+        @Override
+        public void onForegroundProfileSwitch(@UserIdInt int newProfileId) throws RemoteException {
+            final long now = SystemClock.uptimeMillis();
+            synchronized(mLock) {
+                mForegroundProfile = newProfileId;
+                maybeUpdateForegroundProfileLastActivityLocked(now);
+            }
+        }
+    }
+
+    // User id corresponding to activity the user is currently interacting with.
+    private @UserIdInt int mForegroundProfile;
+
+    // Per-profile state to track when a profile should be locked.
+    private final SparseArray<ProfilePowerState> mProfilePowerState = new SparseArray<>();
+
+    private static final class ProfilePowerState {
+        // Profile user id.
+        final @UserIdInt int mUserId;
+        // Maximum time to lock set by admin.
+        long mScreenOffTimeout;
+        // Like top-level mWakeLockSummary, but only for wake locks that affect current profile.
+        int mWakeLockSummary;
+        // Last user activity that happened in an app running in the profile.
+        long mLastUserActivityTime;
+        // Whether profile has been locked last time it timed out.
+        boolean mLockingNotified;
+
+        public ProfilePowerState(@UserIdInt int userId, long screenOffTimeout) {
+            mUserId = userId;
+            mScreenOffTimeout = screenOffTimeout;
+            // Not accurate but at least won't cause immediate locking of the profile.
+            mLastUserActivityTime = SystemClock.uptimeMillis();
+        }
+    }
+
     /**
      * All times are in milliseconds. These constants are kept synchronized with the system
      * global Settings. Any access to this class or its fields should be done while
@@ -752,6 +793,12 @@
             mDisplayManagerInternal.initPowerManagement(
                     mDisplayPowerCallbacks, mHandler, sensorManager);
 
+            try {
+                final ForegroundProfileObserver observer = new ForegroundProfileObserver();
+                ActivityManager.getService().registerUserSwitchObserver(observer, TAG);
+            } catch (RemoteException e) {
+                // Shouldn't happen since in-process.
+            }
 
             // Go.
             readConfigurationLocked();
@@ -1333,6 +1380,8 @@
                 return false;
             }
 
+            maybeUpdateForegroundProfileLastActivityLocked(eventTime);
+
             if ((flags & PowerManager.USER_ACTIVITY_FLAG_NO_CHANGE_LIGHTS) != 0) {
                 if (eventTime > mLastUserActivityTimeNoChangeLights
                         && eventTime > mLastUserActivityTime) {
@@ -1360,6 +1409,13 @@
         return false;
     }
 
+    private void maybeUpdateForegroundProfileLastActivityLocked(long eventTime) {
+        final ProfilePowerState profile = mProfilePowerState.get(mForegroundProfile);
+        if (profile != null && eventTime > profile.mLastUserActivityTime) {
+            profile.mLastUserActivityTime = eventTime;
+        }
+    }
+
     private void wakeUpInternal(long eventTime, String reason, int uid, String opPackageName,
             int opUid) {
         synchronized (mLock) {
@@ -1648,16 +1704,19 @@
                 }
             }
 
-            // Phase 2: Update display power state.
-            boolean displayBecameReady = updateDisplayPowerStateLocked(dirtyPhase2);
+            // Phase 2: Lock profiles that became inactive/not kept awake.
+            updateProfilesLocked(now);
 
-            // Phase 3: Update dream state (depends on display ready signal).
+            // Phase 3: Update display power state.
+            final boolean displayBecameReady = updateDisplayPowerStateLocked(dirtyPhase2);
+
+            // Phase 4: Update dream state (depends on display ready signal).
             updateDreamLocked(dirtyPhase2, displayBecameReady);
 
-            // Phase 4: Send notifications, if needed.
+            // Phase 5: Send notifications, if needed.
             finishWakefulnessChangeIfNeededLocked();
 
-            // Phase 5: Update suspend blocker.
+            // Phase 6: Update suspend blocker.
             // Because we might release the last suspend blocker here, we need to make sure
             // we finished everything else first!
             updateSuspendBlockerLocked();
@@ -1667,6 +1726,29 @@
     }
 
     /**
+     * Check profile timeouts and notify profiles that should be locked.
+     */
+    private void updateProfilesLocked(long now) {
+        final int numProfiles = mProfilePowerState.size();
+        for (int i = 0; i < numProfiles; i++) {
+            final ProfilePowerState profile = mProfilePowerState.valueAt(i);
+            if (isProfileBeingKeptAwakeLocked(profile, now)) {
+                profile.mLockingNotified = false;
+            } else if (!profile.mLockingNotified) {
+                profile.mLockingNotified = true;
+                mNotifier.onProfileTimeout(profile.mUserId);
+            }
+        }
+    }
+
+    private boolean isProfileBeingKeptAwakeLocked(ProfilePowerState profile, long now) {
+        return (profile.mLastUserActivityTime + profile.mScreenOffTimeout > now)
+                || (profile.mWakeLockSummary & WAKE_LOCK_STAY_AWAKE) != 0
+                || (mProximityPositive &&
+                    (profile.mWakeLockSummary & WAKE_LOCK_PROXIMITY_SCREEN_OFF) != 0);
+    }
+
+    /**
      * Updates the value of mIsPowered.
      * Sets DIRTY_IS_POWERED if a change occurred.
      */
@@ -1800,60 +1882,28 @@
         if ((dirty & (DIRTY_WAKE_LOCKS | DIRTY_WAKEFULNESS)) != 0) {
             mWakeLockSummary = 0;
 
+            final int numProfiles = mProfilePowerState.size();
+            for (int i = 0; i < numProfiles; i++) {
+                mProfilePowerState.valueAt(i).mWakeLockSummary = 0;
+            }
+
             final int numWakeLocks = mWakeLocks.size();
             for (int i = 0; i < numWakeLocks; i++) {
                 final WakeLock wakeLock = mWakeLocks.get(i);
-                switch (wakeLock.mFlags & PowerManager.WAKE_LOCK_LEVEL_MASK) {
-                    case PowerManager.PARTIAL_WAKE_LOCK:
-                        if (!wakeLock.mDisabled) {
-                            // We only respect this if the wake lock is not disabled.
-                            mWakeLockSummary |= WAKE_LOCK_CPU;
-                        }
-                        break;
-                    case PowerManager.FULL_WAKE_LOCK:
-                        mWakeLockSummary |= WAKE_LOCK_SCREEN_BRIGHT | WAKE_LOCK_BUTTON_BRIGHT;
-                        break;
-                    case PowerManager.SCREEN_BRIGHT_WAKE_LOCK:
-                        mWakeLockSummary |= WAKE_LOCK_SCREEN_BRIGHT;
-                        break;
-                    case PowerManager.SCREEN_DIM_WAKE_LOCK:
-                        mWakeLockSummary |= WAKE_LOCK_SCREEN_DIM;
-                        break;
-                    case PowerManager.PROXIMITY_SCREEN_OFF_WAKE_LOCK:
-                        mWakeLockSummary |= WAKE_LOCK_PROXIMITY_SCREEN_OFF;
-                        break;
-                    case PowerManager.DOZE_WAKE_LOCK:
-                        mWakeLockSummary |= WAKE_LOCK_DOZE;
-                        break;
-                    case PowerManager.DRAW_WAKE_LOCK:
-                        mWakeLockSummary |= WAKE_LOCK_DRAW;
-                        break;
+                final int wakeLockFlags = getWakeLockSummaryFlags(wakeLock);
+                mWakeLockSummary |= wakeLockFlags;
+                for (int j = 0; j < numProfiles; j++) {
+                    final ProfilePowerState profile = mProfilePowerState.valueAt(j);
+                    if (wakeLockAffectsUser(wakeLock, profile.mUserId)) {
+                        profile.mWakeLockSummary |= wakeLockFlags;
+                    }
                 }
             }
 
-            // Cancel wake locks that make no sense based on the current state.
-            if (mWakefulness != WAKEFULNESS_DOZING) {
-                mWakeLockSummary &= ~(WAKE_LOCK_DOZE | WAKE_LOCK_DRAW);
-            }
-            if (mWakefulness == WAKEFULNESS_ASLEEP
-                    || (mWakeLockSummary & WAKE_LOCK_DOZE) != 0) {
-                mWakeLockSummary &= ~(WAKE_LOCK_SCREEN_BRIGHT | WAKE_LOCK_SCREEN_DIM
-                        | WAKE_LOCK_BUTTON_BRIGHT);
-                if (mWakefulness == WAKEFULNESS_ASLEEP) {
-                    mWakeLockSummary &= ~WAKE_LOCK_PROXIMITY_SCREEN_OFF;
-                }
-            }
-
-            // Infer implied wake locks where necessary based on the current state.
-            if ((mWakeLockSummary & (WAKE_LOCK_SCREEN_BRIGHT | WAKE_LOCK_SCREEN_DIM)) != 0) {
-                if (mWakefulness == WAKEFULNESS_AWAKE) {
-                    mWakeLockSummary |= WAKE_LOCK_CPU | WAKE_LOCK_STAY_AWAKE;
-                } else if (mWakefulness == WAKEFULNESS_DREAMING) {
-                    mWakeLockSummary |= WAKE_LOCK_CPU;
-                }
-            }
-            if ((mWakeLockSummary & WAKE_LOCK_DRAW) != 0) {
-                mWakeLockSummary |= WAKE_LOCK_CPU;
+            mWakeLockSummary = adjustWakeLockSummaryLocked(mWakeLockSummary);
+            for (int i = 0; i < numProfiles; i++) {
+                final ProfilePowerState profile = mProfilePowerState.valueAt(i);
+                profile.mWakeLockSummary = adjustWakeLockSummaryLocked(profile.mWakeLockSummary);
             }
 
             if (DEBUG_SPEW) {
@@ -1864,6 +1914,72 @@
         }
     }
 
+    private int adjustWakeLockSummaryLocked(int wakeLockSummary) {
+        // Cancel wake locks that make no sense based on the current state.
+        if (mWakefulness != WAKEFULNESS_DOZING) {
+            wakeLockSummary &= ~(WAKE_LOCK_DOZE | WAKE_LOCK_DRAW);
+        }
+        if (mWakefulness == WAKEFULNESS_ASLEEP
+                || (wakeLockSummary & WAKE_LOCK_DOZE) != 0) {
+            wakeLockSummary &= ~(WAKE_LOCK_SCREEN_BRIGHT | WAKE_LOCK_SCREEN_DIM
+                    | WAKE_LOCK_BUTTON_BRIGHT);
+            if (mWakefulness == WAKEFULNESS_ASLEEP) {
+                wakeLockSummary &= ~WAKE_LOCK_PROXIMITY_SCREEN_OFF;
+            }
+        }
+
+        // Infer implied wake locks where necessary based on the current state.
+        if ((wakeLockSummary & (WAKE_LOCK_SCREEN_BRIGHT | WAKE_LOCK_SCREEN_DIM)) != 0) {
+            if (mWakefulness == WAKEFULNESS_AWAKE) {
+                wakeLockSummary |= WAKE_LOCK_CPU | WAKE_LOCK_STAY_AWAKE;
+            } else if (mWakefulness == WAKEFULNESS_DREAMING) {
+                wakeLockSummary |= WAKE_LOCK_CPU;
+            }
+        }
+        if ((wakeLockSummary & WAKE_LOCK_DRAW) != 0) {
+            wakeLockSummary |= WAKE_LOCK_CPU;
+        }
+
+        return wakeLockSummary;
+    }
+
+    /** Get wake lock summary flags that correspond to the given wake lock. */
+    private int getWakeLockSummaryFlags(WakeLock wakeLock) {
+        switch (wakeLock.mFlags & PowerManager.WAKE_LOCK_LEVEL_MASK) {
+            case PowerManager.PARTIAL_WAKE_LOCK:
+                if (!wakeLock.mDisabled) {
+                    // We only respect this if the wake lock is not disabled.
+                    return WAKE_LOCK_CPU;
+                }
+                break;
+            case PowerManager.FULL_WAKE_LOCK:
+                return WAKE_LOCK_SCREEN_BRIGHT | WAKE_LOCK_BUTTON_BRIGHT;
+            case PowerManager.SCREEN_BRIGHT_WAKE_LOCK:
+                return WAKE_LOCK_SCREEN_BRIGHT;
+            case PowerManager.SCREEN_DIM_WAKE_LOCK:
+                return WAKE_LOCK_SCREEN_DIM;
+            case PowerManager.PROXIMITY_SCREEN_OFF_WAKE_LOCK:
+                return WAKE_LOCK_PROXIMITY_SCREEN_OFF;
+            case PowerManager.DOZE_WAKE_LOCK:
+                return WAKE_LOCK_DOZE;
+            case PowerManager.DRAW_WAKE_LOCK:
+                return WAKE_LOCK_DRAW;
+        }
+        return 0;
+    }
+
+    private boolean wakeLockAffectsUser(WakeLock wakeLock, @UserIdInt int userId) {
+        if (wakeLock.mWorkSource != null) {
+            for (int k = 0; k < wakeLock.mWorkSource.size(); k++) {
+                final int uid = wakeLock.mWorkSource.get(k);
+                if (userId == UserHandle.getUserId(uid)) {
+                    return true;
+                }
+            }
+        }
+        return userId == UserHandle.getUserId(wakeLock.mOwnerUid);
+    }
+
     void checkForLongWakeLocks() {
         synchronized (mLock) {
             final long now = SystemClock.uptimeMillis();
@@ -1917,10 +2033,11 @@
             if (mWakefulness == WAKEFULNESS_AWAKE
                     || mWakefulness == WAKEFULNESS_DREAMING
                     || mWakefulness == WAKEFULNESS_DOZING) {
-                final int sleepTimeout = getSleepTimeoutLocked();
-                final int screenOffTimeout = getScreenOffTimeoutLocked(sleepTimeout);
-                final int screenDimDuration = getScreenDimDurationLocked(screenOffTimeout);
+                final long sleepTimeout = getSleepTimeoutLocked();
+                final long screenOffTimeout = getScreenOffTimeoutLocked(sleepTimeout);
+                final long screenDimDuration = getScreenDimDurationLocked(screenOffTimeout);
                 final boolean userInactiveOverride = mUserInactiveOverrideFromWindowManager;
+                final long nextProfileTimeout = getNextProfileTimeoutLocked(now);
 
                 mUserActivitySummary = 0;
                 if (mLastUserActivityTime >= mLastWakeTime) {
@@ -1977,10 +2094,12 @@
                     nextTimeout = -1;
                 }
 
+                if (nextProfileTimeout > 0) {
+                    nextTimeout = Math.min(nextTimeout, nextProfileTimeout);
+                }
+
                 if (mUserActivitySummary != 0 && nextTimeout >= 0) {
-                    Message msg = mHandler.obtainMessage(MSG_USER_ACTIVITY_TIMEOUT);
-                    msg.setAsynchronous(true);
-                    mHandler.sendMessageAtTime(msg, nextTimeout);
+                    scheduleUserInactivityTimeout(nextTimeout);
                 }
             } else {
                 mUserActivitySummary = 0;
@@ -1995,6 +2114,28 @@
         }
     }
 
+    private void scheduleUserInactivityTimeout(long timeMs) {
+        final Message msg = mHandler.obtainMessage(MSG_USER_ACTIVITY_TIMEOUT);
+        msg.setAsynchronous(true);
+        mHandler.sendMessageAtTime(msg, timeMs);
+    }
+
+    /**
+     * Finds the next profile timeout time or returns -1 if there are no profiles to be locked.
+     */
+    private long getNextProfileTimeoutLocked(long now) {
+        long nextTimeout = -1;
+        final int numProfiles = mProfilePowerState.size();
+        for (int i = 0; i < numProfiles; i++) {
+            final ProfilePowerState profile = mProfilePowerState.valueAt(i);
+            final long timeout = profile.mLastUserActivityTime + profile.mScreenOffTimeout;
+            if (timeout > now && (nextTimeout == -1 || timeout < nextTimeout)) {
+                nextTimeout = timeout;
+            }
+        }
+        return nextTimeout;
+    }
+
     /**
      * Called when a user activity timeout has occurred.
      * Simply indicates that something about user activity has changed so that the new
@@ -2014,21 +2155,21 @@
         }
     }
 
-    private int getSleepTimeoutLocked() {
-        int timeout = mSleepTimeoutSetting;
+    private long getSleepTimeoutLocked() {
+        final long timeout = mSleepTimeoutSetting;
         if (timeout <= 0) {
             return -1;
         }
         return Math.max(timeout, mMinimumScreenOffTimeoutConfig);
     }
 
-    private int getScreenOffTimeoutLocked(int sleepTimeout) {
-        int timeout = mScreenOffTimeoutSetting;
+    private long getScreenOffTimeoutLocked(long sleepTimeout) {
+        long timeout = mScreenOffTimeoutSetting;
         if (isMaximumScreenOffTimeoutFromDeviceAdminEnforcedLocked()) {
             timeout = Math.min(timeout, mMaximumScreenOffTimeoutFromDeviceAdmin);
         }
         if (mUserActivityTimeoutOverrideFromWindowManager >= 0) {
-            timeout = (int)Math.min(timeout, mUserActivityTimeoutOverrideFromWindowManager);
+            timeout = Math.min(timeout, mUserActivityTimeoutOverrideFromWindowManager);
         }
         if (sleepTimeout >= 0) {
             timeout = Math.min(timeout, sleepTimeout);
@@ -2036,9 +2177,9 @@
         return Math.max(timeout, mMinimumScreenOffTimeoutConfig);
     }
 
-    private int getScreenDimDurationLocked(int screenOffTimeout) {
+    private long getScreenDimDurationLocked(long screenOffTimeout) {
         return Math.min(mMaximumScreenDimDurationConfig,
-                (int)(screenOffTimeout * mMaximumScreenDimRatioConfig));
+                (long)(screenOffTimeout * mMaximumScreenDimRatioConfig));
     }
 
     /**
@@ -2781,9 +2922,27 @@
                 Settings.Global.STAY_ON_WHILE_PLUGGED_IN, val);
     }
 
-    void setMaximumScreenOffTimeoutFromDeviceAdminInternal(int timeMs) {
+    void setMaximumScreenOffTimeoutFromDeviceAdminInternal(@UserIdInt int userId, long timeMs) {
+        if (userId < 0) {
+            Slog.wtf(TAG, "Attempt to set screen off timeout for invalid user: " + userId);
+            return;
+        }
         synchronized (mLock) {
-            mMaximumScreenOffTimeoutFromDeviceAdmin = timeMs;
+            // System-wide timeout
+            if (userId == UserHandle.USER_SYSTEM) {
+                mMaximumScreenOffTimeoutFromDeviceAdmin = timeMs;
+            } else if (timeMs == Long.MAX_VALUE || timeMs == 0) {
+                mProfilePowerState.delete(userId);
+            } else {
+                final ProfilePowerState profile = mProfilePowerState.get(userId);
+                if (profile != null) {
+                    profile.mScreenOffTimeout = timeMs;
+                } else {
+                    mProfilePowerState.put(userId, new ProfilePowerState(userId, timeMs));
+                    // We need to recalculate wake locks for the new profile state.
+                    mDirty |= DIRTY_WAKE_LOCKS;
+                }
+            }
             mDirty |= DIRTY_SETTINGS;
             updatePowerStateLocked();
         }
@@ -2981,7 +3140,7 @@
 
     private boolean isMaximumScreenOffTimeoutFromDeviceAdminEnforcedLocked() {
         return mMaximumScreenOffTimeoutFromDeviceAdmin >= 0
-                && mMaximumScreenOffTimeoutFromDeviceAdmin < Integer.MAX_VALUE;
+                && mMaximumScreenOffTimeoutFromDeviceAdmin < Long.MAX_VALUE;
     }
 
     private void setAttentionLightInternal(boolean on, int color) {
@@ -3325,10 +3484,11 @@
             pw.println("  mScreenBrightnessForVrSetting=" + mScreenBrightnessForVrSetting);
             pw.println("  mDoubleTapWakeEnabled=" + mDoubleTapWakeEnabled);
             pw.println("  mIsVrModeEnabled=" + mIsVrModeEnabled);
+            pw.println("  mForegroundProfile=" + mForegroundProfile);
 
-            final int sleepTimeout = getSleepTimeoutLocked();
-            final int screenOffTimeout = getScreenOffTimeoutLocked(sleepTimeout);
-            final int screenDimDuration = getScreenDimDurationLocked(screenOffTimeout);
+            final long sleepTimeout = getSleepTimeoutLocked();
+            final long screenOffTimeout = getScreenOffTimeoutLocked(sleepTimeout);
+            final long screenDimDuration = getScreenDimDurationLocked(screenOffTimeout);
             pw.println();
             pw.println("Sleep timeout: " + sleepTimeout + " ms");
             pw.println("Screen off timeout: " + screenOffTimeout + " ms");
@@ -3373,6 +3533,23 @@
 
             mBatterySaverPolicy.dump(pw);
 
+            pw.println();
+            final int numProfiles = mProfilePowerState.size();
+            pw.println("Profile power states: size=" + numProfiles);
+            for (int i = 0; i < numProfiles; i++) {
+                final ProfilePowerState profile = mProfilePowerState.valueAt(i);
+                pw.print("  mUserId=");
+                pw.print(profile.mUserId);
+                pw.print(" mScreenOffTimeout=");
+                pw.print(profile.mScreenOffTimeout);
+                pw.print(" mWakeLockSummary=");
+                pw.print(profile.mWakeLockSummary);
+                pw.print(" mLastUserActivityTime=");
+                pw.print(profile.mLastUserActivityTime);
+                pw.print(" mLockingNotified=");
+                pw.println(profile.mLockingNotified);
+            }
+
             wcd = mWirelessChargerDetector;
         }
 
@@ -3590,7 +3767,8 @@
             proto.write(
                     PowerServiceSettingsAndConfigurationDumpProto
                             .MAXIMUM_SCREEN_OFF_TIMEOUT_FROM_DEVICE_ADMIN_MS,
-                    mMaximumScreenOffTimeoutFromDeviceAdmin);
+                    // Clamp to int32
+                    Math.min(mMaximumScreenOffTimeoutFromDeviceAdmin, Integer.MAX_VALUE));
             proto.write(
                     PowerServiceSettingsAndConfigurationDumpProto
                             .IS_MAXIMUM_SCREEN_OFF_TIMEOUT_FROM_DEVICE_ADMIN_ENFORCED_LOCKED,
@@ -3686,9 +3864,9 @@
                     mIsVrModeEnabled);
             proto.end(settingsAndConfigurationToken);
 
-            final int sleepTimeout = getSleepTimeoutLocked();
-            final int screenOffTimeout = getScreenOffTimeoutLocked(sleepTimeout);
-            final int screenDimDuration = getScreenDimDurationLocked(screenOffTimeout);
+            final long sleepTimeout = getSleepTimeoutLocked();
+            final long screenOffTimeout = getScreenOffTimeoutLocked(sleepTimeout);
+            final long screenDimDuration = getScreenDimDurationLocked(screenOffTimeout);
             proto.write(PowerManagerServiceDumpProto.SLEEP_TIMEOUT_MS, sleepTimeout);
             proto.write(PowerManagerServiceDumpProto.SCREEN_OFF_TIMEOUT_MS, screenOffTimeout);
             proto.write(PowerManagerServiceDumpProto.SCREEN_DIM_DURATION_MS, screenDimDuration);
@@ -4697,8 +4875,8 @@
         }
 
         @Override
-        public void setMaximumScreenOffTimeoutFromDeviceAdmin(int timeMs) {
-            setMaximumScreenOffTimeoutFromDeviceAdminInternal(timeMs);
+        public void setMaximumScreenOffTimeoutFromDeviceAdmin(@UserIdInt int userId, long timeMs) {
+            setMaximumScreenOffTimeoutFromDeviceAdminInternal(userId, timeMs);
         }
 
         @Override
diff --git a/services/core/java/com/android/server/trust/TrustAgentWrapper.java b/services/core/java/com/android/server/trust/TrustAgentWrapper.java
index 1382894..ca0a450 100644
--- a/services/core/java/com/android/server/trust/TrustAgentWrapper.java
+++ b/services/core/java/com/android/server/trust/TrustAgentWrapper.java
@@ -512,7 +512,7 @@
                 } else {
                     mTrustAgentService.onConfigure(Collections.EMPTY_LIST, null);
                 }
-                final long maxTimeToLock = dpm.getMaximumTimeToLockForUserAndProfiles(mUserId);
+                final long maxTimeToLock = dpm.getMaximumTimeToLock(null, mUserId);
                 if (maxTimeToLock != mMaximumTimeToLock) {
                     // If the timeout changes, cancel the alarm and send a timeout event to have
                     // the agent re-evaluate trust.
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
index 663083c..a8d8e5b 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
@@ -4603,56 +4603,56 @@
         }
     }
 
-    void updateMaximumTimeToLockLocked(int userHandle) {
-        // Calculate the min timeout for all profiles - including the ones with a separate
-        // challenge. Ideally if the timeout only affected the profile challenge we'd lock that
-        // challenge only and keep the screen on. However there is no easy way of doing that at the
-        // moment so we set the screen off timeout regardless of whether it affects the parent user
-        // or the profile challenge only.
-        long timeMs = Long.MAX_VALUE;
-        int[] profileIds = mUserManager.getProfileIdsWithDisabled(userHandle);
-        for (int profileId : profileIds) {
-            DevicePolicyData policy = getUserDataUnchecked(profileId);
-            final int N = policy.mAdminList.size();
-            for (int i = 0; i < N; i++) {
-                ActiveAdmin admin = policy.mAdminList.get(i);
-                if (admin.maximumTimeToUnlock > 0
-                        && timeMs > admin.maximumTimeToUnlock) {
-                    timeMs = admin.maximumTimeToUnlock;
-                }
-                // If userInfo.id is a managed profile, we also need to look at
-                // the policies set on the parent.
-                if (admin.hasParentActiveAdmin()) {
-                    final ActiveAdmin parentAdmin = admin.getParentActiveAdmin();
-                    if (parentAdmin.maximumTimeToUnlock > 0
-                            && timeMs > parentAdmin.maximumTimeToUnlock) {
-                        timeMs = parentAdmin.maximumTimeToUnlock;
-                    }
-                }
-            }
+    private void updateMaximumTimeToLockLocked(@UserIdInt int userId) {
+        // Update the profile's timeout
+        if (isManagedProfile(userId)) {
+            updateProfileLockTimeoutLocked(userId);
         }
 
-        // We only store the last maximum time to lock on the parent profile. So if calling from a
-        // managed profile, retrieve the policy for the parent.
-        DevicePolicyData policy = getUserDataUnchecked(getProfileParentId(userHandle));
-        if (policy.mLastMaximumTimeToLock == timeMs) {
-            return;
-        }
-        policy.mLastMaximumTimeToLock = timeMs;
-
+        final long timeMs;
         final long ident = mInjector.binderClearCallingIdentity();
         try {
+            // Update the device timeout
+            final int parentId = getProfileParentId(userId);
+            timeMs = getMaximumTimeToLockPolicyFromAdmins(
+                    getActiveAdminsForLockscreenPoliciesLocked(parentId, false));
+
+            final DevicePolicyData policy = getUserDataUnchecked(parentId);
+            if (policy.mLastMaximumTimeToLock == timeMs) {
+                return;
+            }
+            policy.mLastMaximumTimeToLock = timeMs;
+
             if (policy.mLastMaximumTimeToLock != Long.MAX_VALUE) {
                 // Make sure KEEP_SCREEN_ON is disabled, since that
                 // would allow bypassing of the maximum time to lock.
                 mInjector.settingsGlobalPutInt(Settings.Global.STAY_ON_WHILE_PLUGGED_IN, 0);
             }
-
-            mInjector.getPowerManagerInternal().setMaximumScreenOffTimeoutFromDeviceAdmin(
-                    (int) Math.min(policy.mLastMaximumTimeToLock, Integer.MAX_VALUE));
         } finally {
             mInjector.binderRestoreCallingIdentity(ident);
         }
+
+        mInjector.getPowerManagerInternal().setMaximumScreenOffTimeoutFromDeviceAdmin(
+                UserHandle.USER_SYSTEM, timeMs);
+    }
+
+    private void updateProfileLockTimeoutLocked(@UserIdInt int userId) {
+        final long timeMs;
+        if (isSeparateProfileChallengeEnabled(userId)) {
+            timeMs = getMaximumTimeToLockPolicyFromAdmins(
+                    getActiveAdminsForLockscreenPoliciesLocked(userId, false /* parent */));
+        } else {
+            timeMs = Long.MAX_VALUE;
+        }
+
+        final DevicePolicyData policy = getUserDataUnchecked(userId);
+        if (policy.mLastMaximumTimeToLock == timeMs) {
+            return;
+        }
+        policy.mLastMaximumTimeToLock = timeMs;
+
+        mInjector.getPowerManagerInternal().setMaximumScreenOffTimeoutFromDeviceAdmin(
+                userId, policy.mLastMaximumTimeToLock);
     }
 
     @Override
@@ -4663,50 +4663,21 @@
         enforceFullCrossUsersPermission(userHandle);
         synchronized (this) {
             if (who != null) {
-                ActiveAdmin admin = getActiveAdminUncheckedLocked(who, userHandle, parent);
+                final ActiveAdmin admin = getActiveAdminUncheckedLocked(who, userHandle, parent);
                 return admin != null ? admin.maximumTimeToUnlock : 0;
             }
             // Return the strictest policy across all participating admins.
-            List<ActiveAdmin> admins = getActiveAdminsForLockscreenPoliciesLocked(
+            final List<ActiveAdmin> admins = getActiveAdminsForLockscreenPoliciesLocked(
                     userHandle, parent);
-            return getMaximumTimeToLockPolicyFromAdmins(admins);
-        }
-    }
-
-    @Override
-    public long getMaximumTimeToLockForUserAndProfiles(int userHandle) {
-        if (!mHasFeature) {
-            return 0;
-        }
-        enforceFullCrossUsersPermission(userHandle);
-        synchronized (this) {
-            // All admins for this user.
-            ArrayList<ActiveAdmin> admins = new ArrayList<ActiveAdmin>();
-            for (UserInfo userInfo : mUserManager.getProfiles(userHandle)) {
-                DevicePolicyData policy = getUserData(userInfo.id);
-                admins.addAll(policy.mAdminList);
-                // If it is a managed profile, it may have parent active admins
-                if (userInfo.isManagedProfile()) {
-                    for (ActiveAdmin admin : policy.mAdminList) {
-                        if (admin.hasParentActiveAdmin()) {
-                            admins.add(admin.getParentActiveAdmin());
-                        }
-                    }
-                }
-            }
-            return getMaximumTimeToLockPolicyFromAdmins(admins);
+            final long timeMs = getMaximumTimeToLockPolicyFromAdmins(admins);
+            return timeMs == Long.MAX_VALUE ? 0 : timeMs;
         }
     }
 
     private long getMaximumTimeToLockPolicyFromAdmins(List<ActiveAdmin> admins) {
-        long time = 0;
-        final int N = admins.size();
-        for (int i = 0; i < N; i++) {
-            ActiveAdmin admin = admins.get(i);
-            if (time == 0) {
-                time = admin.maximumTimeToUnlock;
-            } else if (admin.maximumTimeToUnlock != 0
-                    && time > admin.maximumTimeToUnlock) {
+        long time = Long.MAX_VALUE;
+        for (final ActiveAdmin admin : admins) {
+            if (admin.maximumTimeToUnlock > 0 && admin.maximumTimeToUnlock < time) {
                 time = admin.maximumTimeToUnlock;
             }
         }
@@ -9132,7 +9103,7 @@
                 // ignore if it contradicts an existing policy
                 long timeMs = getMaximumTimeToLock(
                         who, mInjector.userHandleGetCallingUserId(), /* parent */ false);
-                if (timeMs > 0 && timeMs < Integer.MAX_VALUE) {
+                if (timeMs > 0 && timeMs < Long.MAX_VALUE) {
                     return;
                 }
             }
@@ -9594,6 +9565,13 @@
         public boolean isUserAffiliatedWithDevice(int userId) {
             return DevicePolicyManagerService.this.isUserAffiliatedWithDeviceLocked(userId);
         }
+
+        @Override
+        public void reportSeparateProfileChallengeChanged(@UserIdInt int userId) {
+            synchronized (DevicePolicyManagerService.this) {
+                updateMaximumTimeToLockLocked(userId);
+            }
+        }
     }
 
     private Intent createShowAdminSupportIntent(ComponentName admin, int userId) {
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 ca918c6..70e4b6d 100644
--- a/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
+++ b/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
@@ -34,7 +34,6 @@
 import static org.mockito.Mockito.atLeast;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doReturn;
-import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.nullable;
 import static org.mockito.Mockito.reset;
@@ -2244,27 +2243,32 @@
         reset(getServices().settings);
 
         dpm.setMaximumTimeToLock(admin1, 0);
-        verifyScreenTimeoutCall(null, false);
+        verifyScreenTimeoutCall(null, UserHandle.USER_SYSTEM);
+        verifyStayOnWhilePluggedCleared(false);
         reset(getServices().powerManagerInternal);
         reset(getServices().settings);
 
         dpm.setMaximumTimeToLock(admin1, 1);
-        verifyScreenTimeoutCall(1, true);
+        verifyScreenTimeoutCall(1L, UserHandle.USER_SYSTEM);
+        verifyStayOnWhilePluggedCleared(true);
         reset(getServices().powerManagerInternal);
         reset(getServices().settings);
 
         dpm.setMaximumTimeToLock(admin2, 10);
-        verifyScreenTimeoutCall(null, false);
+        verifyScreenTimeoutCall(null, UserHandle.USER_SYSTEM);
+        verifyStayOnWhilePluggedCleared(false);
         reset(getServices().powerManagerInternal);
         reset(getServices().settings);
 
         dpm.setMaximumTimeToLock(admin1, 5);
-        verifyScreenTimeoutCall(5, true);
+        verifyScreenTimeoutCall(5L, UserHandle.USER_SYSTEM);
+        verifyStayOnWhilePluggedCleared(true);
         reset(getServices().powerManagerInternal);
         reset(getServices().settings);
 
         dpm.setMaximumTimeToLock(admin2, 4);
-        verifyScreenTimeoutCall(4, true);
+        verifyScreenTimeoutCall(4L, UserHandle.USER_SYSTEM);
+        verifyStayOnWhilePluggedCleared(true);
         reset(getServices().powerManagerInternal);
         reset(getServices().settings);
 
@@ -2272,24 +2276,89 @@
         reset(getServices().powerManagerInternal);
         reset(getServices().settings);
 
-        dpm.setMaximumTimeToLock(admin2, Integer.MAX_VALUE);
-        verifyScreenTimeoutCall(Integer.MAX_VALUE, true);
-        reset(getServices().powerManagerInternal);
-        reset(getServices().settings);
-
-        dpm.setMaximumTimeToLock(admin2, Integer.MAX_VALUE + 1);
-        verifyScreenTimeoutCall(Integer.MAX_VALUE, true);
+        dpm.setMaximumTimeToLock(admin2, Long.MAX_VALUE);
+        verifyScreenTimeoutCall(Long.MAX_VALUE, UserHandle.USER_SYSTEM);
+        verifyStayOnWhilePluggedCleared(true);
         reset(getServices().powerManagerInternal);
         reset(getServices().settings);
 
         dpm.setMaximumTimeToLock(admin2, 10);
-        verifyScreenTimeoutCall(10, true);
+        verifyScreenTimeoutCall(10L, UserHandle.USER_SYSTEM);
+        verifyStayOnWhilePluggedCleared(true);
         reset(getServices().powerManagerInternal);
         reset(getServices().settings);
 
-        // There's no restriction; shold be set to MAX.
+        // There's no restriction; should be set to MAX.
         dpm.setMaximumTimeToLock(admin2, 0);
-        verifyScreenTimeoutCall(Integer.MAX_VALUE, false);
+        verifyScreenTimeoutCall(Long.MAX_VALUE, UserHandle.USER_SYSTEM);
+        verifyStayOnWhilePluggedCleared(false);
+    }
+
+    // Test if lock timeout on managed profile is handled correctly depending on whether profile
+    // uses separate challenge.
+    public void testSetMaximumTimeToLockProfile() throws Exception {
+        final int PROFILE_USER = 15;
+        final int PROFILE_ADMIN = UserHandle.getUid(PROFILE_USER, 19436);
+        addManagedProfile(admin1, PROFILE_ADMIN, admin1);
+        mContext.binder.callingUid = PROFILE_ADMIN;
+        final DevicePolicyManagerInternal dpmi =
+                LocalServices.getService(DevicePolicyManagerInternal.class);
+
+        dpm.setMaximumTimeToLock(admin1, 0);
+
+        reset(getServices().powerManagerInternal);
+        reset(getServices().settings);
+
+        // First add timeout for the profile.
+        dpm.setMaximumTimeToLock(admin1, 10);
+        verifyScreenTimeoutCall(10L, UserHandle.USER_SYSTEM);
+
+        reset(getServices().powerManagerInternal);
+        reset(getServices().settings);
+
+        // Add separate challenge
+        when(getServices().lockPatternUtils
+                .isSeparateProfileChallengeEnabled(eq(PROFILE_USER))).thenReturn(true);
+        dpmi.reportSeparateProfileChallengeChanged(PROFILE_USER);
+
+        verifyScreenTimeoutCall(10L, PROFILE_USER);
+        verifyScreenTimeoutCall(Long.MAX_VALUE, UserHandle.USER_SYSTEM);
+
+        reset(getServices().powerManagerInternal);
+        reset(getServices().settings);
+
+        // Remove the timeout.
+        dpm.setMaximumTimeToLock(admin1, 0);
+        verifyScreenTimeoutCall(Long.MAX_VALUE, PROFILE_USER);
+        verifyScreenTimeoutCall(null , UserHandle.USER_SYSTEM);
+
+        reset(getServices().powerManagerInternal);
+        reset(getServices().settings);
+
+        // Add it back.
+        dpm.setMaximumTimeToLock(admin1, 10);
+        verifyScreenTimeoutCall(10L, PROFILE_USER);
+        verifyScreenTimeoutCall(null, UserHandle.USER_SYSTEM);
+
+        reset(getServices().powerManagerInternal);
+        reset(getServices().settings);
+
+        // Remove separate challenge.
+        reset(getServices().lockPatternUtils);
+        when(getServices().lockPatternUtils
+                .isSeparateProfileChallengeEnabled(eq(PROFILE_USER))).thenReturn(false);
+        dpmi.reportSeparateProfileChallengeChanged(PROFILE_USER);
+
+        verifyScreenTimeoutCall(Long.MAX_VALUE, PROFILE_USER);
+        verifyScreenTimeoutCall(10L , UserHandle.USER_SYSTEM);
+
+        reset(getServices().powerManagerInternal);
+        reset(getServices().settings);
+
+        // Remove the timeout.
+        dpm.setMaximumTimeToLock(admin1, 0);
+        verifyScreenTimeoutCall(null, PROFILE_USER);
+        verifyScreenTimeoutCall(Long.MAX_VALUE, UserHandle.USER_SYSTEM);
     }
 
     public void testSetRequiredStrongAuthTimeout_DeviceOwner() throws Exception {
@@ -2365,15 +2434,17 @@
                 () -> dpm.setRequiredStrongAuthTimeout(admin1, -ONE_MINUTE));
     }
 
-    private void verifyScreenTimeoutCall(Integer expectedTimeout,
-            boolean shouldStayOnWhilePluggedInBeCleared) {
+    private void verifyScreenTimeoutCall(Long expectedTimeout, int userId) {
         if (expectedTimeout == null) {
             verify(getServices().powerManagerInternal, times(0))
-                    .setMaximumScreenOffTimeoutFromDeviceAdmin(anyInt());
+                    .setMaximumScreenOffTimeoutFromDeviceAdmin(eq(userId), anyLong());
         } else {
             verify(getServices().powerManagerInternal, times(1))
-                    .setMaximumScreenOffTimeoutFromDeviceAdmin(eq(expectedTimeout));
+                    .setMaximumScreenOffTimeoutFromDeviceAdmin(eq(userId), eq(expectedTimeout));
         }
+    }
+
+    private void verifyStayOnWhilePluggedCleared(boolean cleared) {
         // TODO Verify calls to settingsGlobalPutInt.  Tried but somehow mockito threw
         // UnfinishedVerificationException.
     }
diff --git a/services/tests/servicestests/src/com/android/server/devicepolicy/MockSystemServices.java b/services/tests/servicestests/src/com/android/server/devicepolicy/MockSystemServices.java
index 4ee5ba6..d7df1da 100644
--- a/services/tests/servicestests/src/com/android/server/devicepolicy/MockSystemServices.java
+++ b/services/tests/servicestests/src/com/android/server/devicepolicy/MockSystemServices.java
@@ -179,6 +179,13 @@
                     return getUserInfo(userId1);
                 }
         );
+        when(userManager.getProfileParent(anyInt())).thenAnswer(
+                invocation -> {
+                    final int userId1 = (int) invocation.getArguments()[0];
+                    final UserInfo ui = getUserInfo(userId1);
+                    return ui == null ? null : getUserInfo(ui.profileGroupId);
+                }
+        );
         when(userManager.getProfiles(anyInt())).thenAnswer(
                 invocation -> {
                     final int userId12 = (int) invocation.getArguments()[0];