From 1412fce30f045279c806287015ee866a4a572bc5 Mon Sep 17 00:00:00 2001
From: Silvio Giebl <gieblsilvio@outlook.com>
Date: Mon, 22 Jul 2024 18:19:23 +0200
Subject: [PATCH] Support registry tokens that are not JWTs This is required
 for the GitHub Container Registry

---
 .../oci/internal/registry/OciRegistryApi.kt   | 19 +++++++++++--------
 .../oci/internal/registry/OciRegistryToken.kt | 16 ++++++++++++----
 2 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/src/main/kotlin/io/github/sgtsilvio/gradle/oci/internal/registry/OciRegistryApi.kt b/src/main/kotlin/io/github/sgtsilvio/gradle/oci/internal/registry/OciRegistryApi.kt
index 7977a08a..97fa2aa3 100644
--- a/src/main/kotlin/io/github/sgtsilvio/gradle/oci/internal/registry/OciRegistryApi.kt
+++ b/src/main/kotlin/io/github/sgtsilvio/gradle/oci/internal/registry/OciRegistryApi.kt
@@ -32,6 +32,7 @@ import reactor.util.retry.RetrySpec
 import java.net.URI
 import java.security.DigestException
 import java.security.MessageDigest
+import java.time.Duration
 import java.time.Instant
 import java.time.temporal.ChronoUnit
 import java.util.*
@@ -80,8 +81,10 @@ internal class OciRegistryApi(httpClient: HttpClient) {
     )
 
     private object TokenCacheExpiry : Expiry<TokenCacheKey, OciRegistryToken> {
-        override fun expireAfterCreate(key: TokenCacheKey, value: OciRegistryToken, currentTime: Long) =
-            Instant.now().until(value.payload.expirationTime!!.minusSeconds(30), ChronoUnit.NANOS).coerceAtLeast(0)
+        override fun expireAfterCreate(key: TokenCacheKey, value: OciRegistryToken, currentTime: Long): Long {
+            val expirationTime = value.claims?.expirationTime ?: return Duration.ofMinutes(3).toNanos()
+            return Instant.now().until(expirationTime.minusSeconds(30), ChronoUnit.NANOS).coerceAtLeast(0)
+        }
 
         override fun expireAfterUpdate(
             key: TokenCacheKey,
@@ -569,12 +572,12 @@ internal class OciRegistryApi(httpClient: HttpClient) {
                     else -> createError(response, body)
                 }
             }.retryWhen(RETRY_SPEC).map { response ->
-                val jws = jsonObject(response).run {
+                val token = jsonObject(response).run {
                     if (hasKey("token")) getString("token") else getString("access_token")
                 }
-                val registryToken = OciRegistryToken(jws)
-                val grantedScopes = registryToken.payload.scopes
-                if (grantedScopes == key.scopes) {
+                val registryToken = OciRegistryToken(token)
+                val grantedScopes = registryToken.claims?.scopes
+                if ((grantedScopes == null) || (grantedScopes == key.scopes)) {
                     registryToken
                 } else {
                     tokenCache.asMap().putIfAbsent(
@@ -584,7 +587,7 @@ internal class OciRegistryApi(httpClient: HttpClient) {
                     throw InsufficientScopesException(key.scopes, grantedScopes)
                 }
             }
-        }.map { encodeBearerAuthorization(it.jws) }
+        }.map { encodeBearerAuthorization(it.token) }
     }
 
     private fun getAuthorization(
@@ -593,7 +596,7 @@ internal class OciRegistryApi(httpClient: HttpClient) {
         credentials: Credentials?,
     ): Mono<String> {
         return tokenCache.getIfPresentMono(TokenCacheKey(registryUrl, scopes, credentials?.hashed()))
-            .map { encodeBearerAuthorization(it.jws) }
+            .map { encodeBearerAuthorization(it.token) }
             .run { if (credentials == null) this else defaultIfEmpty(credentials.encodeBasicAuthorization()) }
     }
 
diff --git a/src/main/kotlin/io/github/sgtsilvio/gradle/oci/internal/registry/OciRegistryToken.kt b/src/main/kotlin/io/github/sgtsilvio/gradle/oci/internal/registry/OciRegistryToken.kt
index fd583016..38c0d7fa 100644
--- a/src/main/kotlin/io/github/sgtsilvio/gradle/oci/internal/registry/OciRegistryToken.kt
+++ b/src/main/kotlin/io/github/sgtsilvio/gradle/oci/internal/registry/OciRegistryToken.kt
@@ -10,11 +10,19 @@ import java.time.Instant
 /**
  * @author Silvio Giebl
  */
-internal class OciRegistryToken(val jws: String) {
-    val payload = jsonObject(jws.decodeToJWS().payload).decodeOciRegistryTokenPayload()
+internal class OciRegistryToken(val token: String, val claims: OciRegistryTokenClaims?)
+
+internal fun OciRegistryToken(token: String): OciRegistryToken {
+    val jws = try {
+        token.decodeToJWS()
+    } catch (e: IllegalArgumentException) {
+        return OciRegistryToken(token, null)
+    }
+    val claims = jsonObject(jws.payload).decodeOciRegistryTokenClaims()
+    return OciRegistryToken(token, claims)
 }
 
-internal data class OciRegistryTokenPayload(
+internal data class OciRegistryTokenClaims(
     val issuer: String?,
     val subject: String?,
     val audience: List<String>,
@@ -25,7 +33,7 @@ internal data class OciRegistryTokenPayload(
     val scopes: Set<OciRegistryResourceScope>,
 )
 
-internal fun JsonObject.decodeOciRegistryTokenPayload() = OciRegistryTokenPayload(
+private fun JsonObject.decodeOciRegistryTokenClaims() = OciRegistryTokenClaims(
     getStringOrNull("iss"),
     getStringOrNull("sub"),
     getOrNull("aud") { if (isString()) listOf(asString()) else asArray().toStringList() } ?: emptyList(),