Bytter ut thread-safe token cache med en som er coroutine-safe

Co-authored-by: Richard Borge <[email protected]> Co-authored-by: Vetle Hollund <[email protected]> Co-authored-by: Sturle Helland <[email protected]>
navikt · Dec 7, 2023 · 488c07e · 488c07e
1 parent e5606c0
commit 488c07e
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 11 deletions.
diff --git a/ktor-auth-azuread/build.gradle.kts b/ktor-auth-azuread/build.gradle.kts
@@ -1,7 +1,7 @@
 val ktorVersion = "2.3.6"
 
 dependencies {
-    implementation(project(":cache"))
+//    implementation(project(":cache"))
 
     implementation("io.ktor:ktor-client-content-negotiation:$ktorVersion")
     implementation("io.ktor:ktor-client-cio:$ktorVersion")

diff --git a/ktor-auth-azuread/main/no/nav/aap/ktor/client/AzureAdTokenProvider.kt b/ktor-auth-azuread/main/no/nav/aap/ktor/client/AzureAdTokenProvider.kt
@@ -10,7 +10,6 @@ import io.ktor.client.request.*
 import io.ktor.client.statement.*
 import io.ktor.http.*
 import io.ktor.serialization.jackson.*
-import no.nav.aap.cache.Cache
 import org.slf4j.LoggerFactory
 
 private val secureLog = LoggerFactory.getLogger("secureLog")
@@ -32,20 +31,20 @@ class AzureAdTokenProvider(
         "client_id=${config.clientId}&client_secret=${config.clientSecret}&assertion=$accessToken&scope=$scope&grant_type=urn:ietf:params:oauth:grant-type:jwt-bearer&requested_token_use=on_behalf_of"
     }
 
-    private val tokenCache = Cache<String, Token>()
+    private val cache = TokenCache<String>()
 
     private suspend fun getAccessToken(cacheKey: String, body: () -> String): String {
-        val token = tokenCache[cacheKey]
+        val token = cache.get(cacheKey)
             ?: client.post(config.tokenEndpoint) {
                 accept(ContentType.Application.Json)
                 contentType(ContentType.Application.FormUrlEncoded)
                 setBody(body())
             }.also {
-                if(!it.status.isSuccess()) {
+                if (!it.status.isSuccess()) {
                     secureLog.warn("Feilet token-kall {}: {}", it.status.value, it.bodyAsText())
                 }
-            }.body<Token>().also { fetchedToken ->
-                fetchedToken.addToCache(tokenCache, cacheKey)
+            }.body<Token>().also {
+                cache.add(cacheKey, it)
             }
 
         return token.access_token

diff --git a/ktor-auth-azuread/main/no/nav/aap/ktor/client/Token.kt b/ktor-auth-azuread/main/no/nav/aap/ktor/client/Token.kt
@@ -1,16 +1,43 @@
 package no.nav.aap.ktor.client
 
-import no.nav.aap.cache.Cache
+import kotlinx.coroutines.sync.Mutex
+import kotlinx.coroutines.sync.withLock
 import java.time.Instant
 
 internal data class Token(val expires_in: Long, val access_token: String) {
     private val expiry: Instant = Instant.now().plusSeconds(expires_in - LEEWAY_SECONDS)
 
-    internal fun addToCache(cache: Cache<String, Token>, cacheKey: String) {
-        cache.set(cacheKey, this, expiry)
-    }
+    internal fun expired() = Instant.now().isAfter(expiry)
 
     private companion object {
         const val LEEWAY_SECONDS = 60
     }
 }
+
+internal class TokenCache<K> {
+    private val tokens: HashMap<K, Token> = hashMapOf()
+    private val mutex = Mutex()
+
+    internal suspend fun add(key: K, token: Token) {
+        mutex.withLock {
+            tokens[key] = token
+        }
+    }
+
+    internal suspend fun get(key: K): Token? {
+        tokens[key]?.let {
+            if (it.expired()) {
+                rm(key)
+            }
+        }
+        return mutex.withLock {
+            tokens[key]
+        }
+    }
+
+    private suspend fun rm(key: K) {
+        mutex.withLock {
+            tokens.remove(key)
+        }
+    }
+}