Merge "Profile lock timeout."
diff --git a/core/java/android/app/admin/DevicePolicyManager.java b/core/java/android/app/admin/DevicePolicyManager.java
index 9816297..70e1a96 100644
--- a/core/java/android/app/admin/DevicePolicyManager.java
+++ b/core/java/android/app/admin/DevicePolicyManager.java
@@ -3213,23 +3213,6 @@
     }
 
     /**
-     * Returns maximum time to lock that applied by all profiles in this user. We do this because we
-     * do not have a separate timeout to lock for work challenge only.
-     *
-     * @hide
-     */
-    public long getMaximumTimeToLockForUserAndProfiles(int userHandle) {
-        if (mService != null) {
-            try {
-                return mService.getMaximumTimeToLockForUserAndProfiles(userHandle);
-            } catch (RemoteException e) {
-                throw e.rethrowFromSystemServer();
-            }
-        }
-        return 0;
-    }
-
-    /**
      * Called by a device/profile owner to set the timeout after which unlocking with secondary, non
      * strong auth (e.g. fingerprint, trust agents) times out, i.e. the user has to use a strong
      * authentication method like password, pin or pattern.
diff --git a/core/java/android/app/admin/DevicePolicyManagerInternal.java b/core/java/android/app/admin/DevicePolicyManagerInternal.java
index 05f6c2a..b692ffd9 100644
--- a/core/java/android/app/admin/DevicePolicyManagerInternal.java
+++ b/core/java/android/app/admin/DevicePolicyManagerInternal.java
@@ -16,6 +16,7 @@
 
 package android.app.admin;
 
+import android.annotation.UserIdInt;
 import android.content.Intent;
 
 import java.util.List;
@@ -115,4 +116,11 @@
      * device owner to be affiliated with.
      */
     public abstract boolean isUserAffiliatedWithDevice(int userId);
+
+    /**
+     * Reports that a profile has changed to use a unified or separate credential.
+     *
+     * @param userId User ID of the profile.
+     */
+    public abstract void reportSeparateProfileChallengeChanged(@UserIdInt int userId);
 }
diff --git a/core/java/android/app/admin/IDevicePolicyManager.aidl b/core/java/android/app/admin/IDevicePolicyManager.aidl
index 5b02c22..7cf19ee 100644
--- a/core/java/android/app/admin/IDevicePolicyManager.aidl
+++ b/core/java/android/app/admin/IDevicePolicyManager.aidl
@@ -95,7 +95,6 @@
 
     void setMaximumTimeToLock(in ComponentName who, long timeMs, boolean parent);
     long getMaximumTimeToLock(in ComponentName who, int userHandle, boolean parent);
-    long getMaximumTimeToLockForUserAndProfiles(int userHandle);
 
     void setRequiredStrongAuthTimeout(in ComponentName who, long timeMs, boolean parent);
     long getRequiredStrongAuthTimeout(in ComponentName who, int userId, boolean parent);
diff --git a/core/java/android/os/PowerManagerInternal.java b/core/java/android/os/PowerManagerInternal.java
index 77ac2651..3ef0961 100644
--- a/core/java/android/os/PowerManagerInternal.java
+++ b/core/java/android/os/PowerManagerInternal.java
@@ -110,7 +110,7 @@
      *
      * This method must only be called by the device administration policy manager.
      */
