Add unit test coverage for unlock attempts in TrustTests

The downstream effects of an unlock attempt using a pin/password were
not previously tested within the TrustTests. Add these tests in
preparation for further changes to pin/password unlock reporting to
TrustManagerService.

Add options within ScreenLockRule and TrustAgentRule to allow for test
cases to start with strong auth being required, and the device locked.

BUG: b/323086607
Test: `atest TrustTests`
Flag: TEST_ONLY
(cherry picked from https://googleplex-android-review.googlesource.com/q/commit:45ca5feb1120c6300cec11352eb9f21d37697eca)
Merged-In: I1c8eaacc9c20284ef2783bcd65c14f68b0b444e8
Change-Id: I1c8eaacc9c20284ef2783bcd65c14f68b0b444e8

5 files changed, 362 insertions, 15 deletions

diff --git a/tests/TrustTests/AndroidManifest.xml b/tests/TrustTests/AndroidManifest.xml
index 30cf345db34d..2f6c0dd14adc 100644
--- a/tests/TrustTests/AndroidManifest.xml
+++ b/tests/TrustTests/AndroidManifest.xml
@@ -78,6 +78,7 @@
                 <action android:name="android.service.trust.TrustAgentService" />
             </intent-filter>
         </service>
+
         <service
             android:name=".IsActiveUnlockRunningTrustAgent"
             android:exported="true"
@@ -88,6 +89,16 @@
             </intent-filter>
         </service>
 
+        <service
+            android:name=".UnlockAttemptTrustAgent"
+            android:exported="true"
+            android:label="Test Agent"
+            android:permission="android.permission.BIND_TRUST_AGENT">
+            <intent-filter>
+                <action android:name="android.service.trust.TrustAgentService" />
+            </intent-filter>
+        </service>
+
     </application>
 
     <!--  self-instrumenting test package. -->
diff --git a/tests/TrustTests/src/android/trust/test/UnlockAttemptTest.kt b/tests/TrustTests/src/android/trust/test/UnlockAttemptTest.kt
new file mode 100644
index 000000000000..2c9361df63fd
--- /dev/null
+++ b/tests/TrustTests/src/android/trust/test/UnlockAttemptTest.kt
@@ -0,0 +1,227 @@
+/*
+ * 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 android.trust.test
+
+import android.app.trust.TrustManager
+import android.content.Context
+import android.trust.BaseTrustAgentService
+import android.trust.TrustTestActivity
+import android.trust.test.lib.LockStateTrackingRule
+import android.trust.test.lib.ScreenLockRule
+import android.trust.test.lib.TestTrustListener
+import android.trust.test.lib.TrustAgentRule
+import android.util.Log
+import androidx.test.core.app.ApplicationProvider.getApplicationContext
+import androidx.test.ext.junit.rules.ActivityScenarioRule
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.google.common.truth.Truth.assertThat
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.rules.RuleChain
+import org.junit.runner.RunWith
+
+/**
+ * Test for the impacts of reporting unlock attempts.
+ *
+ * atest TrustTests:UnlockAttemptTest
+ */
+@RunWith(AndroidJUnit4::class)
+class UnlockAttemptTest {
+    private val context = getApplicationContext<Context>()
+    private val trustManager = context.getSystemService(TrustManager::class.java) as TrustManager
+    private val userId = context.userId
+    private val activityScenarioRule = ActivityScenarioRule(TrustTestActivity::class.java)
+    private val screenLockRule = ScreenLockRule(requireStrongAuth = true)
+    private val lockStateTrackingRule = LockStateTrackingRule()
+    private val trustAgentRule =
+        TrustAgentRule<UnlockAttemptTrustAgent>(startUnlocked = false, startEnabled = false)
+
+    private val trustListener = UnlockAttemptTrustListener()
+    private val agent get() = trustAgentRule.agent
+
+    @get:Rule
+    val rule: RuleChain =
+        RuleChain.outerRule(activityScenarioRule)
+            .around(screenLockRule)
+            .around(lockStateTrackingRule)
+            .around(trustAgentRule)
+
+    @Before
+    fun setUp() {
+        trustManager.registerTrustListener(trustListener)
+    }
+
+    @Test
+    fun successfulUnlockAttempt_allowsTrustAgentToStart() =
+        runUnlockAttemptTest(enableAndVerifyTrustAgent = false, managingTrust = false) {
+            trustAgentRule.enableTrustAgent()
+
+            triggerSuccessfulUnlock()
+
+            trustAgentRule.verifyAgentIsRunning(MAX_WAIT_FOR_ENABLED_TRUST_AGENT_TO_START)
+        }
+
+    @Test
+    fun successfulUnlockAttempt_notifiesTrustAgent() =
+        runUnlockAttemptTest(enableAndVerifyTrustAgent = true, managingTrust = true) {
+            val oldSuccessfulCount = agent.successfulUnlockCallCount
+            val oldFailedCount = agent.failedUnlockCallCount
+
+            triggerSuccessfulUnlock()
+
+            assertThat(agent.successfulUnlockCallCount).isEqualTo(oldSuccessfulCount + 1)
+            assertThat(agent.failedUnlockCallCount).isEqualTo(oldFailedCount)
+        }
+
+    @Test
+    fun successfulUnlockAttempt_notifiesTrustListenerOfManagedTrust() =
+        runUnlockAttemptTest(enableAndVerifyTrustAgent = true, managingTrust = true) {
+            val oldTrustManagedChangedCount = trustListener.onTrustManagedChangedCount[userId] ?: 0
+
+            triggerSuccessfulUnlock()
+
+            assertThat(trustListener.onTrustManagedChangedCount[userId] ?: 0).isEqualTo(
+                oldTrustManagedChangedCount + 1
+            )
+        }
+
+    @Test
+    fun failedUnlockAttempt_doesNotAllowTrustAgentToStart() =
+        runUnlockAttemptTest(enableAndVerifyTrustAgent = false, managingTrust = false) {
+            trustAgentRule.enableTrustAgent()
+
+            triggerFailedUnlock()
+
+            trustAgentRule.ensureAgentIsNotRunning(MAX_WAIT_FOR_ENABLED_TRUST_AGENT_TO_START)
+        }
+
+    @Test
+    fun failedUnlockAttempt_notifiesTrustAgent() =
+        runUnlockAttemptTest(enableAndVerifyTrustAgent = true, managingTrust = true) {
+            val oldSuccessfulCount = agent.successfulUnlockCallCount
+            val oldFailedCount = agent.failedUnlockCallCount
+
+            triggerFailedUnlock()
+
+            assertThat(agent.successfulUnlockCallCount).isEqualTo(oldSuccessfulCount)
+            assertThat(agent.failedUnlockCallCount).isEqualTo(oldFailedCount + 1)
+        }
+
+    @Test
+    fun failedUnlockAttempt_doesNotNotifyTrustListenerOfManagedTrust() =
+        runUnlockAttemptTest(enableAndVerifyTrustAgent = true, managingTrust = true) {
+            val oldTrustManagedChangedCount = trustListener.onTrustManagedChangedCount[userId] ?: 0
+
+            triggerFailedUnlock()
+
+            assertThat(trustListener.onTrustManagedChangedCount[userId] ?: 0).isEqualTo(
+                oldTrustManagedChangedCount
+            )
+        }
+
+    private fun runUnlockAttemptTest(
+        enableAndVerifyTrustAgent: Boolean,
+        managingTrust: Boolean,
+        testBlock: () -> Unit,
+    ) {
+        if (enableAndVerifyTrustAgent) {
+            Log.i(TAG, "Triggering successful unlock")
+            triggerSuccessfulUnlock()
+            Log.i(TAG, "Enabling and waiting for trust agent")
+            trustAgentRule.enableAndVerifyTrustAgentIsRunning(
+                MAX_WAIT_FOR_ENABLED_TRUST_AGENT_TO_START
+            )
+            Log.i(TAG, "Managing trust: $managingTrust")
+            agent.setManagingTrust(managingTrust)
+            await()
+        }
+        testBlock()
+    }
+
+    private fun triggerSuccessfulUnlock() {
+        screenLockRule.successfulScreenLockAttempt()
+        trustAgentRule.reportSuccessfulUnlock()
+        await()
+    }
+
+    private fun triggerFailedUnlock() {
+        screenLockRule.failedScreenLockAttempt()
+        trustAgentRule.reportFailedUnlock()
+        await()
+    }
+
+    companion object {
+        private const val TAG = "UnlockAttemptTest"
+        private fun await(millis: Long = 500) = Thread.sleep(millis)
+        private const val MAX_WAIT_FOR_ENABLED_TRUST_AGENT_TO_START = 10000L
+    }
+}
+
+class UnlockAttemptTrustAgent : BaseTrustAgentService() {
+    var successfulUnlockCallCount: Long = 0
+        private set
+    var failedUnlockCallCount: Long = 0
+        private set
+
+    override fun onUnlockAttempt(successful: Boolean) {
+        super.onUnlockAttempt(successful)
+        if (successful) {
+            successfulUnlockCallCount++
+        } else {
+            failedUnlockCallCount++
+        }
+    }
+}
+
+private class UnlockAttemptTrustListener : TestTrustListener() {
+    var enabledTrustAgentsChangedCount = mutableMapOf<Int, Int>()
+    var onTrustManagedChangedCount = mutableMapOf<Int, Int>()
+
+    override fun onEnabledTrustAgentsChanged(userId: Int) {
+        enabledTrustAgentsChangedCount.compute(userId) { _: Int, curr: Int? ->
+            if (curr == null) 0 else curr + 1
+        }
+    }
+
+    data class TrustChangedParams(
+        val enabled: Boolean,
+        val newlyUnlocked: Boolean,
+        val userId: Int,
+        val flags: Int,
+        val trustGrantedMessages: MutableList<String>?
+    )
+
+    val onTrustChangedCalls = mutableListOf<TrustChangedParams>()
+
+    override fun onTrustChanged(
+        enabled: Boolean,
+        newlyUnlocked: Boolean,
+        userId: Int,
+        flags: Int,
+        trustGrantedMessages: MutableList<String>
+    ) {
+        onTrustChangedCalls += TrustChangedParams(
+            enabled, newlyUnlocked, userId, flags, trustGrantedMessages
+        )
+    }
+
+    override fun onTrustManagedChanged(enabled: Boolean, userId: Int) {
+        onTrustManagedChangedCount.compute(userId) { _: Int, curr: Int? ->
+            if (curr == null) 0 else curr + 1
+        }
+    }
+}
diff --git a/tests/TrustTests/src/android/trust/test/lib/ScreenLockRule.kt b/tests/TrustTests/src/android/trust/test/lib/ScreenLockRule.kt
index f1edca3ff86e..1ccdcc623c5b 100644
--- a/tests/TrustTests/src/android/trust/test/lib/ScreenLockRule.kt
+++ b/tests/TrustTests/src/android/trust/test/lib/ScreenLockRule.kt
@@ -24,6 +24,8 @@ import androidx.test.core.app.ApplicationProvider.getApplicationContext
 import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
 import androidx.test.uiautomator.UiDevice
 import com.android.internal.widget.LockPatternUtils
+import com.android.internal.widget.LockPatternUtils.StrongAuthTracker.STRONG_AUTH_NOT_REQUIRED
+import com.android.internal.widget.LockPatternUtils.StrongAuthTracker.STRONG_AUTH_REQUIRED_AFTER_USER_LOCKDOWN
 import com.android.internal.widget.LockscreenCredential
 import com.google.common.truth.Truth.assertWithMessage
 import org.junit.rules.TestRule
@@ -32,13 +34,18 @@ import org.junit.runners.model.Statement
 
 /**
  * Sets a screen lock on the device for the duration of the test.
+ *
+ * @param requireStrongAuth Whether a strong auth is required at the beginning.
+ * If true, trust agents will not be available until the user verifies their credentials.
  */
-class ScreenLockRule : TestRule {
+class ScreenLockRule(val requireStrongAuth: Boolean = false) : TestRule {
     private val context: Context = getApplicationContext()
+    private val userId = context.userId
     private val uiDevice = UiDevice.getInstance(getInstrumentation())
     private val windowManager = checkNotNull(WindowManagerGlobal.getWindowManagerService())
     private val lockPatternUtils = LockPatternUtils(context)
     private var instantLockSavedValue = false
+    private var strongAuthSavedValue: Int = 0
 
     override fun apply(base: Statement, description: Description) = object : Statement() {
         override fun evaluate() {
@@ -46,10 +53,12 @@ class ScreenLockRule : TestRule {
             dismissKeyguard()
             setScreenLock()
             setLockOnPowerButton()
+            configureStrongAuthState()
 
             try {
                 base.evaluate()
             } finally {
+                restoreStrongAuthState()
                 removeScreenLock()
                 revertLockOnPowerButton()
                 dismissKeyguard()
@@ -57,6 +66,22 @@ class ScreenLockRule : TestRule {
         }
     }
 
+    private fun configureStrongAuthState() {
+        strongAuthSavedValue = lockPatternUtils.getStrongAuthForUser(userId)
+        if (requireStrongAuth) {
+            Log.d(TAG, "Triggering strong auth due to simulated lockdown")
+            lockPatternUtils.requireStrongAuth(STRONG_AUTH_REQUIRED_AFTER_USER_LOCKDOWN, userId)
+            wait("strong auth required after lockdown") {
+                lockPatternUtils.getStrongAuthForUser(userId) ==
+                        STRONG_AUTH_REQUIRED_AFTER_USER_LOCKDOWN
+            }
+        }
+    }
+
+    private fun restoreStrongAuthState() {
+        lockPatternUtils.requireStrongAuth(strongAuthSavedValue, userId)
+    }
+
     private fun verifyNoScreenLockAlreadySet() {
         assertWithMessage("Screen Lock must not already be set on device")
                 .that(lockPatternUtils.isSecure(context.userId))
@@ -82,6 +107,22 @@ class ScreenLockRule : TestRule {
         }
     }
 
+    fun successfulScreenLockAttempt() {
+        lockPatternUtils.verifyCredential(LockscreenCredential.createPin(PIN), context.userId, 0)
+        lockPatternUtils.userPresent(context.userId)
+        wait("strong auth not required") {
+            lockPatternUtils.getStrongAuthForUser(context.userId) == STRONG_AUTH_NOT_REQUIRED
+        }
+    }
+
+    fun failedScreenLockAttempt() {
+        lockPatternUtils.verifyCredential(
+            LockscreenCredential.createPin(WRONG_PIN),
+            context.userId,
+            0
+        )
+    }
+
     private fun setScreenLock() {
         lockPatternUtils.setLockCredential(
                 LockscreenCredential.createPin(PIN),
@@ -121,5 +162,6 @@ class ScreenLockRule : TestRule {
     companion object {
         private const val TAG = "ScreenLockRule"
         private const val PIN = "0000"
+        private const val WRONG_PIN = "0001"
     }
 }
diff --git a/tests/TrustTests/src/android/trust/test/lib/TrustAgentRule.kt b/tests/TrustTests/src/android/trust/test/lib/TrustAgentRule.kt
index 18bc029b6845..404c6d968b3a 100644
--- a/tests/TrustTests/src/android/trust/test/lib/TrustAgentRule.kt
+++ b/tests/TrustTests/src/android/trust/test/lib/TrustAgentRule.kt
@@ -20,14 +20,15 @@ import android.app.trust.TrustManager
 import android.content.ComponentName
 import android.content.Context
 import android.trust.BaseTrustAgentService
+import android.trust.test.lib.TrustAgentRule.Companion.invoke
 import android.util.Log
 import androidx.test.core.app.ApplicationProvider.getApplicationContext
 import com.android.internal.widget.LockPatternUtils
 import com.google.common.truth.Truth.assertWithMessage
+import kotlin.reflect.KClass
 import org.junit.rules.TestRule
 import org.junit.runner.Description
 import org.junit.runners.model.Statement
-import kotlin.reflect.KClass
 
 /**
  * Enables a trust agent and causes the system service to bind to it.
@@ -37,7 +38,9 @@ import kotlin.reflect.KClass
  * @constructor Creates the rule. Do not use; instead, use [invoke].
  */
 class TrustAgentRule<T : BaseTrustAgentService>(
-    private val serviceClass: KClass<T>
+    private val serviceClass: KClass<T>,
+    private val startUnlocked: Boolean,
+    private val startEnabled: Boolean,
 ) : TestRule {
     private val context: Context = getApplicationContext()
     private val trustManager = context.getSystemService(TrustManager::class.java) as TrustManager
@@ -48,11 +51,18 @@ class TrustAgentRule<T : BaseTrustAgentService>(
     override fun apply(base: Statement, description: Description) = object : Statement() {
         override fun evaluate() {
             verifyTrustServiceRunning()
-            unlockDeviceWithCredential()
-            enableTrustAgent()
+            if (startUnlocked) {
+                reportSuccessfulUnlock()
+            } else {
+                Log.i(TAG, "Trust manager not starting in unlocked state")
+            }
 
             try {
-                verifyAgentIsRunning()
+                if (startEnabled) {
+                    enableAndVerifyTrustAgentIsRunning()
+                } else {
+                    Log.i(TAG, "Trust agent ${serviceClass.simpleName} not enabled")
+                }
                 base.evaluate()
             } finally {
                 disableTrustAgent()
@@ -64,12 +74,22 @@ class TrustAgentRule<T : BaseTrustAgentService>(
         assertWithMessage("Trust service is not running").that(trustManager).isNotNull()
     }
 
-    private fun unlockDeviceWithCredential() {
-        Log.d(TAG, "Unlocking device with credential")
+    fun reportSuccessfulUnlock() {
+        Log.i(TAG, "Reporting successful unlock")
         trustManager.reportUnlockAttempt(true, context.userId)
     }
 
-    private fun enableTrustAgent() {
+    fun reportFailedUnlock() {
+        Log.i(TAG, "Reporting failed unlock")
+        trustManager.reportUnlockAttempt(false, context.userId)
+    }
+
+    fun enableAndVerifyTrustAgentIsRunning(maxWait: Long = 30000L) {
+        enableTrustAgent()
+        verifyAgentIsRunning(maxWait)
+    }
+
+    fun enableTrustAgent() {
         val componentName = ComponentName(context, serviceClass.java)
         val userId = context.userId
         Log.i(TAG, "Enabling trust agent ${componentName.flattenToString()} for user $userId")
@@ -79,12 +99,18 @@ class TrustAgentRule<T : BaseTrustAgentService>(
         lockPatternUtils.setEnabledTrustAgents(agents, userId)
     }
 
-    private fun verifyAgentIsRunning() {
-        wait("${serviceClass.simpleName} to be running") {
+    fun verifyAgentIsRunning(maxWait: Long = 30000L) {
+        wait("${serviceClass.simpleName} to be running", maxWait) {
             BaseTrustAgentService.instance(serviceClass) != null
         }
     }
 
+    fun ensureAgentIsNotRunning(window: Long = 30000L) {
+        ensure("${serviceClass.simpleName} is not running", window) {
+            BaseTrustAgentService.instance(serviceClass) == null
+        }
+    }
+
     private fun disableTrustAgent() {
         val componentName = ComponentName(context, serviceClass.java)
         val userId = context.userId
@@ -97,13 +123,23 @@ class TrustAgentRule<T : BaseTrustAgentService>(
 
     companion object {
         /**
-         * Creates a new rule for the specified agent class. Example usage:
+         * Creates a new rule for the specified agent class. Starts with the device unlocked and
+         * the trust agent enabled. Example usage:
          * ```
          *   @get:Rule val rule = TrustAgentRule<MyTestAgent>()
          * ```
+         *
+         * Also supports setting different device lock and trust agent enablement states:
+         * ```
+         *   @get:Rule val rule = TrustAgentRule<MyTestAgent>(startUnlocked = false, startEnabled = false)
+         * ```
          */
-        inline operator fun <reified T : BaseTrustAgentService> invoke() =
-            TrustAgentRule(T::class)
+        inline operator fun <reified T : BaseTrustAgentService> invoke(
+            startUnlocked: Boolean = true,
+            startEnabled: Boolean = true,
+        ) =
+            TrustAgentRule(T::class, startUnlocked, startEnabled)
+
 
         private const val TAG = "TrustAgentRule"
     }
diff --git a/tests/TrustTests/src/android/trust/test/lib/utils.kt b/tests/TrustTests/src/android/trust/test/lib/Utils.kt
index e047202f6740..3b32b47a6160 100644
--- a/tests/TrustTests/src/android/trust/test/lib/utils.kt
+++ b/tests/TrustTests/src/android/trust/test/lib/Utils.kt
@@ -39,7 +39,7 @@ internal fun wait(
 ) {
     var waited = 0L
     var count = 0
-    while (!conditionFunction.invoke(count)) {
+    while (!conditionFunction(count)) {
         assertWithMessage("Condition exceeded maximum wait time of $maxWait ms: $description")
             .that(waited <= maxWait)
             .isTrue()
@@ -49,3 +49,34 @@ internal fun wait(
         Thread.sleep(rate)
     }
 }
+
+/**
+ * Ensures that [conditionFunction] is true with a failed assertion if it is not within [window]
+ * ms.
+ *
+ * The condition function can perform additional logic (for example, logging or attempting to make
+ * the condition become true).
+ *
+ * @param conditionFunction function which takes the attempt count & returns whether the condition
+ *                          is met
+ */
+internal fun ensure(
+    description: String? = null,
+    window: Long = 30000L,
+    rate: Long = 50L,
+    conditionFunction: (count: Int) -> Boolean
+) {
+    var waited = 0L
+    var count = 0
+    while (waited <= window) {
+        assertWithMessage("Condition failed within $window ms: $description").that(
+                conditionFunction(
+                    count
+                )
+            ).isTrue()
+        waited += rate
+        count++
+        Log.i(TAG, "Ensuring $description ($waited/$window) #$count")
+        Thread.sleep(rate)
+    }
+}