Add KeyguardStateCallbackInteractor.

Failing to update this callback correctly was resulting in key input being restricted after unlock in some cases, which in certain orders breaks test setup and teardown.

Bug: 278086361
Test: atest com.android.server.wm.flicker.notification
Flag: com.android.systemui.keyguard_wm_state_refactor
Change-Id: Icc96984ca2edde8f75a7dde6dc462e44f1661103

4 files changed, 106 insertions, 1 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java
index 1bc91cac1a84..67625d04b467 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java
@@ -81,6 +81,7 @@ import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.keyguard.domain.interactor.KeyguardDismissInteractor;
 import com.android.systemui.keyguard.domain.interactor.KeyguardEnabledInteractor;
 import com.android.systemui.keyguard.domain.interactor.KeyguardInteractor;
+import com.android.systemui.keyguard.domain.interactor.KeyguardStateCallbackInteractor;
 import com.android.systemui.keyguard.domain.interactor.KeyguardWakeDirectlyToGoneInteractor;
 import com.android.systemui.keyguard.ui.binder.KeyguardSurfaceBehindParamsApplier;
 import com.android.systemui.keyguard.ui.binder.KeyguardSurfaceBehindViewBinder;
@@ -126,6 +127,7 @@ public class KeyguardService extends Service {
     private final Lazy<DeviceEntryInteractor> mDeviceEntryInteractorLazy;
     private final Executor mMainExecutor;
     private final Lazy<KeyguardStateCallbackStartable> mKeyguardStateCallbackStartableLazy;
+    private final KeyguardStateCallbackInteractor mKeyguardStateCallbackInteractor;
 
     private static RemoteAnimationTarget[] wrap(TransitionInfo info, boolean wallpapers,
             SurfaceControl.Transaction t, ArrayMap<SurfaceControl, SurfaceControl> leashMap,
@@ -350,7 +352,8 @@ public class KeyguardService extends Service {
             Lazy<KeyguardStateCallbackStartable> keyguardStateCallbackStartableLazy,
             KeyguardWakeDirectlyToGoneInteractor keyguardWakeDirectlyToGoneInteractor,
             KeyguardDismissInteractor keyguardDismissInteractor,
-            Lazy<DeviceEntryInteractor> deviceEntryInteractorLazy) {
+            Lazy<DeviceEntryInteractor> deviceEntryInteractorLazy,
+            KeyguardStateCallbackInteractor keyguardStateCallbackInteractor) {
         super();
         mKeyguardViewMediator = keyguardViewMediator;
         mKeyguardLifecyclesDispatcher = keyguardLifecyclesDispatcher;
@@ -363,6 +366,7 @@ public class KeyguardService extends Service {
         mSceneInteractorLazy = sceneInteractorLazy;
         mMainExecutor = mainExecutor;
         mKeyguardStateCallbackStartableLazy = keyguardStateCallbackStartableLazy;
+        mKeyguardStateCallbackInteractor = keyguardStateCallbackInteractor;
         mDeviceEntryInteractorLazy = deviceEntryInteractorLazy;
 
         if (KeyguardWmStateRefactor.isEnabled()) {
@@ -455,6 +459,8 @@ public class KeyguardService extends Service {
             checkPermission();
             if (SceneContainerFlag.isEnabled()) {
                 mKeyguardStateCallbackStartableLazy.get().addCallback(callback);
+            } else if (KeyguardWmStateRefactor.isEnabled()) {
+                mKeyguardStateCallbackInteractor.addCallback(callback);
             } else {
                 mKeyguardViewMediator.addStateMonitorCallback(callback);
             }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
index d38c9520eb7a..3b1569d7f79b 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
@@ -2288,6 +2288,10 @@ public class KeyguardViewMediator implements CoreStartable, Dumpable,
     }
 
     private void updateInputRestrictedLocked() {
+        if (KeyguardWmStateRefactor.isEnabled()) {
+            return;
+        }
+
         boolean inputRestricted = isInputRestricted();
         if (mInputRestricted != inputRestricted) {
             mInputRestricted = inputRestricted;
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardStateCallbackInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardStateCallbackInteractor.kt
new file mode 100644
index 000000000000..420fbd4ae48d
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardStateCallbackInteractor.kt
@@ -0,0 +1,93 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.keyguard.domain.interactor
+
+import android.os.DeadObjectException
+import android.os.RemoteException
+import com.android.internal.policy.IKeyguardStateCallback
+import com.android.systemui.CoreStartable
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.dagger.qualifiers.Background
+import com.android.systemui.keyguard.KeyguardWmStateRefactor
+import com.android.systemui.keyguard.shared.model.KeyguardState
+import com.android.systemui.scene.shared.flag.SceneContainerFlag
+import com.android.systemui.user.domain.interactor.SelectedUserInteractor
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.flow.collectLatest
+import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.withContext
+import javax.inject.Inject
+
+/**
+ * Updates KeyguardStateCallbacks provided to KeyguardService with KeyguardTransitionInteractor
+ * state.
+ *
+ * This borrows heavily from [KeyguardStateCallbackStartable], which requires Flexiglass. This class
+ * can be removed after Flexiglass launches.
+ */
+@SysUISingleton
+class KeyguardStateCallbackInteractor
+@Inject
+constructor(
+    @Application private val applicationScope: CoroutineScope,
+    @Background private val backgroundDispatcher: CoroutineDispatcher,
+    private val selectedUserInteractor: SelectedUserInteractor,
+    private val keyguardTransitionInteractor: KeyguardTransitionInteractor,
+) : CoreStartable {
+    private val callbacks = mutableListOf<IKeyguardStateCallback>()
+
+    override fun start() {
+        if (!KeyguardWmStateRefactor.isEnabled || SceneContainerFlag.isEnabled) {
+            return
+        }
+
+        applicationScope.launch {
+            combine(
+                selectedUserInteractor.selectedUser,
+                keyguardTransitionInteractor.currentKeyguardState,
+                ::Pair
+            ).collectLatest { (selectedUser, currentState) ->
+                val iterator = callbacks.iterator()
+                    withContext(backgroundDispatcher) {
+                        while (iterator.hasNext()) {
+                            val callback = iterator.next()
+                            try {
+                                 callback.onShowingStateChanged(
+                                    currentState != KeyguardState.GONE,
+                                    selectedUser
+                                )
+                                callback.onInputRestrictedStateChanged(
+                                    currentState != KeyguardState.GONE)
+                            } catch (e: RemoteException) {
+                                if (e is DeadObjectException) {
+                                    iterator.remove()
+                                }
+                            }
+                        }
+                    }
+                }
+        }
+    }
+
+    fun addCallback(callback: IKeyguardStateCallback) {
+        KeyguardWmStateRefactor.isUnexpectedlyInLegacyMode()
+        callbacks.add(callback)
+    }
+}
\ No newline at end of file
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionCoreStartable.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionCoreStartable.kt
index d9c48fa7e581..25b8fd32e82a 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionCoreStartable.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionCoreStartable.kt
@@ -29,6 +29,7 @@ constructor(
     private val auditLogger: KeyguardTransitionAuditLogger,
     private val bootInteractor: KeyguardTransitionBootInteractor,
     private val statusBarDisableFlagsInteractor: StatusBarDisableFlagsInteractor,
+    private val keyguardStateCallbackInteractor: KeyguardStateCallbackInteractor,
 ) : CoreStartable {
 
     override fun start() {
@@ -55,6 +56,7 @@ constructor(
         auditLogger.start()
         bootInteractor.start()
         statusBarDisableFlagsInteractor.start()
+        keyguardStateCallbackInteractor.start()
     }
 
     companion object {