-    public abstract void setMaximumScreenOffTimeoutFromDeviceAdmin(int timeMs);
+    public abstract void setMaximumScreenOffTimeoutFromDeviceAdmin(int userId, long timeMs);
 
     /**
      * Used by the dream manager to override certain properties while dozing.
diff --git a/packages/SettingsLib/src/com/android/settingslib/RestrictedLockUtils.java b/packages/SettingsLib/src/com/android/settingslib/RestrictedLockUtils.java
index c3a36e9..fce5dd9 100644
--- a/packages/SettingsLib/src/com/android/settingslib/RestrictedLockUtils.java
+++ b/packages/SettingsLib/src/com/android/settingslib/RestrictedLockUtils.java
@@ -432,53 +432,9 @@
      * the admin component will be set to {@code null} and userId to {@link UserHandle#USER_NULL}
      */
     public static EnforcedAdmin checkIfMaximumTimeToLockIsSet(Context context) {
-        final DevicePolicyManager dpm = (DevicePolicyManager) context.getSystemService(
-                Context.DEVICE_POLICY_SERVICE);
-        if (dpm == null) {
-            return null;
-        }
-        EnforcedAdmin enforcedAdmin = null;
-        final int userId = UserHandle.myUserId();
-        final UserManager um = UserManager.get(context);
-        final List<UserInfo> profiles = um.getProfiles(userId);
-        final int profilesSize = profiles.size();
-        // As we do not have a separate screen lock timeout settings for work challenge,
-        // we need to combine all profiles maximum time to lock even work challenge is
-        // enabled.
-        for (int i = 0; i < profilesSize; i++) {
-            final UserInfo userInfo = profiles.get(i);
-            final List<ComponentName> admins = dpm.getActiveAdminsAsUser(userInfo.id);
-            if (admins == null) {
-                continue;
-            }
-            for (ComponentName admin : admins) {
-                if (dpm.getMaximumTimeToLock(admin, userInfo.id) > 0) {
-                    if (enforcedAdmin == null) {
-                        enforcedAdmin = new EnforcedAdmin(admin, userInfo.id);
-                    } else {
-                        return EnforcedAdmin.MULTIPLE_ENFORCED_ADMIN;
-                    }
-                    // This same admins could have set policies both on the managed profile
-                    // and on the parent. So, if the admin has set the policy on the
-                    // managed profile here, we don't need to further check if that admin
-                    // has set policy on the parent admin.
-                    continue;
-                }
-                if (userInfo.isManagedProfile()) {
-                    // If userInfo.id is a managed profile, we also need to look at
-                    // the policies set on the parent.
-                    DevicePolicyManager parentDpm = sProxy.getParentProfileInstance(dpm, userInfo);
-                    if (parentDpm.getMaximumTimeToLock(admin, userInfo.id) > 0) {
-                        if (enforcedAdmin == null) {
-                            enforcedAdmin = new EnforcedAdmin(admin, userInfo.id);
-                        } else {
-                            return EnforcedAdmin.MULTIPLE_ENFORCED_ADMIN;
-                        }
-                    }
-                }
-            }
-        }
-        return enforcedAdmin;
+        return checkForLockSetting(context, UserHandle.myUserId(),
+                (DevicePolicyManager dpm, ComponentName admin, @UserIdInt int userId) ->
+                        dpm.getMaximumTimeToLock(admin, userId) > 0);
     }
 
     private interface LockSettingCheck {
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
index c4d9cf5..91ae448 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
@@ -872,7 +872,7 @@
 
         // From DevicePolicyAdmin
         final long policyTimeout = mLockPatternUtils.getDevicePolicyManager()
-                .getMaximumTimeToLockForUserAndProfiles(userId);
+                .getMaximumTimeToLock(null, userId);
 
         long timeout;
 
diff --git a/services/core/java/com/android/server/locksettings/LockSettingsService.java b/services/core/java/com/android/server/locksettings/LockSettingsService.java
index d2baa56..eef4d9b 100644
--- a/services/core/java/com/android/server/locksettings/LockSettingsService.java
+++ b/services/core/java/com/android/server/locksettings/LockSettingsService.java
@@ -38,6 +38,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;
@@ -90,6 +91,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;
@@ -100,6 +102,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.LockSettingsStorage.PersistentData;
@@ -891,14 +894,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
@@ -1234,9 +1249,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,
@@ -2442,9 +2458,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 387818b..e5351b4 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
@@ -4674,56 +4674,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
@@ -4734,50 +4734,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;
             }
         }
@@ -9451,7 +9422,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;
                 }
             }
@@ -9913,6 +9884,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 4779474..60783db 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 8cb0459..4232c44 100644
--- a/services/tests/servicestests/src/com/android/server/devicepolicy/MockSystemServices.java
+++ b/services/tests/servicestests/src/com/android/server/devicepolicy/MockSystemServices.java
@@ -181,6 +181,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];