Adding tests for temporary and renewable trust

Adding unit tests for granting and revoking temporary and renewable
trust. These unit tests should cover basic cases, and that the phone
should be locked/unlocked when expected. Discovered a bug: revoking
trust now also downgrades from trustable and will no longer listen for a
trustable downgrade. Also modified assertLocked and assertUnlocked in an
attempt to reduce flakiness.

Test: atest TrustTests --iterations
Bug: 221155933
Change-Id: I4e0e213427111dbe25a76ff67b6f36e57e295793
diff --git a/services/core/java/com/android/server/trust/TrustAgentWrapper.java b/services/core/java/com/android/server/trust/TrustAgentWrapper.java
index 20cd8f5..adca216 100644
--- a/services/core/java/com/android/server/trust/TrustAgentWrapper.java
+++ b/services/core/java/com/android/server/trust/TrustAgentWrapper.java
@@ -198,6 +198,8 @@
                     // Fall through.
                 case MSG_REVOKE_TRUST:
                     mTrusted = false;
+                    mTrustable = false;
+                    mWaitingForTrustableDowngrade = false;
                     mDisplayTrustGrantedMessage = false;
                     mMessage = null;
                     mHandler.removeMessages(MSG_TRUST_TIMEOUT);
diff --git a/tests/TrustTests/AndroidManifest.xml b/tests/TrustTests/AndroidManifest.xml
index 68bc1f69..8b4cbfd 100644
--- a/tests/TrustTests/AndroidManifest.xml
+++ b/tests/TrustTests/AndroidManifest.xml
@@ -68,6 +68,16 @@
                 <action android:name="android.service.trust.TrustAgentService" />
             </intent-filter>
         </service>
+
+        <service
+            android:name=".TemporaryAndRenewableTrustAgent"
+            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/GrantAndRevokeTrustTest.kt b/tests/TrustTests/src/android/trust/test/GrantAndRevokeTrustTest.kt
index 790afd3..af7a98c 100644
--- a/tests/TrustTests/src/android/trust/test/GrantAndRevokeTrustTest.kt
+++ b/tests/TrustTests/src/android/trust/test/GrantAndRevokeTrustTest.kt
@@ -60,7 +60,6 @@
     @Test
     fun sleepingDeviceWithoutGrantLocksDevice() {
         uiDevice.sleep()
-        await()
 
         lockStateTrackingRule.assertLocked()
     }
@@ -69,7 +68,6 @@
     fun grantKeepsDeviceUnlocked() {
         trustAgentRule.agent.grantTrust(GRANT_MESSAGE, 10000, 0)
         uiDevice.sleep()
-        await()
 
         lockStateTrackingRule.assertUnlocked()
     }
@@ -80,7 +78,6 @@
         await()
         uiDevice.sleep()
         trustAgentRule.agent.revokeTrust()
-        await()
 
         lockStateTrackingRule.assertLocked()
     }
diff --git a/tests/TrustTests/src/android/trust/test/LockUserTest.kt b/tests/TrustTests/src/android/trust/test/LockUserTest.kt
index 8f200a6..a7dd41ad 100644
--- a/tests/TrustTests/src/android/trust/test/LockUserTest.kt
+++ b/tests/TrustTests/src/android/trust/test/LockUserTest.kt
@@ -24,7 +24,6 @@
 import android.util.Log
 import androidx.test.ext.junit.rules.ActivityScenarioRule
 import androidx.test.ext.junit.runners.AndroidJUnit4
-import com.google.common.truth.Truth.assertThat
 import org.junit.Rule
 import org.junit.Test
 import org.junit.rules.RuleChain
