From 34e4065b18cf39c1f5f620982c884bb4e6ad89a0 Mon Sep 17 00:00:00 2001 From: jrhee17 Date: Thu, 9 Jul 2026 10:49:36 +0900 Subject: [PATCH 1/3] minimal impl --- .../client/athenz/DefaultZtsBaseClient.java | 139 ++++++++++++++++++ .../armeria/client/athenz/ZtsBaseClient.java | 123 +++------------- .../client/athenz/ZtsBaseClientBuilder.java | 8 +- .../server/athenz/MinifiedAuthZpeClient.java | 56 +------ .../jwts/ArmeriaJwtsSigningKeyResolver.java | 67 +++++++++ .../token/jwts/ArmeriaResourceRetriever.java | 57 +++++++ .../athenz/auth/token/jwts/package-info.java | 25 ++++ .../ArmeriaJwtsSigningKeyResolverTest.java | 94 ++++++++++++ 8 files changed, 413 insertions(+), 156 deletions(-) create mode 100644 athenz/src/main/java/com/linecorp/armeria/client/athenz/DefaultZtsBaseClient.java create mode 100644 athenz/src/main/java/com/yahoo/athenz/auth/token/jwts/ArmeriaJwtsSigningKeyResolver.java create mode 100644 athenz/src/main/java/com/yahoo/athenz/auth/token/jwts/ArmeriaResourceRetriever.java create mode 100644 athenz/src/main/java/com/yahoo/athenz/auth/token/jwts/package-info.java create mode 100644 athenz/src/test/java/com/linecorp/armeria/client/athenz/ArmeriaJwtsSigningKeyResolverTest.java diff --git a/athenz/src/main/java/com/linecorp/armeria/client/athenz/DefaultZtsBaseClient.java b/athenz/src/main/java/com/linecorp/armeria/client/athenz/DefaultZtsBaseClient.java new file mode 100644 index 00000000000..2963f8ff3f6 --- /dev/null +++ b/athenz/src/main/java/com/linecorp/armeria/client/athenz/DefaultZtsBaseClient.java @@ -0,0 +1,139 @@ +/* + * Copyright 2026 LY Corporation + * + * LY Corporation licenses this file to you under the Apache License, + * version 2.0 (the "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at: + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT + * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the + * License for the specific language governing permissions and limitations + * under the License. + */ + +package com.linecorp.armeria.client.athenz; + +import static java.util.Objects.requireNonNull; + +import java.net.URI; +import java.security.cert.X509Certificate; +import java.time.Duration; +import java.util.List; +import java.util.function.Consumer; +import java.util.function.Supplier; + +import org.slf4j.LoggerFactory; + +import com.linecorp.armeria.client.ClientFactory; +import com.linecorp.armeria.client.ClientFactoryBuilder; +import com.linecorp.armeria.client.WebClient; +import com.linecorp.armeria.client.WebClientBuilder; +import com.linecorp.armeria.client.logging.LoggingClient; +import com.linecorp.armeria.client.retry.RetryRule; +import com.linecorp.armeria.client.retry.RetryingClient; +import com.linecorp.armeria.common.CommonPools; +import com.linecorp.armeria.common.TlsKeyPair; +import com.linecorp.armeria.common.TlsProvider; +import com.linecorp.armeria.common.annotation.Nullable; +import com.linecorp.armeria.common.logging.LogLevel; +import com.linecorp.armeria.common.logging.LogWriter; +import com.linecorp.armeria.common.util.AbstractListenable; + +/** + * The default {@link ZtsBaseClient} implementation that manages a {@link ClientFactory} with + * {@link TlsKeyPair} rotation and proxy configuration. + */ +final class DefaultZtsBaseClient implements ZtsBaseClient { + + private final TlsKeyPairListener tlsKeyPairListener = new TlsKeyPairListener(); + private final URI ztsUri; + private final ClientFactory clientFactory; + @Nullable + private final Consumer webClientConfigurer; + private final WebClient defaultWebClient; + + DefaultZtsBaseClient(URI ztsUri, Supplier keyPairSupplier, + List trustedCertificates, int autoKeyRefreshIntervalMillis, + @Nullable Consumer clientFactoryConfigurer, + @Nullable Consumer webClientConfigurer) { + + this.ztsUri = ztsUri; + final TlsProvider tlsProvider = + TlsProvider.ofScheduled(keyPairSupplier, + trustedCertificates, + tlsKeyPairListener::onTlsKeyPairUpdated, + Duration.ofMillis(autoKeyRefreshIntervalMillis), + CommonPools.blockingTaskExecutor()); + + final ClientFactoryBuilder factoryBuilder = ClientFactory.builder().tlsProvider(tlsProvider); + if (clientFactoryConfigurer != null) { + clientFactoryConfigurer.accept(factoryBuilder); + } + clientFactory = factoryBuilder.build(); + this.webClientConfigurer = webClientConfigurer; + defaultWebClient = webClient(builder -> {}); + } + + @Override + public ClientFactory clientFactory() { + return clientFactory; + } + + @Override + public WebClient webClient() { + return defaultWebClient; + } + + @Override + public WebClient webClient(Consumer configurer) { + requireNonNull(configurer, "configurer"); + + final WebClientBuilder clientBuilder = + WebClient.builder(ztsUri) + .decorator(RetryingClient.newDecorator(RetryRule.failsafe())); + + clientBuilder.factory(clientFactory); + + if (LoggerFactory.getLogger(DefaultZtsBaseClient.class).isTraceEnabled()) { + final LogWriter logWriter = LogWriter.builder() + .logger(DefaultZtsBaseClient.class.getName()) + .requestLogLevel(LogLevel.TRACE) + .successfulResponseLogLevel(LogLevel.TRACE) + .build(); + clientBuilder.decorator(LoggingClient.builder() + .logWriter(logWriter) + .newDecorator()); + } + if (webClientConfigurer != null) { + webClientConfigurer.accept(clientBuilder); + } + configurer.accept(clientBuilder); + return clientBuilder.build(); + } + + @Override + public void addTlsKeyPairListener(Consumer listener) { + requireNonNull(listener, "listener"); + tlsKeyPairListener.addListener(listener); + } + + @Override + public void removeTlsKeyPairListener(Consumer listener) { + requireNonNull(listener, "listener"); + tlsKeyPairListener.removeListener(listener); + } + + @Override + public void close() { + defaultWebClient.options().factory().closeAsync(); + } + + private static final class TlsKeyPairListener extends AbstractListenable { + void onTlsKeyPairUpdated(TlsKeyPair newTlsKeyPair) { + notifyListeners(newTlsKeyPair); + } + } +} diff --git a/athenz/src/main/java/com/linecorp/armeria/client/athenz/ZtsBaseClient.java b/athenz/src/main/java/com/linecorp/armeria/client/athenz/ZtsBaseClient.java index c02cb60fae8..53f5f7feb67 100644 --- a/athenz/src/main/java/com/linecorp/armeria/client/athenz/ZtsBaseClient.java +++ b/athenz/src/main/java/com/linecorp/armeria/client/athenz/ZtsBaseClient.java @@ -19,30 +19,14 @@ import static java.util.Objects.requireNonNull; import java.net.URI; -import java.net.URL; -import java.security.cert.X509Certificate; -import java.time.Duration; -import java.util.List; import java.util.function.Consumer; -import java.util.function.Supplier; - -import org.slf4j.LoggerFactory; import com.linecorp.armeria.client.ClientFactory; -import com.linecorp.armeria.client.ClientFactoryBuilder; import com.linecorp.armeria.client.WebClient; import com.linecorp.armeria.client.WebClientBuilder; -import com.linecorp.armeria.client.logging.LoggingClient; -import com.linecorp.armeria.client.retry.RetryRule; -import com.linecorp.armeria.client.retry.RetryingClient; -import com.linecorp.armeria.common.CommonPools; import com.linecorp.armeria.common.TlsKeyPair; -import com.linecorp.armeria.common.TlsProvider; import com.linecorp.armeria.common.annotation.Nullable; import com.linecorp.armeria.common.annotation.UnstableApi; -import com.linecorp.armeria.common.logging.LogLevel; -import com.linecorp.armeria.common.logging.LogWriter; -import com.linecorp.armeria.common.util.AbstractListenable; import com.linecorp.armeria.common.util.SafeCloseable; import com.linecorp.armeria.server.athenz.AthenzService; @@ -61,12 +45,12 @@ * } */ @UnstableApi -public final class ZtsBaseClient implements SafeCloseable { +public interface ZtsBaseClient extends SafeCloseable { /** * Returns a new {@link ZtsBaseClientBuilder} for the specified ZTS URI. */ - public static ZtsBaseClientBuilder builder(String ztsUri) { + static ZtsBaseClientBuilder builder(String ztsUri) { requireNonNull(ztsUri, "ztsUri"); return builder(URI.create(ztsUri)); } @@ -74,7 +58,7 @@ public static ZtsBaseClientBuilder builder(String ztsUri) { /** * Returns a new {@link ZtsBaseClientBuilder} for the specified ZTS {@link URI}. */ - public static ZtsBaseClientBuilder builder(URI ztsUri) { + static ZtsBaseClientBuilder builder(URI ztsUri) { requireNonNull(ztsUri, "ztsUri"); return new ZtsBaseClientBuilder(normalizeZtsUri(ztsUri)); } @@ -91,120 +75,59 @@ private static URI normalizeZtsUri(URI ztsUri) { return URI.create(rawUri); } - private final TlsKeyPairListener tlsKeyPairListener = new TlsKeyPairListener(); - private final URI ztsUri; - @Nullable - private final URI proxyUri; - private final ClientFactory clientFactory; - @Nullable - private final Consumer webClientConfigurer; - private final WebClient defaultWebClient; - - ZtsBaseClient(URI ztsUri, @Nullable URI proxyUri, Supplier keyPairSupplier, - List trustedCertificates, int autoKeyRefreshIntervalMillis, - @Nullable Consumer clientFactoryConfigurer, - @Nullable Consumer webClientConfigurer) { - - this.ztsUri = ztsUri; - this.proxyUri = proxyUri; - final TlsProvider tlsProvider = - TlsProvider.ofScheduled(keyPairSupplier, - trustedCertificates, - tlsKeyPairListener::onTlsKeyPairUpdated, - Duration.ofMillis(autoKeyRefreshIntervalMillis), - CommonPools.blockingTaskExecutor()); - - final ClientFactoryBuilder factoryBuilder = ClientFactory.builder().tlsProvider(tlsProvider); - if (clientFactoryConfigurer != null) { - clientFactoryConfigurer.accept(factoryBuilder); - } - clientFactory = factoryBuilder.build(); - this.webClientConfigurer = webClientConfigurer; - defaultWebClient = webClient(builder -> {}); - } - /** - * Returns the ZTS {@link URL} that this client connects to. + * Returns the ZTS {@link URI} that this client connects to. + * + * @deprecated Use {@code webClient().uri()} instead. */ - public URI ztsUri() { - return ztsUri; + @Deprecated + default URI ztsUri() { + return webClient().uri(); } /** - * Returns the proxy {@link URL} if it was configured, or {@code null} if no proxy is set. + * Returns the proxy {@link URI} if it was configured, or {@code null} if no proxy is set. + * + * @deprecated The proxy is configured at the {@link ClientFactory} level via + * {@link ZtsBaseClientBuilder#proxyUri(URI)} and is already applied to the {@link WebClient}. */ + @Deprecated @Nullable - public URI proxyUri() { - return proxyUri; + default URI proxyUri() { + return null; } /** * Returns the {@link ClientFactory} that is used to create the {@link WebClient} instances. */ - public ClientFactory clientFactory() { - return clientFactory; + default ClientFactory clientFactory() { + return webClient().options().factory(); } /** * Returns the {@link WebClient} that can connect to the ZTS server. */ - public WebClient webClient() { - return defaultWebClient; - } + WebClient webClient(); /** * Returns a new {@link WebClient} that can connect to the ZTS server with the specified configurer. */ - public WebClient webClient(Consumer configurer) { + default WebClient webClient(Consumer configurer) { requireNonNull(configurer, "configurer"); - - final WebClientBuilder clientBuilder = - WebClient.builder(ztsUri) - .decorator(RetryingClient.newDecorator(RetryRule.failsafe())); - - clientBuilder.factory(clientFactory); - - if (LoggerFactory.getLogger(ZtsBaseClient.class).isTraceEnabled()) { - final LogWriter logWriter = LogWriter.builder() - .logger(ZtsBaseClient.class.getName()) - .requestLogLevel(LogLevel.TRACE) - .successfulResponseLogLevel(LogLevel.TRACE) - .build(); - clientBuilder.decorator(LoggingClient.builder() - .logWriter(logWriter) - .newDecorator()); - } - if (webClientConfigurer != null) { - webClientConfigurer.accept(clientBuilder); - } - configurer.accept(clientBuilder); - return clientBuilder.build(); + return webClient(); } /** * Adds a listener that will be notified when the {@link TlsKeyPair} is updated. */ - public void addTlsKeyPairListener(Consumer listener) { + default void addTlsKeyPairListener(Consumer listener) { requireNonNull(listener, "listener"); - tlsKeyPairListener.addListener(listener); } /** * Removes a listener that was previously added with {@link #addTlsKeyPairListener(Consumer)}. */ - public void removeTlsKeyPairListener(Consumer listener) { + default void removeTlsKeyPairListener(Consumer listener) { requireNonNull(listener, "listener"); - tlsKeyPairListener.removeListener(listener); - } - - @Override - public void close() { - defaultWebClient.options().factory().closeAsync(); - } - - private static final class TlsKeyPairListener extends AbstractListenable { - void onTlsKeyPairUpdated(TlsKeyPair newTlsKeyPair) { - notifyListeners(newTlsKeyPair); - } } } diff --git a/athenz/src/main/java/com/linecorp/armeria/client/athenz/ZtsBaseClientBuilder.java b/athenz/src/main/java/com/linecorp/armeria/client/athenz/ZtsBaseClientBuilder.java index f6cf63d0ebb..1156074ba72 100644 --- a/athenz/src/main/java/com/linecorp/armeria/client/athenz/ZtsBaseClientBuilder.java +++ b/athenz/src/main/java/com/linecorp/armeria/client/athenz/ZtsBaseClientBuilder.java @@ -53,8 +53,6 @@ public final class ZtsBaseClientBuilder { private final URI ztsUri; - @Nullable - private URI proxyUri; @Nullable private Supplier athenzKeyPairSupplier; @@ -155,7 +153,6 @@ public ZtsBaseClientBuilder proxyUri(String proxyUrl) { */ public ZtsBaseClientBuilder proxyUri(URI proxyUri) { requireNonNull(proxyUri, "proxyUri"); - this.proxyUri = proxyUri; final boolean isTls = "https".equalsIgnoreCase(proxyUri.getScheme()); final int port = proxyUri.getPort() == -1 ? (isTls ? 443 : 80) : proxyUri.getPort(); final InetSocketAddress proxyAddress = InetSocketAddress.createUnresolved(proxyUri.getHost(), port); @@ -211,7 +208,8 @@ public ZtsBaseClient build() { if (athenzKeyPairSupplier == null) { throw new IllegalStateException("Athenz key pair must be set"); } - return new ZtsBaseClient(ztsUri, proxyUri, athenzKeyPairSupplier, trustedCertificates, - autoKeyRefreshIntervalMillis, clientFactoryConfigurer, webClientConfigurer); + return new DefaultZtsBaseClient(ztsUri, athenzKeyPairSupplier, trustedCertificates, + autoKeyRefreshIntervalMillis, clientFactoryConfigurer, + webClientConfigurer); } } diff --git a/athenz/src/main/java/com/linecorp/armeria/server/athenz/MinifiedAuthZpeClient.java b/athenz/src/main/java/com/linecorp/armeria/server/athenz/MinifiedAuthZpeClient.java index 986cd2b5fcd..bb51a46c0fe 100644 --- a/athenz/src/main/java/com/linecorp/armeria/server/athenz/MinifiedAuthZpeClient.java +++ b/athenz/src/main/java/com/linecorp/armeria/server/athenz/MinifiedAuthZpeClient.java @@ -32,7 +32,6 @@ import static com.yahoo.athenz.zpe.ZpeConsts.ZPE_PROP_MILLIS_BETWEEN_ZTS_CALLS; -import java.net.URI; import java.security.PublicKey; import java.security.cert.X509Certificate; import java.util.List; @@ -49,6 +48,7 @@ import com.oath.auth.Utils; import com.yahoo.athenz.auth.token.AccessToken; import com.yahoo.athenz.auth.token.RoleToken; +import com.yahoo.athenz.auth.token.jwts.ArmeriaJwtsSigningKeyResolver; import com.yahoo.athenz.auth.token.jwts.JwtsSigningKeyResolver; import com.yahoo.athenz.auth.util.CryptoException; import com.yahoo.athenz.zpe.ZpeConsts; @@ -56,17 +56,8 @@ import com.yahoo.athenz.zpe.pkey.PublicKeyStore; import com.yahoo.rdl.Struct; -import com.linecorp.armeria.client.ClientFactory; -import com.linecorp.armeria.client.ClientTlsSpec; -import com.linecorp.armeria.client.ClientTlsSpecBuilder; import com.linecorp.armeria.client.athenz.ZtsBaseClient; -import com.linecorp.armeria.common.TlsKeyPair; -import com.linecorp.armeria.common.TlsProvider; import com.linecorp.armeria.common.annotation.Nullable; -import com.linecorp.armeria.common.util.TlsEngineType; -import com.linecorp.armeria.internal.common.util.SslContextUtil; - -import io.netty.handler.ssl.JdkSslContext; final class MinifiedAuthZpeClient { @@ -111,11 +102,8 @@ final class MinifiedAuthZpeClient { private void initializeAccessTokenSignKeyResolver(ZtsBaseClient ztsBaseClient, String oauth2KeysPath) { final String serverUrl = System.getProperty(ZpeConsts.ZPE_PROP_JWK_URI); if (serverUrl == null || serverUrl.isEmpty()) { - accessSignKeyResolver = newDefaultJwtsSigningKeyResolver(ztsBaseClient, oauth2KeysPath); - ztsBaseClient.addTlsKeyPairListener(tlsKeyPair -> { - // Refresh the JwtsSigningKeyResolver when the TLS key pair changes - accessSignKeyResolver = newDefaultJwtsSigningKeyResolver(ztsBaseClient, oauth2KeysPath); - }); + accessSignKeyResolver = ArmeriaJwtsSigningKeyResolver.create( + ztsBaseClient.webClient(), oauth2KeysPath); return; } @@ -132,42 +120,8 @@ private void initializeAccessTokenSignKeyResolver(ZtsBaseClient ztsBaseClient, S logger.warn("Unable to initialize key refresher: {}", ex.getMessage()); } } - final URI proxyUri = ztsBaseClient.proxyUri(); - accessSignKeyResolver = new JwtsSigningKeyResolver(serverUrl, sslContext, - proxyUri != null ? proxyUri.toString() : null); - } - - private static JwtsSigningKeyResolver newDefaultJwtsSigningKeyResolver(ZtsBaseClient ztsBaseClient, - String oauth2KeysPath) { - final URI ztsUri = ztsBaseClient.ztsUri(); - final URI proxyUri = ztsBaseClient.proxyUri(); - String proxyUriStr = null; - if (proxyUri != null) { - proxyUriStr = proxyUri.toString(); - } - final ClientFactory clientFactory = ztsBaseClient.clientFactory(); - final TlsProvider tlsProvider = clientFactory.options().tlsProvider(); - final boolean allowUnsafeCiphers = clientFactory.options().tlsConfig().allowsUnsafeCiphers(); - final ClientTlsSpec clientTlsSpec = toTlsSpec(tlsProvider, allowUnsafeCiphers); - final JdkSslContext sslContext = - (JdkSslContext) SslContextUtil.toSslContext(clientTlsSpec); - return new JwtsSigningKeyResolver(ztsUri + oauth2KeysPath, sslContext.context(), proxyUriStr); - } - - private static ClientTlsSpec toTlsSpec(TlsProvider tlsProvider, boolean allowUnsafeCiphers) { - final ClientTlsSpecBuilder builder = ClientTlsSpec.builder(); - final TlsKeyPair tlsKeyPair = tlsProvider.keyPair("*"); - if (tlsKeyPair != null) { - builder.tlsKeyPair(tlsKeyPair); - } - final List trustedCertificates = tlsProvider.trustedCertificates("*"); - if (trustedCertificates != null) { - builder.trustedCertificates(trustedCertificates); - } - builder.engineType(TlsEngineType.JDK); - builder.alpnProtocols(SslContextUtil.DEFAULT_ALPN_PROTOCOLS); - builder.allowUnsafeCiphers(allowUnsafeCiphers); - return builder.build(); + final String proxyUrl = System.getProperty(ZpeConsts.ZPE_PROP_JWK_PROXY_URI); + accessSignKeyResolver = new JwtsSigningKeyResolver(serverUrl, sslContext, proxyUrl); } /** diff --git a/athenz/src/main/java/com/yahoo/athenz/auth/token/jwts/ArmeriaJwtsSigningKeyResolver.java b/athenz/src/main/java/com/yahoo/athenz/auth/token/jwts/ArmeriaJwtsSigningKeyResolver.java new file mode 100644 index 00000000000..fadd73bf8e0 --- /dev/null +++ b/athenz/src/main/java/com/yahoo/athenz/auth/token/jwts/ArmeriaJwtsSigningKeyResolver.java @@ -0,0 +1,67 @@ +/* + * Copyright 2026 LY Corporation + * + * LY Corporation licenses this file to you under the Apache License, + * version 2.0 (the "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at: + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT + * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the + * License for the specific language governing permissions and limitations + * under the License. + */ + +package com.yahoo.athenz.auth.token.jwts; + +import javax.net.ssl.SSLContext; + +import com.nimbusds.jose.util.ResourceRetriever; + +import com.linecorp.armeria.client.WebClient; +import com.linecorp.armeria.common.annotation.Nullable; +import com.linecorp.armeria.common.annotation.UnstableApi; + +/** + * A {@link JwtsSigningKeyResolver} that uses an Armeria {@link WebClient} to fetch JWKS keys + * instead of the default JDK HTTP client. This avoids the need to extract an {@link SSLContext} + * from the Armeria {@link com.linecorp.armeria.client.ClientFactory}. + * + *

Use the static factory method {@link #create(WebClient, String)} to create an instance. + */ +@UnstableApi +public final class ArmeriaJwtsSigningKeyResolver extends JwtsSigningKeyResolver { + + private static final ThreadLocal INIT_RETRIEVER = new ThreadLocal<>(); + + /** + * Creates a new {@link ArmeriaJwtsSigningKeyResolver}. + * + * @param webClient the {@link WebClient} pre-configured to connect to the ZTS server + * @param oauth2KeysPath the path to the OAuth2 keys endpoint (e.g., "/oauth2/keys?rfc=true") + */ + public static ArmeriaJwtsSigningKeyResolver create(WebClient webClient, String oauth2KeysPath) { + INIT_RETRIEVER.set(new ArmeriaResourceRetriever(webClient, oauth2KeysPath)); + try { + final String jwksUri = webClient.uri() + oauth2KeysPath; + return new ArmeriaJwtsSigningKeyResolver(jwksUri); + } finally { + INIT_RETRIEVER.remove(); + } + } + + private ArmeriaJwtsSigningKeyResolver(String jwksUri) { + super(jwksUri, null, true); + } + + @Override + ResourceRetriever getResourceRetriever(@Nullable String proxyUrl, @Nullable SSLContext sslContext) { + final ArmeriaResourceRetriever retriever = INIT_RETRIEVER.get(); + if (retriever != null) { + return retriever; + } + return super.getResourceRetriever(proxyUrl, sslContext); + } +} diff --git a/athenz/src/main/java/com/yahoo/athenz/auth/token/jwts/ArmeriaResourceRetriever.java b/athenz/src/main/java/com/yahoo/athenz/auth/token/jwts/ArmeriaResourceRetriever.java new file mode 100644 index 00000000000..4f7f8a37ee1 --- /dev/null +++ b/athenz/src/main/java/com/yahoo/athenz/auth/token/jwts/ArmeriaResourceRetriever.java @@ -0,0 +1,57 @@ +/* + * Copyright 2026 LY Corporation + * + * LY Corporation licenses this file to you under the Apache License, + * version 2.0 (the "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at: + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT + * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the + * License for the specific language governing permissions and limitations + * under the License. + */ + +package com.yahoo.athenz.auth.token.jwts; + +import java.io.IOException; +import java.net.URL; + +import com.nimbusds.jose.util.Resource; +import com.nimbusds.jose.util.ResourceRetriever; + +import com.linecorp.armeria.client.WebClient; +import com.linecorp.armeria.common.AggregatedHttpResponse; +import com.linecorp.armeria.common.HttpStatus; +import com.linecorp.armeria.common.MediaType; +import com.linecorp.armeria.common.annotation.Nullable; + +/** + * A Nimbus {@link ResourceRetriever} implementation that uses an Armeria {@link WebClient} + * to fetch JWKS resources. The {@code url} parameter passed to {@link #retrieveResource(URL)} + * is ignored because the {@link WebClient} is already configured to target the correct ZTS host. + */ +final class ArmeriaResourceRetriever implements ResourceRetriever { + + private final WebClient webClient; + private final String path; + + ArmeriaResourceRetriever(WebClient webClient, String path) { + this.webClient = webClient; + this.path = path; + } + + @Override + public Resource retrieveResource(URL url) throws IOException { + final AggregatedHttpResponse res = webClient.blocking().get(path); + if (!res.status().equals(HttpStatus.OK)) { + throw new IOException("Failed to retrieve JWKS resource: " + res.status()); + } + @Nullable + final MediaType contentType = res.headers().contentType(); + final String contentTypeStr = contentType != null ? contentType.toString() : "application/json"; + return new Resource(res.contentUtf8(), contentTypeStr); + } +} diff --git a/athenz/src/main/java/com/yahoo/athenz/auth/token/jwts/package-info.java b/athenz/src/main/java/com/yahoo/athenz/auth/token/jwts/package-info.java new file mode 100644 index 00000000000..bfc4d435e73 --- /dev/null +++ b/athenz/src/main/java/com/yahoo/athenz/auth/token/jwts/package-info.java @@ -0,0 +1,25 @@ +/* + * Copyright 2026 LY Corporation + * + * LY Corporation licenses this file to you under the Apache License, + * version 2.0 (the "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at: + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT + * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the + * License for the specific language governing permissions and limitations + * under the License. + */ + +/** + * Provides workarounds for Athenz integration. + */ +@UnstableApi +@NonNullByDefault +package com.yahoo.athenz.auth.token.jwts; + +import com.linecorp.armeria.common.annotation.NonNullByDefault; +import com.linecorp.armeria.common.annotation.UnstableApi; diff --git a/athenz/src/test/java/com/linecorp/armeria/client/athenz/ArmeriaJwtsSigningKeyResolverTest.java b/athenz/src/test/java/com/linecorp/armeria/client/athenz/ArmeriaJwtsSigningKeyResolverTest.java new file mode 100644 index 00000000000..59895b93cf7 --- /dev/null +++ b/athenz/src/test/java/com/linecorp/armeria/client/athenz/ArmeriaJwtsSigningKeyResolverTest.java @@ -0,0 +1,94 @@ +/* + * Copyright 2026 LY Corporation + * + * LY Corporation licenses this file to you under the Apache License, + * version 2.0 (the "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at: + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT + * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the + * License for the specific language governing permissions and limitations + * under the License. + */ + +package com.linecorp.armeria.client.athenz; + +import static org.assertj.core.api.Assertions.assertThat; + +import java.security.interfaces.RSAPublicKey; + +import org.junit.jupiter.api.Order; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; + +import com.nimbusds.jose.jwk.JWKSet; +import com.nimbusds.jose.jwk.RSAKey; + +import com.linecorp.armeria.client.WebClient; +import com.linecorp.armeria.common.HttpResponse; +import com.linecorp.armeria.common.HttpStatus; +import com.linecorp.armeria.common.MediaType; +import com.linecorp.armeria.server.ServerBuilder; +import com.linecorp.armeria.testing.junit5.server.SelfSignedCertificateExtension; +import com.linecorp.armeria.testing.junit5.server.ServerExtension; + +import com.yahoo.athenz.auth.token.jwts.ArmeriaJwtsSigningKeyResolver; + +class ArmeriaJwtsSigningKeyResolverTest { + + private static final String KEY_ID = "test-key-1"; + + @Order(0) + @RegisterExtension + static final SelfSignedCertificateExtension cert = new SelfSignedCertificateExtension(); + + @Order(1) + @RegisterExtension + static ServerExtension server = new ServerExtension() { + @Override + protected void configure(ServerBuilder sb) { + final RSAKey rsaJwk = new RSAKey.Builder( + (RSAPublicKey) cert.certificate().getPublicKey()) + .keyID(KEY_ID) + .build(); + final String jwksJson = new JWKSet(rsaJwk).toString(); + + sb.service("/oauth2/keys", (ctx, req) -> + HttpResponse.of(HttpStatus.OK, MediaType.JSON, jwksJson)); + sb.service("/oauth2/error", (ctx, req) -> + HttpResponse.of(HttpStatus.INTERNAL_SERVER_ERROR)); + } + }; + + @Test + void shouldResolvePublicKey() { + final WebClient webClient = WebClient.of(server.httpUri()); + final ArmeriaJwtsSigningKeyResolver resolver = + ArmeriaJwtsSigningKeyResolver.create(webClient, "/oauth2/keys"); + + assertThat(resolver.getPublicKey(KEY_ID)) + .isEqualTo(cert.certificate().getPublicKey()); + } + + @Test + void shouldReturnNullForUnknownKeyId() { + final WebClient webClient = WebClient.of(server.httpUri()); + final ArmeriaJwtsSigningKeyResolver resolver = + ArmeriaJwtsSigningKeyResolver.create(webClient, "/oauth2/keys"); + + assertThat(resolver.getPublicKey("nonexistent-key")).isNull(); + } + + @Test + void shouldHandleServerError() { + final WebClient webClient = WebClient.of(server.httpUri()); + final ArmeriaJwtsSigningKeyResolver resolver = + ArmeriaJwtsSigningKeyResolver.create(webClient, "/oauth2/error"); + + // The resolver logs errors internally but does not throw. + assertThat(resolver.getPublicKey(KEY_ID)).isNull(); + } +} From 37622c629be672b0eb0411e196c9094c85a7cfa3 Mon Sep 17 00:00:00 2001 From: jrhee17 Date: Thu, 9 Jul 2026 10:58:05 +0900 Subject: [PATCH 2/3] lint --- .../client/athenz/ArmeriaJwtsSigningKeyResolverTest.java | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/athenz/src/test/java/com/linecorp/armeria/client/athenz/ArmeriaJwtsSigningKeyResolverTest.java b/athenz/src/test/java/com/linecorp/armeria/client/athenz/ArmeriaJwtsSigningKeyResolverTest.java index 59895b93cf7..a9647fd5d4f 100644 --- a/athenz/src/test/java/com/linecorp/armeria/client/athenz/ArmeriaJwtsSigningKeyResolverTest.java +++ b/athenz/src/test/java/com/linecorp/armeria/client/athenz/ArmeriaJwtsSigningKeyResolverTest.java @@ -26,6 +26,7 @@ import com.nimbusds.jose.jwk.JWKSet; import com.nimbusds.jose.jwk.RSAKey; +import com.yahoo.athenz.auth.token.jwts.ArmeriaJwtsSigningKeyResolver; import com.linecorp.armeria.client.WebClient; import com.linecorp.armeria.common.HttpResponse; @@ -35,8 +36,6 @@ import com.linecorp.armeria.testing.junit5.server.SelfSignedCertificateExtension; import com.linecorp.armeria.testing.junit5.server.ServerExtension; -import com.yahoo.athenz.auth.token.jwts.ArmeriaJwtsSigningKeyResolver; - class ArmeriaJwtsSigningKeyResolverTest { private static final String KEY_ID = "test-key-1"; From bf1e1ca9fcab639001d5b99b356aa7af3da0d33c Mon Sep 17 00:00:00 2001 From: jrhee17 Date: Thu, 9 Jul 2026 15:34:51 +0900 Subject: [PATCH 3/3] address comments by @minwoox and @ikhoon --- .../armeria/client/athenz/DefaultZtsBaseClient.java | 6 +++--- .../com/linecorp/armeria/client/athenz/ZtsBaseClient.java | 5 +---- .../auth/token/jwts/ArmeriaJwtsSigningKeyResolver.java | 7 +++++++ 3 files changed, 11 insertions(+), 7 deletions(-) diff --git a/athenz/src/main/java/com/linecorp/armeria/client/athenz/DefaultZtsBaseClient.java b/athenz/src/main/java/com/linecorp/armeria/client/athenz/DefaultZtsBaseClient.java index 2963f8ff3f6..78af2708da3 100644 --- a/athenz/src/main/java/com/linecorp/armeria/client/athenz/DefaultZtsBaseClient.java +++ b/athenz/src/main/java/com/linecorp/armeria/client/athenz/DefaultZtsBaseClient.java @@ -97,9 +97,9 @@ public WebClient webClient(Consumer configurer) { clientBuilder.factory(clientFactory); - if (LoggerFactory.getLogger(DefaultZtsBaseClient.class).isTraceEnabled()) { + if (LoggerFactory.getLogger(ZtsBaseClient.class).isTraceEnabled()) { final LogWriter logWriter = LogWriter.builder() - .logger(DefaultZtsBaseClient.class.getName()) + .logger(ZtsBaseClient.class.getName()) .requestLogLevel(LogLevel.TRACE) .successfulResponseLogLevel(LogLevel.TRACE) .build(); @@ -128,7 +128,7 @@ public void removeTlsKeyPairListener(Consumer listener) { @Override public void close() { - defaultWebClient.options().factory().closeAsync(); + clientFactory.closeAsync(); } private static final class TlsKeyPairListener extends AbstractListenable { diff --git a/athenz/src/main/java/com/linecorp/armeria/client/athenz/ZtsBaseClient.java b/athenz/src/main/java/com/linecorp/armeria/client/athenz/ZtsBaseClient.java index 53f5f7feb67..dc11a40d9ee 100644 --- a/athenz/src/main/java/com/linecorp/armeria/client/athenz/ZtsBaseClient.java +++ b/athenz/src/main/java/com/linecorp/armeria/client/athenz/ZtsBaseClient.java @@ -112,10 +112,7 @@ default ClientFactory clientFactory() { /** * Returns a new {@link WebClient} that can connect to the ZTS server with the specified configurer. */ - default WebClient webClient(Consumer configurer) { - requireNonNull(configurer, "configurer"); - return webClient(); - } + WebClient webClient(Consumer configurer); /** * Adds a listener that will be notified when the {@link TlsKeyPair} is updated. diff --git a/athenz/src/main/java/com/yahoo/athenz/auth/token/jwts/ArmeriaJwtsSigningKeyResolver.java b/athenz/src/main/java/com/yahoo/athenz/auth/token/jwts/ArmeriaJwtsSigningKeyResolver.java index fadd73bf8e0..33da259b821 100644 --- a/athenz/src/main/java/com/yahoo/athenz/auth/token/jwts/ArmeriaJwtsSigningKeyResolver.java +++ b/athenz/src/main/java/com/yahoo/athenz/auth/token/jwts/ArmeriaJwtsSigningKeyResolver.java @@ -18,6 +18,9 @@ import javax.net.ssl.SSLContext; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + import com.nimbusds.jose.util.ResourceRetriever; import com.linecorp.armeria.client.WebClient; @@ -34,6 +37,8 @@ @UnstableApi public final class ArmeriaJwtsSigningKeyResolver extends JwtsSigningKeyResolver { + private static final Logger logger = LoggerFactory.getLogger(ArmeriaJwtsSigningKeyResolver.class); + private static final ThreadLocal INIT_RETRIEVER = new ThreadLocal<>(); /** @@ -62,6 +67,8 @@ ResourceRetriever getResourceRetriever(@Nullable String proxyUrl, @Nullable SSLC if (retriever != null) { return retriever; } + logger.warn("ArmeriaResourceRetriever is not available. " + + "Falling back to the default resource retriever."); return super.getResourceRetriever(proxyUrl, sslContext); } }