Add InMemoryLruCache ChallengeChecker implementation

PiperOrigin-RevId: 780246310

Author: suzannajiwani

src/main/kotlin/challengecheckers/InMemoryLruCache.kt

@@ -0,0 +1,29 @@
+package com.android.keyattestation.verifier.challengecheckers
+
+import com.android.keyattestation.verifier.ChallengeChecker
+import com.google.protobuf.ByteString
+
+/**
+ * A [ChallengeChecker] which checks for replay of challenges via an in-memory LRU cache which holds
+ * up to `maxCacheSize` challenges. Challenges are considered invalid if they are already present in
+ * the cache, which prevents replay (reuse of challenges). Checking a challenge will affect the
+ * ordering of the cache, making it more-recently-used.
+ *
+ * @property maxCacheSize the maximum number of challenges to cache
+ */
+class InMemoryLruCache(private val maxCacheSize: Int) : ChallengeChecker {
+  // Use a LinkedHashMap instead of LinkedHashSet even though we don't care about the values since
+  // it can order entries by access-order. Use default initial capacity and load factor.
+  private val cache: LinkedHashMap<ByteString, Int> =
+    object : LinkedHashMap<ByteString, Int>(16, 0.75f, true) {
+
+      // Used to query whether the oldest entry should be removed from the cache.
+      override fun removeEldestEntry(eldest: MutableMap.MutableEntry<ByteString, Int>) =
+        size > maxCacheSize
+    }
+
+  override fun checkChallenge(challenge: ByteString): Boolean {
+    val previousValue = cache.putIfAbsent(challenge, 1)
+    return previousValue == null
+  }
+}

src/test/kotlin/challengecheckers/InMemoryLruCacheTest.kt

@@ -0,0 +1,78 @@
+package com.android.keyattestation.verifier.challengecheckers
+
+import com.google.common.truth.Truth.assertThat
+import com.google.protobuf.ByteString
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@RunWith(JUnit4::class)
+class InMemoryLruCacheTest {
+
+  @Test
+  fun checkChallenge_firstChallenge_returnsTrue() {
+    val challengeChecker = InMemoryLruCache(1)
+    assertThat(challengeChecker.checkChallenge(ByteString.copyFromUtf8("challenge"))).isTrue()
+  }
+
+  @Test
+  fun checkChallenge_partialCacheCheckNewChallenge_returnsTrue() {
+    val challengeChecker = InMemoryLruCache(10)
+    for (i in 1..9) {
+      assertThat(challengeChecker.checkChallenge(ByteString.copyFromUtf8("challenge$i"))).isTrue()
+    }
+
+    assertThat(challengeChecker.checkChallenge(ByteString.copyFromUtf8("foo"))).isTrue()
+  }
+
+  @Test
+  fun checkChallenge_fullCacheCheckExistingChallenge_returnsFalse() {
+    val challengeChecker = InMemoryLruCache(10)
+    for (i in 1..10) {
+      assertThat(challengeChecker.checkChallenge(ByteString.copyFromUtf8("challenge$i"))).isTrue()
+    }
+
+    assertThat(challengeChecker.checkChallenge(ByteString.copyFromUtf8("challenge1"))).isFalse()
+  }
+
+  @Test
+  fun checkChallenge_overflowCacheCheckOldestChallenge_returnsTrue() {
+    val challengeChecker = InMemoryLruCache(10)
+
+    // Fill cache with 10 challenges and overflow with the 11th challenge.
+    for (i in 1..11) {
+      assertThat(challengeChecker.checkChallenge(ByteString.copyFromUtf8("challenge$i"))).isTrue()
+    }
+
+    assertThat(challengeChecker.checkChallenge(ByteString.copyFromUtf8("challenge1"))).isTrue()
+  }
+
+  @Test
+  fun checkChallenge_overflowCacheCheckNewerChallenge_returnsFalse() {
+    val challengeChecker = InMemoryLruCache(10)
+    for (i in 1..11) {
+      assertThat(challengeChecker.checkChallenge(ByteString.copyFromUtf8("challenge$i"))).isTrue()
+    }
+
+    assertThat(challengeChecker.checkChallenge(ByteString.copyFromUtf8("challenge2"))).isFalse()
+  }
+
+  @Test
+  fun checkChallenge_checkingChallenge_affectsCacheOrder() {
+    // fill cache
+    val challengeChecker = InMemoryLruCache(3)
+    for (i in 1..3) {
+      assertThat(challengeChecker.checkChallenge(ByteString.copyFromUtf8("challenge$i"))).isTrue()
+    }
+
+    // check oldest challenge
+    assertThat(challengeChecker.checkChallenge(ByteString.copyFromUtf8("challenge1"))).isFalse()
+
+    // add new challenge to overflow cache + kick out least-recently-used challenge
+    assertThat(challengeChecker.checkChallenge(ByteString.copyFromUtf8("challenge4"))).isTrue()
+
+    // check that challenge1 is still in the cache + challenge2 is kicked out
+    assertThat(challengeChecker.checkChallenge(ByteString.copyFromUtf8("challenge1"))).isFalse()
+    assertThat(challengeChecker.checkChallenge(ByteString.copyFromUtf8("challenge2"))).isTrue()
+  }
+}