@@ -52,9 +51,8 @@
     fun lockUser_locksTheDevice() {
         Log.i(TAG, "Locking user")
         trustAgentRule.agent.lockUser()
-        await()
 
-        assertThat(lockStateTrackingRule.lockState.locked).isTrue()
+        lockStateTrackingRule.assertLocked()
     }
 
     companion object {
diff --git a/tests/TrustTests/src/android/trust/test/TemporaryAndRenewableTrustTest.kt b/tests/TrustTests/src/android/trust/test/TemporaryAndRenewableTrustTest.kt
new file mode 100644
index 0000000..14c227b
--- /dev/null
+++ b/tests/TrustTests/src/android/trust/test/TemporaryAndRenewableTrustTest.kt
@@ -0,0 +1,124 @@
+/*
+ * Copyright (C) 2022 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.service.trust.TrustAgentService.FLAG_GRANT_TRUST_TEMPORARY_AND_RENEWABLE
+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.TrustAgentRule
+import androidx.test.ext.junit.rules.ActivityScenarioRule
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
+import androidx.test.uiautomator.UiDevice
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.rules.RuleChain
+import org.junit.runner.RunWith
+
+/**
+ * Test for testing revokeTrust & grantTrust for renewable trust.
+ *
+ * atest TrustTests:TemporaryAndRenewableTrustTest
+ */
+@RunWith(AndroidJUnit4::class)
+class TemporaryAndRenewableTrustTest {
+    private val uiDevice = UiDevice.getInstance(getInstrumentation())
+    private val activityScenarioRule = ActivityScenarioRule(TrustTestActivity::class.java)
+    private val lockStateTrackingRule = LockStateTrackingRule()
+    private val trustAgentRule = TrustAgentRule<TemporaryAndRenewableTrustAgent>()
+
+    @get:Rule
+    val rule: RuleChain = RuleChain
+        .outerRule(activityScenarioRule)
+        .around(ScreenLockRule())
+        .around(lockStateTrackingRule)
+        .around(trustAgentRule)
+
+    @Before
+    fun manageTrust() {
+        trustAgentRule.agent.setManagingTrust(true)
+    }
+
+    // This test serves a baseline for Grant tests, verifying that the default behavior of the
+    // device is to lock when put to sleep
+    @Test
+    fun sleepingDeviceWithoutGrantLocksDevice() {
+        uiDevice.sleep()
+
+        lockStateTrackingRule.assertLocked()
+    }
+
+    @Test
+    fun grantTrustLockedDevice_deviceStaysLocked() {
+        uiDevice.sleep()
+        lockStateTrackingRule.assertLocked()
+
+        trustAgentRule.agent.grantTrust(GRANT_MESSAGE, 0, FLAG_GRANT_TRUST_TEMPORARY_AND_RENEWABLE)
+        uiDevice.wakeUp()
+
+        lockStateTrackingRule.assertLocked()
+    }
+
+    @Test
+    fun grantTrustUnlockedDevice_deviceLocksOnScreenOff() {
+        trustAgentRule.agent.grantTrust(GRANT_MESSAGE, 0, FLAG_GRANT_TRUST_TEMPORARY_AND_RENEWABLE)
+        uiDevice.sleep()
+
+        lockStateTrackingRule.assertLocked()
+    }
+
+    @Test
+    fun grantTrustLockedDevice_grantTrustOnLockedDeviceUnlocksDevice() {
+        trustAgentRule.agent.grantTrust(GRANT_MESSAGE, 0, FLAG_GRANT_TRUST_TEMPORARY_AND_RENEWABLE)
+        uiDevice.sleep()
+
+        lockStateTrackingRule.assertLocked()
+
+        trustAgentRule.agent.grantTrust(GRANT_MESSAGE, 0, FLAG_GRANT_TRUST_TEMPORARY_AND_RENEWABLE)
+        uiDevice.wakeUp()
+
+        lockStateTrackingRule.assertUnlocked()
+    }
+
+    @Test
+    fun grantTrustLockedDevice_revokeTrustPreventsSubsequentUnlock() {
+        trustAgentRule.agent.grantTrust(GRANT_MESSAGE, 0, FLAG_GRANT_TRUST_TEMPORARY_AND_RENEWABLE)
+        uiDevice.sleep()
+
+        lockStateTrackingRule.assertLocked()
+
+        trustAgentRule.agent.revokeTrust()
+        await(500)
+        uiDevice.wakeUp()
+        await(500)
+
+        trustAgentRule.agent.grantTrust(GRANT_MESSAGE, 0, FLAG_GRANT_TRUST_TEMPORARY_AND_RENEWABLE)
+
+        lockStateTrackingRule.assertLocked()
+    }
+
+    companion object {
+        private const val TAG = "TemporaryAndRenewableTrustTest"
+        private const val GRANT_MESSAGE = "granted by test"
+        private fun await(millis: Long) = Thread.sleep(millis)
+    }
+}
+
+class TemporaryAndRenewableTrustAgent : BaseTrustAgentService()
diff --git a/tests/TrustTests/src/android/trust/test/lib/LockStateTrackingRule.kt b/tests/TrustTests/src/android/trust/test/lib/LockStateTrackingRule.kt
index 0023af8..834f212 100644
--- a/tests/TrustTests/src/android/trust/test/lib/LockStateTrackingRule.kt
+++ b/tests/TrustTests/src/android/trust/test/lib/LockStateTrackingRule.kt
@@ -52,8 +52,29 @@
         }
     }
 
-    fun assertLocked() = assertThat(lockState.locked).isTrue()
-    fun assertUnlocked() = assertThat(lockState.locked).isFalse()
+    fun assertLocked() {
+        val maxWaits = 50
+        var waitCount = 0
+
+        while ((lockState.locked == false) && waitCount < maxWaits) {
+            Log.i(TAG, "phone still locked, wait 50ms more ($waitCount)")
+            Thread.sleep(50)
+            waitCount++
+        }
+        assertThat(lockState.locked).isTrue()
+    }
+
+    fun assertUnlocked() {
+        val maxWaits = 50
+        var waitCount = 0
+
+        while ((lockState.locked == true) && waitCount < maxWaits) {
+            Log.i(TAG, "phone still unlocked, wait 50ms more ($waitCount)")
+            Thread.sleep(50)
+            waitCount++
+        }
+        assertThat(lockState.locked).isFalse()
+    }
 
     inner class Listener : TrustListener {
         override fun onTrustChanged(