diff --git a/example/src/androidTest/java/org/wordpress/android/fluxc/JWTTokenTests.kt b/example/src/androidTest/java/org/wordpress/android/fluxc/JWTTokenTests.kt index 3ff1e1f91d..162592a8cf 100644 --- a/example/src/androidTest/java/org/wordpress/android/fluxc/JWTTokenTests.kt +++ b/example/src/androidTest/java/org/wordpress/android/fluxc/JWTTokenTests.kt @@ -7,19 +7,19 @@ import org.wordpress.android.fluxc.model.JWTToken class JWTTokenTests { @Test - fun given_a_valid_token__when_takeIfValid_is_called__then_return_it() { + fun given_a_valid_token__when_validateExpiryDate_is_called__then_return_it() { val token = generateToken(expired = false) - val result = token.takeIfValid() + val result = token.validateExpiryDate() Assert.assertNotNull(result) } @Test - fun given_an_expired_token__when_takeIfValid_is_called__then_return_null() { + fun given_an_expired_token__when_validateExpiryDate_is_called__then_return_null() { val token = generateToken(expired = true) - val result = token.takeIfValid() + val result = token.validateExpiryDate() Assert.assertNull(result) } diff --git a/fluxc/src/main/java/org/wordpress/android/fluxc/model/JWTToken.kt b/fluxc/src/main/java/org/wordpress/android/fluxc/model/JWTToken.kt index 168fc0bb36..5b1ba89750 100644 --- a/fluxc/src/main/java/org/wordpress/android/fluxc/model/JWTToken.kt +++ b/fluxc/src/main/java/org/wordpress/android/fluxc/model/JWTToken.kt @@ -1,7 +1,6 @@ package org.wordpress.android.fluxc.model import android.util.Base64 -import org.json.JSONException import org.json.JSONObject @JvmInline @@ -11,21 +10,26 @@ value class JWTToken( /** * Returns the token if it is still valid, or null if it is expired. */ - @Suppress("SwallowedException", "MagicNumber") - fun takeIfValid(): JWTToken? { - fun JSONObject.getLongOrNull(name: String) = try { - this.getLong(name) - } catch (e: JSONException) { - null - } + @Suppress("MagicNumber") + fun validateExpiryDate(): JWTToken? { + fun JSONObject.getLongOrNull(name: String) = this.optLong(name, Long.MAX_VALUE).takeIf { it != Long.MAX_VALUE } - val payload = this.value.split(".")[1] - val claimsJson = String(Base64.decode(payload, Base64.DEFAULT)) - val claims = JSONObject(claimsJson) + val payloadJson = getPayloadJson() + val expiration = payloadJson.getLongOrNull("exp") + ?: payloadJson.getLongOrNull("expires") + ?: return null - val expiration = claims.getLongOrNull("exp") ?: claims.getLongOrNull("expires") ?: return null val now = System.currentTimeMillis() / 1000 return if (expiration > now) this else null } + + fun getPayloadItem(key: String): String? { + return getPayloadJson().optString(key) + } + + private fun getPayloadJson(): JSONObject { + val payloadEncoded = this.value.split(".")[1] + return JSONObject(String(Base64.decode(payloadEncoded, Base64.DEFAULT))) + } } diff --git a/fluxc/src/main/java/org/wordpress/android/fluxc/store/jetpackai/JetpackAIStore.kt b/fluxc/src/main/java/org/wordpress/android/fluxc/store/jetpackai/JetpackAIStore.kt index ed8452a490..82ae4344d3 100644 --- a/fluxc/src/main/java/org/wordpress/android/fluxc/store/jetpackai/JetpackAIStore.kt +++ b/fluxc/src/main/java/org/wordpress/android/fluxc/store/jetpackai/JetpackAIStore.kt @@ -76,21 +76,22 @@ class JetpackAIStore @Inject constructor( caller = this, loggedMessage = "fetch Jetpack AI completions" ) { - val token = token?.takeIfValid() ?: fetchJetpackAIJWTToken(site).let { tokenResponse -> - when (tokenResponse) { - is Error -> { - return@withDefaultContext JetpackAICompletionsResponse.Error( - type = AUTH_ERROR, - message = tokenResponse.message, - ) - } + val token = token?.validateExpiryDate()?.validateBlogId(site.siteId) + ?: fetchJetpackAIJWTToken(site).let { tokenResponse -> + when (tokenResponse) { + is Error -> { + return@withDefaultContext JetpackAICompletionsResponse.Error( + type = AUTH_ERROR, + message = tokenResponse.message, + ) + } - is Success -> { - token = tokenResponse.token - tokenResponse.token + is Success -> { + token = tokenResponse.token + tokenResponse.token + } } } - } val result = jetpackAIRestClient.fetchJetpackAITextCompletion(token, prompt, feature) @@ -114,4 +115,7 @@ class JetpackAIStore @Inject constructor( ) { jetpackAIRestClient.fetchJetpackAIJWTToken(site) } + + private fun JWTToken.validateBlogId(blogId: Long): JWTToken? = + if (getPayloadItem("blog_id")?.toLong() == blogId) this else null }