Fix JWT token refresh to Fusion validation request (#6231) [ci fast]

Signed-off-by: Paolo Di Tommaso <[email protected]>

Author: pditommaso

plugins/nf-tower/src/main/io/seqera/tower/plugin/TowerFactory.groovy

@@ -100,7 +100,7 @@ class TowerFactory implements TraceObserverFactory {
     @Memoized
     static TowerClient client(Session session, Map<String,String> env) {
         final config = session.config
-        Boolean isEnabled = config.navigate('tower.enabled') as Boolean || env.get('TOWER_WORKFLOW_ID')
+        Boolean isEnabled = config.navigate('tower.enabled') as Boolean || env.get('TOWER_WORKFLOW_ID') || config.navigate('fusion.enabled') as Boolean
         return isEnabled
             ? createTowerClient0(session, config, env)
             : null

plugins/nf-tower/src/main/io/seqera/tower/plugin/TowerFusionToken.groovy

@@ -15,6 +15,7 @@ import com.google.common.util.concurrent.UncheckedExecutionException
 import com.google.gson.Gson
 import com.google.gson.JsonSyntaxException
 import dev.failsafe.Failsafe
+import dev.failsafe.FailsafeException
 import dev.failsafe.RetryPolicy
 import dev.failsafe.event.EventListener
 import dev.failsafe.event.ExecutionAttemptedEvent
@@ -61,6 +62,8 @@ class TowerFusionToken implements FusionToken {
     private static final int DEFAULT_RETRY_POLICY_MAX_ATTEMPTS = 10
     private static final double DEFAULT_RETRY_POLICY_JITTER = 0.5
 
+    private CookieManager cookieManager = new CookieManager()
+
     // The HttpClient instance used to send requests
     private final HttpClient httpClient = newDefaultHttpClient()
 
@@ -79,7 +82,9 @@ class TowerFusionToken implements FusionToken {
     private String endpoint
 
     // Platform access token to use for requests
-    private String accessToken
+    private volatile String accessToken
+
+    private volatile String refreshToken
 
     // Platform workflowId
     private String workspaceId
@@ -92,6 +97,7 @@ class TowerFusionToken implements FusionToken {
         final env = SysEnv.get()
         this.endpoint = PlatformHelper.getEndpoint(config, env)
         this.accessToken = PlatformHelper.getAccessToken(config, env)
+        this.refreshToken = PlatformHelper.getRefreshToken(config, env)
         this.workflowId = env.get('TOWER_WORKFLOW_ID')
         this.workspaceId = PlatformHelper.getWorkspaceId(config, env)
     }
@@ -181,11 +187,11 @@ class TowerFusionToken implements FusionToken {
      * Create a new HttpClient instance with default settings
      * @return The new HttpClient instance
      */
-    private static HttpClient newDefaultHttpClient() {
+    private HttpClient newDefaultHttpClient() {
         final builder = HttpClient.newBuilder()
             .version(HttpClient.Version.HTTP_1_1)
             .followRedirects(HttpClient.Redirect.NEVER)
-            .cookieHandler(new CookieManager())
+            .cookieHandler(cookieManager)
             .connectTimeout(DEFAULT_CONNECTION_TIMEOUT)
         // use virtual threads executor if enabled
         if ( Threads.useVirtual() ) {
@@ -235,8 +241,17 @@ class TowerFusionToken implements FusionToken {
      * @param req The HttpRequest to send
      * @return The HttpResponse received
      */
-    private <T> HttpResponse<String> safeHttpSend(HttpRequest req, RetryPolicy<T> policy) {
-        return Failsafe.with(policy).get(
+    private <T> HttpResponse<String> safeHttpSend(HttpRequest req) {
+        try {
+            safeApply(req)
+        }
+        catch (FailsafeException e) {
+            throw e.cause
+        }
+    }
+
+    private <T> HttpResponse<String> safeApply(HttpRequest req) {
+        return Failsafe.with(retryPolicy).get(
             () -> {
                 log.debug "Http request: method=${req.method()}; uri=${req.uri()}; request=${req}"
                 final resp = httpClient.send(req, HttpResponse.BodyHandlers.ofString())
@@ -287,28 +302,35 @@ class TowerFusionToken implements FusionToken {
     /**
      * Request a license token from Platform.
      *
-     * @param req The LicenseTokenRequest object
+     * @param request The LicenseTokenRequest object
      * @return The LicenseTokenResponse object
-     *
-     * @throws AbortOperationException if a Platform access token cannot be found
-     * @throws UnauthorizedException if the access token is invalid
-     * @throws BadResponseException if the response is not as expected
-     * @throws IllegalStateException if the request cannot be sent
      */
-    private GetLicenseTokenResponse sendRequest(GetLicenseTokenRequest req) throws AbortOperationException, UnauthorizedException, BadResponseException, IllegalStateException {
+    private GetLicenseTokenResponse sendRequest(GetLicenseTokenRequest request) {
+        return sendRequest0(request, 1)
+    }
 
-        final httpReq = makeHttpRequest(req)
+    private GetLicenseTokenResponse sendRequest0(GetLicenseTokenRequest request, int attempt) {
+
+        final httpReq = makeHttpRequest(request)
 
         try {
-            final resp = safeHttpSend(httpReq, retryPolicy)
+            final resp = safeHttpSend(httpReq)
 
             if( resp.statusCode() == 200 ) {
                 final ret = parseLicenseTokenResponse(resp.body())
                 return ret
             }
 
             if( resp.statusCode() == 401 ) {
-                throw new UnauthorizedException("Unauthorized [401] - Verify you have provided a valid access token")
+                final shouldRetry = accessToken
+                    && refreshToken
+                    && attempt==1
+                    && refreshJwtToken0(refreshToken)
+                if( shouldRetry ) {
+                    return sendRequest0(request, attempt+1)
+                }
+                else
+                    throw new UnauthorizedException("Unauthorized [401] - Verify you have provided a valid Seqera Platform access token")
             }
 
             throw new BadResponseException("Invalid response: ${httpReq.method()} ${httpReq.uri()} [${resp.statusCode()}] ${resp.body()}")
@@ -317,4 +339,52 @@ class TowerFusionToken implements FusionToken {
             throw new IllegalStateException("Unable to send request to '${httpReq.uri()}' : ${e.message}")
         }
     }
+
+    protected boolean refreshJwtToken0(String refresh) {
+        log.debug "Token refresh request >> $refresh"
+
+        final req = HttpRequest.newBuilder()
+            .uri(new URI("${endpoint}/oauth/access_token"))
+            .headers('Content-Type',"application/x-www-form-urlencoded")
+            .POST(HttpRequest.BodyPublishers.ofString("grant_type=refresh_token&refresh_token=${URLEncoder.encode(refresh, 'UTF-8')}"))
+            .build()
+
+        final resp = safeHttpSend(req)
+        final code = resp.statusCode()
+        final body = resp.body()
+        log.debug "Refresh cookie response: [${code}] ${body}"
+        if( resp.statusCode() != 200 )
+            return false
+
+        final authCookie = getCookie('JWT')
+        final refreshCookie = getCookie('JWT_REFRESH_TOKEN')
+
+        // set the new bearer token in the current client session
+        if( authCookie?.value ) {
+            log.trace "Updating http client bearer token=$authCookie.value"
+            accessToken = authCookie.value
+        }
+        else {
+            log.warn "Missing JWT cookie from refresh token response ~ $authCookie"
+        }
+
+        // set the new refresh token
+        if( refreshCookie?.value ) {
+            log.trace "Updating http client refresh token=$refreshCookie.value"
+            refreshToken = refreshCookie.value
+        }
+        else {
+            log.warn "Missing JWT_REFRESH_TOKEN cookie from refresh token response ~ $refreshCookie"
+        }
+
+        return true
+    }
+
+    private HttpCookie getCookie(final String cookieName) {
+        for( HttpCookie it : cookieManager.cookieStore.cookies ) {
+            if( it.name == cookieName )
+                return it
+        }
+        return null
+    }
 }

plugins/nf-tower/src/test/io/seqera/tower/plugin/TowerFusionEnvTest.groovy

@@ -5,6 +5,7 @@ import java.time.temporal.ChronoUnit
 
 import com.github.tomakehurst.wiremock.WireMockServer
 import com.github.tomakehurst.wiremock.client.WireMock
+import com.github.tomakehurst.wiremock.stubbing.Scenario
 import io.seqera.tower.plugin.exception.UnauthorizedException
 import nextflow.Global
 import nextflow.Session
@@ -43,7 +44,6 @@ class TowerFusionEnvTest extends Specification {
         SysEnv.pop()      // <-- restore the system host env
     }
 
-
     def 'should return the endpoint from the config'() {
         given: 'a session'
         Global.session = Mock(Session) {
@@ -287,7 +287,6 @@ class TowerFusionEnvTest extends Specification {
         and: 'the request is correct'
         wireMockServer.verify(1, WireMock.postRequestedFor(WireMock.urlEqualTo("/license/token/"))
             .withHeader('Authorization', WireMock.equalTo('Bearer abc123')))
-
     }
 
     def 'should get a license token with environment'() {
@@ -336,7 +335,94 @@ class TowerFusionEnvTest extends Specification {
         SysEnv.pop()
     }
 
+    def 'should refresh the auth token on 401 and retry the request'() {
+        given:
+        SysEnv.push([
+            TOWER_WORKFLOW_ID: '12345',
+            TOWER_ACCESS_TOKEN: 'abc-token',
+            TOWER_REFRESH_TOKEN: 'xyz-refresh',
+            TOWER_WORKSPACE_ID: '67890',
+            TOWER_API_ENDPOINT: wireMockServer.baseUrl()
+        ])
+        def PRODUCT = 'some-product'
+        def VERSION = 'some-version'
+        and:
+        Global.session = Mock(Session) { getConfig() >> [:] }
+        and:
+        def provider = new TowerFusionToken()
+
+        and: 'prepare stubs'
+
+        final now = Instant.now()
+        final expirationDate = GsonHelper.toJson(now.plus(1, ChronoUnit.DAYS))
+
+        // 1️⃣ First attempt: /license/token/ fails with 401
+        wireMockServer.stubFor(
+            WireMock.post(urlEqualTo("/license/token/"))
+                .inScenario("Refresh flow")
+                .whenScenarioStateIs(Scenario.STARTED)
+                .willReturn(WireMock.aResponse().withStatus(401))
+                .willSetStateTo("Token Refreshed")
+        )
+
+        // 2️⃣ Refresh token call
+        wireMockServer.stubFor(
+            WireMock.post(urlEqualTo("/oauth/access_token"))
+                .inScenario("Refresh flow")
+                .whenScenarioStateIs("Token Refreshed")
+                .withHeader('Content-Type', equalTo('application/x-www-form-urlencoded'))
+                .withRequestBody(containing('grant_type=refresh_token'))
+                .withRequestBody(containing(URLEncoder.encode('xyz-refresh', 'UTF-8')))
+                .willReturn(
+                    WireMock.aResponse()
+                        .withStatus(200)
+                        .withHeader('Set-Cookie', 'JWT=new-abc-token')
+                        .withHeader('Set-Cookie', 'JWT_REFRESH_TOKEN=new-refresh-456')
+                )
+                .willSetStateTo("Retry Ready")
+        )
 
+        // 3️⃣ Retry: /license/token/ succeeds
+        wireMockServer.stubFor(
+            WireMock.post(urlEqualTo("/license/token/"))
+                .inScenario("Refresh flow")
+                .whenScenarioStateIs("Retry Ready")
+                .withHeader('Authorization', equalTo('Bearer new-abc-token'))
+                .willReturn(
+                    WireMock.aResponse()
+                        .withStatus(200)
+                        .withHeader('Content-Type', 'application/json')
+                        .withBody('{"signedToken":"xyz789", "expiresAt":' + expirationDate + '}')
+                )
+        )
+
+        when:
+        final token = provider.getLicenseToken(PRODUCT, VERSION)
+
+        then:
+        token == 'xyz789'
+
+        and: 'the initial request was sent with the old token'
+        wireMockServer.verify(1, WireMock.postRequestedFor(urlEqualTo("/license/token/"))
+            .withHeader('Authorization', equalTo('Bearer abc-token'))
+        )
+
+        and: 'the refresh request was sent with correct form data'
+        wireMockServer.verify(1, WireMock.postRequestedFor(WireMock.urlEqualTo("/oauth/access_token"))
+            .withHeader('Content-Type', equalTo('application/x-www-form-urlencoded'))
+            .withRequestBody(containing('grant_type=refresh_token'))
+            .withRequestBody(containing(URLEncoder.encode('xyz-refresh', 'UTF-8')))
+        )
+
+        and: 'the retried request was sent with the new token'
+        wireMockServer.verify(1, WireMock.postRequestedFor(WireMock.urlEqualTo("/license/token/"))
+            .withHeader('Authorization', equalTo('Bearer new-abc-token'))
+        )
+
+        cleanup:
+        SysEnv.pop()
+    }
+    
     def 'should throw UnauthorizedException if getting a token fails with 401'() {
         given: 'a TowerFusionEnv provider'
         Global.session = Mock(Session) {
@@ -369,7 +455,6 @@ class TowerFusionEnvTest extends Specification {
         thrown(UnauthorizedException)
     }
 
-
     def 'should deserialize response' () {
         given:
         def ts = Instant.ofEpochSecond(1738788914)