Skip to content

Commit eb38acf

Browse files
ikhoonikhoon
committed
Allow specifying JWKS path in AthenzService (#6393)
Motivation: In recent Athenz versions, access token validation only works when the curve names in JWKS follow the RFC specification. Therefore, I propose changing the default JWKS endpoint to return RFC-compliant values, while still allowing the default to be overridden in `AthenzServiceBuilder` References: - AthenZ/athenz#2707 - https://bitbucket.org/connect2id/nimbus-jose-jwt/issues/599/add-prime256v1-to-the-supported-curves-in Motifications: - Use `/oauth2/keys?rfc=true` to fetch JWKS from the ZTS server. - Add the same property in `AthenzConfig` for Spring Boot integration. Result: Fixed Athenz access token validation to work with RFC-compliant JWKS.
1 parent cea0ab2 commit eb38acf

File tree

9 files changed

+287
-19
lines changed

9 files changed

+287
-19
lines changed

‎athenz/src/main/java/com/linecorp/armeria/internal/server/athenz/AthenzServiceDecoratorFactoryProvider.java‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -25,11 +25,13 @@
2525

2626
import com.linecorp.armeria.client.athenz.ZtsBaseClient;
2727
import com.linecorp.armeria.client.athenz.ZtsBaseClientBuilder;
28+
import com.linecorp.armeria.common.annotation.Nullable;
2829
import com.linecorp.armeria.server.Server;
2930
import com.linecorp.armeria.server.ServerBuilder;
3031
import com.linecorp.armeria.server.ServerListenerAdapter;
3132
import com.linecorp.armeria.server.athenz.AthenzPolicyConfig;
3233
import com.linecorp.armeria.server.athenz.AthenzServiceDecoratorFactory;
34+
import com.linecorp.armeria.server.athenz.AthenzServiceDecoratorFactoryBuilder;
3335

3436
/**
3537
* A helper class to create an {@link AthenzServiceDecoratorFactory} instance.
@@ -39,7 +41,9 @@ public final class AthenzServiceDecoratorFactoryProvider {
3941

4042
public static AthenzServiceDecoratorFactory create(
4143
ServerBuilder sb, URI ztsUri, File athenzPrivateKey, File athenzPublicKey,
42-
URI proxyUri, File athenzCaCert, List<String> domains, boolean jwsPolicySupport,
44+
@Nullable URI proxyUri, @Nullable File athenzCaCert, @Nullable String oauth2KeysPath,
45+
List<String> domains,
46+
boolean jwsPolicySupport,
4347
Duration policyRefreshInterval) {
4448

4549
final ZtsBaseClientBuilder clientBuilder =
@@ -62,8 +66,13 @@ public void serverStopped(Server server) throws Exception {
6266
ztsBaseClient.close();
6367
}
6468
});
65-
return AthenzServiceDecoratorFactory
66-
.builder(ztsBaseClient)
69+
final AthenzServiceDecoratorFactoryBuilder factoryBuilder =
70+
AthenzServiceDecoratorFactory
71+
.builder(ztsBaseClient);
72+
if (oauth2KeysPath != null) {
73+
factoryBuilder.oauth2KeysPath(oauth2KeysPath);
74+
}
75+
return factoryBuilder
6776
.policyConfig(athenzPolicyConfig)
6877
.build();
6978
}

‎athenz/src/main/java/com/linecorp/armeria/server/athenz/AbstractAthenzServiceBuilder.java‎

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,9 +20,11 @@
2020
import static com.google.common.base.Preconditions.checkState;
2121
import static java.util.Objects.requireNonNull;
2222

23+
import java.net.URI;
2324
import java.time.Duration;
2425

2526
import com.yahoo.athenz.zpe.ZpeClient;
27+
import com.yahoo.athenz.zpe.ZpeConsts;
2628
import com.yahoo.athenz.zpe.pkey.PublicKeyStore;
2729

2830
import com.linecorp.armeria.client.athenz.ZtsBaseClient;
@@ -36,13 +38,15 @@
3638
public abstract class AbstractAthenzServiceBuilder<SELF extends AbstractAthenzServiceBuilder<SELF>> {
3739

3840
private static final Duration DEFAULT_OAUTH2_KEYS_REFRESH_INTERVAL = Duration.ofHours(1);
41+
private static final String DEFAULT_OAUTH2_KEY_PATH = "/oauth2/keys?rfc=true";
3942
private static final int MAX_TOKEN_CACHE_SIZE = 10240;
4043

4144
private Duration oauth2KeysRefreshInterval = DEFAULT_OAUTH2_KEYS_REFRESH_INTERVAL;
4245

4346
private final ZtsBaseClient ztsBaseClient;
4447
@Nullable
4548
private AthenzPolicyConfig policyConfig;
49+
private String oauth2KeysPath = DEFAULT_OAUTH2_KEY_PATH;
4650
private int maxTokenCacheSize = MAX_TOKEN_CACHE_SIZE;
4751

4852
AbstractAthenzServiceBuilder(ZtsBaseClient ztsBaseClient) {
@@ -70,6 +74,22 @@ public SELF oauth2KeysRefreshInterval(Duration oauth2KeysRefreshInterval) {
7074
return self();
7175
}
7276

77+
/**
78+
* Sets the URI to fetch OAuth2 keys (JWK) from the ZTS server. The path is relative to the base {@link URI}
79+
* of {@link ZtsBaseClient}.
80+
* If not set, defaults to {@value #DEFAULT_OAUTH2_KEY_PATH}.
81+
* If {@value ZpeConsts#ZPE_PROP_JWK_URI} is set as a system property, this option is ignored.
82+
*/
83+
public SELF oauth2KeysPath(String oauth2KeysPath) {
84+
requireNonNull(oauth2KeysPath, "oauth2KeysPath");
85+
checkArgument(!oauth2KeysPath.isEmpty(), "oauth2KeysPath must not be empty");
86+
checkArgument(oauth2KeysPath.charAt(0) == '/',
87+
"oauth2KeysPath: %s (expected: a relative path starting with '/')",
88+
oauth2KeysPath);
89+
this.oauth2KeysPath = oauth2KeysPath;
90+
return self();
91+
}
92+
7393
/**
7494
* Set the limit of role and access tokens that are cached to improve the performance of validating
7595
* signatures since the tokens must be re-used by clients until they're about to be expired.
@@ -83,13 +103,13 @@ public SELF maxTokenCacheSize(int maxTokenCacheSize) {
83103

84104
MinifiedAuthZpeClient buildAuthZpeClient() {
85105
checkState(policyConfig != null, "policyConfig must be set before building the service");
86-
final PublicKeyStore publicKeyStore = new AthenzPublicKeyProvider(ztsBaseClient,
87-
oauth2KeysRefreshInterval);
106+
final PublicKeyStore publicKeyStore = new AthenzPublicKeyProvider(
107+
ztsBaseClient, oauth2KeysRefreshInterval, oauth2KeysPath);
88108
final ZpeClient zpeClient = new AthenzPolicyClient(ztsBaseClient, policyConfig, publicKeyStore,
89109
maxTokenCacheSize);
90110
// NB: zpeClient.init() will block until the initial policy data is loaded.
91111
zpeClient.init(null);
92-
return new MinifiedAuthZpeClient(ztsBaseClient, publicKeyStore, zpeClient);
112+
return new MinifiedAuthZpeClient(ztsBaseClient, publicKeyStore, zpeClient, oauth2KeysPath);
93113
}
94114

95115
@SuppressWarnings("unchecked")

‎athenz/src/main/java/com/linecorp/armeria/server/athenz/AthenzPublicKeyProvider.java‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -51,13 +51,15 @@ final class AthenzPublicKeyProvider implements PublicKeyStore {
5151
private final long minRetryInterval;
5252
private final AsyncLoader<Map<String, CompletableFuture<PublicKey>>> ztsKeyLoader;
5353
private final AsyncLoader<Map<String, CompletableFuture<PublicKey>>> zmsKeyLoader;
54+
private final String oauth2KeyPath;
5455
private volatile long lastReloadZtsJwkTime;
5556
private volatile long lastReloadZmsJwkTime;
5657

57-
AthenzPublicKeyProvider(ZtsBaseClient ztsBaseClient, Duration refreshInterval) {
58+
AthenzPublicKeyProvider(ZtsBaseClient ztsBaseClient, Duration refreshInterval, String oauth2KeyPath) {
5859
// TODO(ikhoon): Make minRetryInterval configurable.
5960
minRetryInterval = refreshInterval.toMillis() / 4;
6061
webClient = ztsBaseClient.webClient();
62+
this.oauth2KeyPath = oauth2KeyPath;
6163
ztsKeyLoader = AsyncLoader.<Map<String, CompletableFuture<PublicKey>>>builder(k -> fetchZtsKeys())
6264
.name("athenz-zts-key-loader")
6365
.refreshAfterLoad(refreshInterval)
@@ -108,7 +110,7 @@ private CompletableFuture<PublicKey> getKey(
108110
private CompletableFuture<Map<String, CompletableFuture<PublicKey>>> fetchZtsKeys() {
109111
return webClient
110112
.prepare()
111-
.get("/oauth2/keys")
113+
.get(oauth2KeyPath)
112114
.asJson(Keys.class)
113115
.execute()
114116
.thenApply(res -> {

‎athenz/src/main/java/com/linecorp/armeria/server/athenz/MinifiedAuthZpeClient.java‎

Lines changed: 12 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -180,7 +180,8 @@ public String toString() {
180180
}
181181
}
182182

183-
MinifiedAuthZpeClient(ZtsBaseClient ztsBaseClient, PublicKeyStore publicKeyStore, ZpeClient zpeClt) {
183+
MinifiedAuthZpeClient(ZtsBaseClient ztsBaseClient, PublicKeyStore publicKeyStore, ZpeClient zpeClt,
184+
String oauth2KeysPath) {
184185
this.publicKeyStore = publicKeyStore;
185186
this.zpeClt = zpeClt;
186187

@@ -190,21 +191,21 @@ public String toString() {
190191

191192
// initialize the access token signing key resolver
192193

193-
initializeAccessTokenSignKeyResolver(ztsBaseClient);
194+
initializeAccessTokenSignKeyResolver(ztsBaseClient, oauth2KeysPath);
194195

195196
// save the last zts api call time, and the allowed interval between api calls
196197

197198
setMillisBetweenZtsCalls(Long.parseLong(
198199
System.getProperty(ZPE_PROP_MILLIS_BETWEEN_ZTS_CALLS, Long.toString(30 * 1000 * 60))));
199200
}
200201

201-
private void initializeAccessTokenSignKeyResolver(ZtsBaseClient ztsBaseClient) {
202+
private void initializeAccessTokenSignKeyResolver(ZtsBaseClient ztsBaseClient, String oauth2KeysPath) {
202203
final String serverUrl = System.getProperty(ZpeConsts.ZPE_PROP_JWK_URI);
203204
if (serverUrl == null || serverUrl.isEmpty()) {
204-
accessSignKeyResolver = newDefaultJwtsSigningKeyResolver(ztsBaseClient);
205+
accessSignKeyResolver = newDefaultJwtsSigningKeyResolver(ztsBaseClient, oauth2KeysPath);
205206
ztsBaseClient.addTlsKeyPairListener(tlsKeyPair -> {
206207
// Refresh the JwtsSigningKeyResolver when the TLS key pair changes
207-
accessSignKeyResolver = newDefaultJwtsSigningKeyResolver(ztsBaseClient);
208+
accessSignKeyResolver = newDefaultJwtsSigningKeyResolver(ztsBaseClient, oauth2KeysPath);
208209
});
209210
return;
210211
}
@@ -222,13 +223,15 @@ private void initializeAccessTokenSignKeyResolver(ZtsBaseClient ztsBaseClient) {
222223
logger.warn("Unable to initialize key refresher: {}", ex.getMessage());
223224
}
224225
}
225-
accessSignKeyResolver = new JwtsSigningKeyResolver(serverUrl, sslContext, null);
226+
final URI proxyUri = ztsBaseClient.proxyUri();
227+
accessSignKeyResolver = new JwtsSigningKeyResolver(serverUrl, sslContext,
228+
proxyUri != null ? proxyUri.toString() : null);
226229
}
227230

228-
private static JwtsSigningKeyResolver newDefaultJwtsSigningKeyResolver(ZtsBaseClient ztsBaseClient) {
231+
private static JwtsSigningKeyResolver newDefaultJwtsSigningKeyResolver(ZtsBaseClient ztsBaseClient,
232+
String oauth2KeysPath) {
229233
final URI ztsUri = ztsBaseClient.ztsUri();
230234
final URI proxyUri = ztsBaseClient.proxyUri();
231-
logger.debug("No JWK URI specified, using {}", ztsUri);
232235
String proxyUriStr = null;
233236
if (proxyUri != null) {
234237
proxyUriStr = proxyUri.toString();
@@ -239,7 +242,7 @@ private static JwtsSigningKeyResolver newDefaultJwtsSigningKeyResolver(ZtsBaseCl
239242
null, clientFactory.meterRegistry());
240243
final JdkSslContext sslContext = (JdkSslContext) sslContextFactory.getOrCreate(SslContextMode.CLIENT,
241244
"*");
242-
return new JwtsSigningKeyResolver(ztsUri + "/oauth2/keys", sslContext.context(), proxyUriStr);
245+
return new JwtsSigningKeyResolver(ztsUri + oauth2KeysPath, sslContext.context(), proxyUriStr);
243246
}
244247

245248
/**

‎athenz/src/test/java/com/linecorp/armeria/server/athenz/AthenzPolicyLoaderTest.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,8 @@ class AthenzPolicyLoaderTest {
4949
void loadPolicyFiles(boolean jwsPolicySupport) throws Exception {
5050
try (ZtsBaseClient baseClient = athenzExtension.newZtsBaseClient(TEST_SERVICE)) {
5151
final PublicKeyStore publicKeyStore = new AthenzPublicKeyProvider(baseClient,
52-
Duration.ofSeconds(10));
52+
Duration.ofSeconds(10),
53+
"/oauth2/keys?rfc=true");
5354
final AthenzPolicyConfig policyConfig = new AthenzPolicyConfig(ImmutableList.of(TEST_DOMAIN_NAME),
5455
ImmutableMap.of(), jwsPolicySupport,
5556
Duration.ofSeconds(10));

0 commit comments

Comments
 (0)