From 573677590ef03d440e447442451b5b00d1ea74a6 Mon Sep 17 00:00:00 2001
From: Suzanna Jiwani <suzannaj@google.com>
Date: Tue, 26 Aug 2025 07:10:30 -0700
Subject: [PATCH] Add InMemoryLruCache ChallengeChecker implementation

PiperOrigin-RevId: 799549771
---
 .../challengecheckers/InMemoryLruCache.kt     | 29 +++++++
 .../challengecheckers/InMemoryLruCacheTest.kt | 78 +++++++++++++++++++
 2 files changed, 107 insertions(+)
 create mode 100644 src/main/kotlin/challengecheckers/InMemoryLruCache.kt
 create mode 100644 src/test/kotlin/challengecheckers/InMemoryLruCacheTest.kt

diff --git a/src/main/kotlin/challengecheckers/InMemoryLruCache.kt b/src/main/kotlin/challengecheckers/InMemoryLruCache.kt
new file mode 100644
index 0000000..79064eb
--- /dev/null
+++ b/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
+  }
+}
diff --git a/src/test/kotlin/challengecheckers/InMemoryLruCacheTest.kt b/src/test/kotlin/challengecheckers/InMemoryLruCacheTest.kt
new file mode 100644
index 0000000..379974f
--- /dev/null
+++ b/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()
+  }
+}