Configuring Spring's WebClient for OAuth2 JWT Bearer Client Authentication (client_secret_jwt with private key) - MVC non-reactive version

application.yaml

spring:
  security:
    oauth2:
      client:
        registration:
          someidentifier:
            client-id: your_client_id
            authorization-grant-type: urn:ietf:params:oauth:grant-type:jwt-bearer
            # do not set client-authentication-method=client_secret_jwt, it's not used correctly
        provider:
          someidentifier:
            # Salesforce sample
            token-uri: <your_url>/services/oauth2/token

Customizer.kt

/**
 * Note that consumers will have to add an attribute with registration ID to their requests, see example.
 *
 * @constructor
 *
 * The `failureHandler` is implemented as indicated in the javadoc of [ServletOAuth2AuthorizedClientExchangeFilterFunction]'s constructor.
 */
@Configuration
class GatewayWebClientCustomizer(
    clientRegistrationRepository: ClientRegistrationRepository,
    authorizedClientService: OAuth2AuthorizedClientService,
) : WebClientCustomizer {

    companion object {
        private val jwtHeaders = mapOf(
            "typ" to "JWT",
            "alg" to "RS256",
        )
    }

    private val oAuth2FilterFunction: ServletOAuth2AuthorizedClientExchangeFilterFunction

    init {
        val clientProvider = JwtBearerOAuth2AuthorizedClientProvider().apply {
            // see CustomJwtBearerTokenResponseClient
            //setAccessTokenResponseClient(CustomJwtBearerTokenResponseClient(Duration.ofMinutes(3L)))
            setJwtAssertionResolver { authorizationContext ->
                // compute assertion, see for example https://help.salesforce.com/s/articleView?id=sf.remoteaccess_oauth_jwt_flow.htm
                // you'll probably need authorizationContext.clientRegistration.[registrationId|clientId]
                Jwt(assertion, null, null, jwtHeaders, claims)
            }
        }
        val authorizedClientManager = AuthorizedClientServiceOAuth2AuthorizedClientManager(clientRegistrationRepository, authorizedClientService).apply {
            setAuthorizedClientProvider(
                OAuth2AuthorizedClientProviderBuilder.builder().provider(clientProvider).build()
            )
        }

        oAuth2FilterFunction = ServletOAuth2AuthorizedClientExchangeFilterFunction(authorizedClientManager)
        val failureHandler = RemoveAuthorizedClientOAuth2AuthorizationFailureHandler { clientRegistrationId, principal, _ ->
            authorizedClientService.removeAuthorizedClient(clientRegistrationId, principal.name)
        }
        oAuth2FilterFunction.setAuthorizationFailureHandler(failureHandler)
    }

    override fun customize(webClientBuilder: WebClient.Builder) {
        webClientBuilder
            .apply(oAuth2FilterFunction.oauth2Configuration())
    }
}

CustomJwtBearerTokenResponseClient.kt

// Salesforce doesn't include expires_in in its token response, so Spring assumes it's only valid for 1 second.
// Considering clock skew logic, each token is basically valid for a single call,
// so you might want to force the assumption that it can be valid for a little longer.
class CustomJwtBearerTokenResponseClient(
    private val minimumTokenDuration: Duration
) : OAuth2AccessTokenResponseClient<JwtBearerGrantRequest> {

    private val defaultClient = DefaultJwtBearerTokenResponseClient()

    override fun getTokenResponse(authorizationGrantRequest: JwtBearerGrantRequest): OAuth2AccessTokenResponse {
        val response = defaultClient.getTokenResponse(authorizationGrantRequest)
        return if (Duration.between(Instant.now(), response.accessToken.expiresAt) >= minimumTokenDuration) {
            response
        } else {
            OAuth2AccessTokenResponse.withResponse(response)
                .expiresIn(minimumTokenDuration.seconds)
                .build()
        }
    }
}

SampleConsumer.kt

class Example(builder: WebClient.Builder) {
    private val webClient = builder
        .baseUrl(uri)
        .defaultRequest { it.attributes(ServletOAuth2AuthorizedClientExchangeFilterFunction.clientRegistrationId(salesforceInstanceKey)) }
        .build()
}