Implement logic for validating JWT token's expiry date before using it

example/src/androidTest/java/org/wordpress/android/fluxc/JWTTokenTests.kt

@@ -0,0 +1,47 @@
+package org.wordpress.android.fluxc
+
+import android.util.Base64
+import org.junit.Assert
+import org.junit.Test
+import org.wordpress.android.fluxc.model.JWTToken
+
+class JWTTokenTests {
+    @Test
+    fun given_a_valid_token__when_takeIfValid_is_called__then_return_it() {
+        val token = generateToken(expired = false)
+
+        val result = token.takeIfValid()
+
+        Assert.assertNotNull(result)
+    }
+
+    @Test
+    fun given_an_expired_token__when_takeIfValid_is_called__then_return_null() {
+        val token = generateToken(expired = true)
+
+        val result = token.takeIfValid()
+
+        Assert.assertNull(result)
+    }
+
+    private fun generateToken(expired: Boolean): JWTToken {
+        val expirationTime = System.currentTimeMillis() / 1000 + if (expired) -100 else 100
+
+        // Sample token from https://jwt.io/ modifier with an expiration time
+        val header = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9"
+        val payload = Base64.encode(
+            """
+                {
+                    "sub": "1234567890",
+                    "name": "John Doe",
+                    "iat": 1516239022,
+                    "exp": $expirationTime,
+                    "expires": $expirationTime
+                }
+            """.trimIndent().toByteArray(), Base64.DEFAULT
+        ).decodeToString()
+        val signature = "SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c"
+
+        return JWTToken("$header.$payload.$signature")
+    }
+}

fluxc/src/main/java/org/wordpress/android/fluxc/model/JWTToken.kt

@@ -0,0 +1,31 @@
+package org.wordpress.android.fluxc.model
+
+import android.util.Base64
+import org.json.JSONException
+import org.json.JSONObject
+
+@JvmInline
+value class JWTToken(
+    val value: String
+) {
+    /**
+     * 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
+        }
+
+        val payload = this.value.split(".")[1]
+        val claimsJson = String(Base64.decode(payload, Base64.DEFAULT))
+        val claims = JSONObject(claimsJson)
+
+        val expiration = claims.getLongOrNull("exp") ?: claims.getLongOrNull("expires") ?: return null
+        val now = System.currentTimeMillis() / 1000
+
+        return if (expiration > now) this else null
+    }
+}

fluxc/src/main/java/org/wordpress/android/fluxc/network/rest/wpcom/jetpackai/JetpackAIRestClient.kt

@@ -5,6 +5,7 @@ import com.android.volley.RequestQueue
 import com.google.gson.annotations.SerializedName
 import org.wordpress.android.fluxc.Dispatcher
 import org.wordpress.android.fluxc.generated.endpoint.WPCOMV2
+import org.wordpress.android.fluxc.model.JWTToken
 import org.wordpress.android.fluxc.model.SiteModel
 import org.wordpress.android.fluxc.network.BaseRequest.GenericErrorType
 import org.wordpress.android.fluxc.network.UserAgent
@@ -43,7 +44,7 @@ class JetpackAIRestClient @Inject constructor(
         )
 
         return when (response) {
-            is Response.Success -> JetpackAIJWTTokenResponse.Success(response.data.token)
+            is Response.Success -> JetpackAIJWTTokenResponse.Success(JWTToken(response.data.token))
             is Response.Error -> JetpackAIJWTTokenResponse.Error(
                 response.error.toJetpackAICompletionsError(),
                 response.error.message
@@ -52,14 +53,14 @@ class JetpackAIRestClient @Inject constructor(
     }
 
     suspend fun fetchJetpackAITextCompletion(
-        token: String,
+        token: JWTToken,
         prompt: String,
         feature: String
     ): JetpackAICompletionsResponse {
         val url = WPCOMV2.text_completion.url
         val body = mutableMapOf<String, String>()
         body.apply {
-            put("token", token)
+            put("token", token.value)
             put("prompt", prompt)
             put("feature", feature)
             put("_fields", FIELDS_TO_REQUEST)
@@ -136,7 +137,7 @@ class JetpackAIRestClient @Inject constructor(
     )
 
     sealed class JetpackAIJWTTokenResponse {
-        data class Success(val token: String) : JetpackAIJWTTokenResponse()
+        data class Success(val token: JWTToken) : JetpackAIJWTTokenResponse()
         data class Error(
             val type: JetpackAICompletionsErrorType,
             val message: String? = null

fluxc/src/main/java/org/wordpress/android/fluxc/store/jetpackai/JetpackAIStore.kt

@@ -1,5 +1,6 @@
 package org.wordpress.android.fluxc.store.jetpackai
 
+import org.wordpress.android.fluxc.model.JWTToken
 import org.wordpress.android.fluxc.model.SiteModel
 import org.wordpress.android.fluxc.network.rest.wpcom.jetpackai.JetpackAIRestClient
 import org.wordpress.android.fluxc.network.rest.wpcom.jetpackai.JetpackAIRestClient.JetpackAICompletionsErrorType.AUTH_ERROR
@@ -17,7 +18,7 @@ class JetpackAIStore @Inject constructor(
     private val jetpackAIRestClient: JetpackAIRestClient,
     private val coroutineEngine: CoroutineEngine
 ) {
-    private var token: String? = null
+    private var token: JWTToken? = null
 
     /**
      * Fetches Jetpack AI completions for a given prompt to be used on a particular post.
@@ -75,7 +76,7 @@ class JetpackAIStore @Inject constructor(
         caller = this,
         loggedMessage = "fetch Jetpack AI completions"
     ) {
-        val token = token ?: fetchJetpackAIJWTToken(site).let { tokenResponse ->
+        val token = token?.takeIfValid() ?: fetchJetpackAIJWTToken(site).let { tokenResponse ->
             when (tokenResponse) {
                 is Error -> {
                     return@withDefaultContext JetpackAICompletionsResponse.Error(