From 488c07ee937d605fb06f415738d2ae9b3a3e0300 Mon Sep 17 00:00:00 2001
From: Robin <robin.tordly@nav.no>
Date: Thu, 7 Dec 2023 13:53:24 +0100
Subject: [PATCH] Bytter ut thread-safe token cache med en som er
 coroutine-safe

Co-authored-by: Richard Borge <richard.borge@nav.no>
Co-authored-by: Vetle Hollund <vetle.hollund@nav.no>
Co-authored-by: Sturle Helland <sturle.helland@nav.no>
---
 ktor-auth-azuread/build.gradle.kts            |  2 +-
 .../aap/ktor/client/AzureAdTokenProvider.kt   | 11 +++---
 .../main/no/nav/aap/ktor/client/Token.kt      | 35 ++++++++++++++++---
 3 files changed, 37 insertions(+), 11 deletions(-)

diff --git a/ktor-auth-azuread/build.gradle.kts b/ktor-auth-azuread/build.gradle.kts
index a7445510..527e8d8f 100644
--- 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
index fff43bc0..d97a8124 100644
--- 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
index 122196cc..04ee174e 100644
--- 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)
+        }
+    }
+}