Add Athenz integration module

athenz/build.gradle

@@ -0,0 +1,26 @@
+/*
+ * Copyright 2025 LINE Corporation
+ *
+ * LINE 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.
+ */
+
+dependencies {
+    implementation project(":oauth2")
+    implementation libs.athenz.zts.client
+    implementation libs.athenz.zpe.client
+    implementation libs.caffeine
+
+    testImplementation libs.athenz.zms.client
+    testImplementation libs.jwt
+    testImplementation libs.testcontainers.junit.jupiter
+}

athenz/src/main/java/com/linecorp/armeria/client/athenz/AccessTokenClient.java

@@ -0,0 +1,113 @@
+/*
+ * Copyright 2025 LINE Corporation
+ *
+ * LINE 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 com.linecorp.armeria.client.athenz.RoleTokenClient.ROLE_JOINER;
+
+import java.time.Duration;
+import java.time.Instant;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import com.google.common.collect.ImmutableList;
+import com.yahoo.athenz.auth.AuthorityConsts;
+
+import com.linecorp.armeria.client.auth.oauth2.AccessTokenRequest;
+import com.linecorp.armeria.client.auth.oauth2.OAuth2AuthorizationGrant;
+import com.linecorp.armeria.common.HttpHeadersBuilder;
+import com.linecorp.armeria.common.QueryParamsBuilder;
+import com.linecorp.armeria.common.athenz.AccessDeniedException;
+import com.linecorp.armeria.common.auth.oauth2.ClientAuthentication;
+import com.linecorp.armeria.common.auth.oauth2.GrantedOAuth2AccessToken;
+import com.linecorp.armeria.common.util.Exceptions;
+
+final class AccessTokenClient implements TokenClient {
+
+    private final AtomicBoolean tlsKeyPairUpdated = new AtomicBoolean();
+
+    private final long refreshBeforeMillis;
+    private final String domainName;
+    private final List<String> roleNames;
+    private final OAuth2AuthorizationGrant authorizationGrant;
+
+    AccessTokenClient(ZtsBaseClient ztsBaseClient, String domainName, List<String> roleNames,
+                      Duration refreshBefore) {
+        refreshBeforeMillis = refreshBefore.toMillis();
+        this.domainName = domainName;
+        this.roleNames = roleNames;
+
+        ztsBaseClient.addTlsKeyPairListener(tlsKeyPair -> tlsKeyPairUpdated.set(true));
+
+        // Scope syntax:
+        // - <domain-name>:domain
+        // - <domain-name>:role.<role-name>
+        // https://github.com/AthenZ/athenz/blob/5e064414224eca025c7a4ae1df5b5eb381e71a16/clients/java/zts/src/main/java/com/yahoo/athenz/zts/ZTSClient.java#L1446
+        final ImmutableList.Builder<String> scopeBuilder = ImmutableList.builder();
+        if (roleNames.isEmpty()) {
+            scopeBuilder.add(domainName + ":domain");
+        } else {
+            for (String role : roleNames) {
+                scopeBuilder.add(domainName + AuthorityConsts.ROLE_SEP + role);
+            }
+        }
+        final AccessTokenRequest tokenRequest = AccessTokenRequest.ofClientCredentials(
+                NoopClientAuthentication.INSTANCE, scopeBuilder.build());
+        authorizationGrant = OAuth2AuthorizationGrant.builder(ztsBaseClient.webClient(),
+                                                              "/oauth2/token")
+                                                     .accessTokenRequest(tokenRequest)
+                                                     .refreshIf(this::shouldRefreshToken)
+                                                     .build();
+    }
+
+    private boolean shouldRefreshToken(GrantedOAuth2AccessToken token) {
+        if (tlsKeyPairUpdated.compareAndSet(true, false)) {
+            // If the TLS key pair is updated, we need to refresh the token.
+            return true;
+        }
+        final Duration expiresIn = token.expiresIn();
+        if (expiresIn == null) {
+            return false;
+        }
+
+        return !token.isValid(Instant.now().plusMillis(refreshBeforeMillis));
+    }
+
+    @Override
+    public CompletableFuture<String> getToken() {
+        return authorizationGrant.getAccessToken().handle((token, cause) -> {
+            if (cause != null) {
+                cause = Exceptions.peel(cause);
+                throw new AccessDeniedException("Failed to obtain an Athenz access token. (domain: " +
+                                                domainName + ", roles: " + ROLE_JOINER.join(roleNames) + ')',
+                                                cause);
+            }
+            return token.accessToken();
+        }).toCompletableFuture();
+    }
+
+    private enum NoopClientAuthentication implements ClientAuthentication {
+
+        INSTANCE;
+
+        @Override
+        public void addAsHeaders(HttpHeadersBuilder headersBuilder) {}
+
+        @Override
+        public void addAsBodyParams(QueryParamsBuilder formBuilder) {}
+    }
+}

athenz/src/main/java/com/linecorp/armeria/client/athenz/AthenzClient.java

@@ -0,0 +1,178 @@
+/*
+ * Copyright 2025 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 com.linecorp.armeria.internal.common.athenz.AthenzHeaderNames.YAHOO_ROLE_AUTH;
+import static java.util.Objects.requireNonNull;
+
+import java.time.Duration;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.function.Function;
+
+import com.google.common.collect.ImmutableList;
+
+import com.linecorp.armeria.client.ClientRequestContext;
+import com.linecorp.armeria.client.HttpClient;
+import com.linecorp.armeria.client.SimpleDecoratingHttpClient;
+import com.linecorp.armeria.common.HttpHeaderNames;
+import com.linecorp.armeria.common.HttpRequest;
+import com.linecorp.armeria.common.HttpResponse;
+import com.linecorp.armeria.common.RequestHeadersBuilder;
+import com.linecorp.armeria.common.annotation.UnstableApi;
+import com.linecorp.armeria.common.athenz.TokenType;
+import com.linecorp.armeria.common.util.Exceptions;
+
+/**
+ * An {@link HttpClient} that adds an Athenz token to the request headers.
+ * {@link TokenType#ACCESS_TOKEN} and {@link TokenType#ROLE_TOKEN} are supported.
+ *
+ * <p>The acquired token is cached and automatically refreshed before it expires based on the specified
+ * duration. If not specified, the default refresh duration is 10 minutes before the token expires.
+ *
+ * <p>Example:
+ * <pre>{@code
+ * import com.linecorp.armeria.client.athenz.ZtsBaseClient;
+ * import com.linecorp.armeria.client.athenz.AthenzClient;
+ *
+ * ZtsBaseClient ztsBaseClient =
+ *   ZtsBaseClient
+ *     .builder("https://athenz.example.com:8443/zts/v1")
+ *     .keyPair("/var/lib/athenz/service.key.pem", "/var/lib/athenz/service.cert.pem")
+ *     .build();
+ *
+ * WebClient
+ *   .builder()
+ *   .decorator(AthenzClient.newDecorator(ztsBaseClient, "my-domain",
+ *                                        TokenType.ROLE_TOKEN)
+ *   ...
+ *   .build();
+ * }</pre>
+ */
+@UnstableApi
+public final class AthenzClient extends SimpleDecoratingHttpClient {
+
+    private static final Duration DEFAULT_REFRESH_BEFORE = Duration.ofMinutes(10);
+
+    /**
+     * Returns a new {@link HttpClient} decorator that obtains an Athenz token for the specified domain and
+     * adds it to the request headers.
+     *
+     * @param ztsBaseClient the ZTS base client to use to communicate with the ZTS server
+     * @param domainName the Athenz domain name
+     * @param tokenType the type of Athenz token to obtain
+     */
+    public static Function<HttpClient, AthenzClient> newDecorator(ZtsBaseClient ztsBaseClient,
+                                                                  String domainName, TokenType tokenType) {
+        return newDecorator(ztsBaseClient, domainName, ImmutableList.of(), tokenType);
+    }
+
+    /**
+     * Returns a new {@link HttpClient} decorator that obtains an Athenz token for the specified domain and
+     * role name, and adds it to the request headers.
+     *
+     * @param ztsBaseClient the ZTS base client to use to communicate with the ZTS server
+     * @param domainName the Athenz domain name
+     * @param roleName the Athenz role name
+     * @param tokenType the type of Athenz token to obtain
+     */
+    public static Function<HttpClient, AthenzClient> newDecorator(ZtsBaseClient ztsBaseClient,
+                                                                  String domainName, String roleName,
+                                                                  TokenType tokenType) {
+        return newDecorator(ztsBaseClient, domainName, ImmutableList.of(roleName), tokenType);
+    }
+
+    /**
+     * Returns a new {@link HttpClient} decorator that obtains an Athenz token for the specified domain and
+     * role names, and adds it to the request headers.
+     *
+     * @param ztsBaseClient the ZTS base client to use to communicate with the ZTS server
+     * @param domainName the Athenz domain name
+     * @param roleNames the list of Athenz role names
+     * @param tokenType the type of Athenz token to obtain
+     */
+    public static Function<HttpClient, AthenzClient> newDecorator(ZtsBaseClient ztsBaseClient,
+                                                                  String domainName, List<String> roleNames,
+                                                                  TokenType tokenType) {
+        return newDecorator(ztsBaseClient, domainName, roleNames, tokenType, DEFAULT_REFRESH_BEFORE);
+    }
+
+    /**
+     * Returns a new {@link HttpClient} decorator that obtains an Athenz token for the specified domain and
+     * role names, and adds it to the request headers.
+     *
+     * @param ztsBaseClient the ZTS base client to use to communicate with the ZTS server
+     * @param domainName the Athenz domain name
+     * @param roleNames the list of Athenz role names
+     * @param tokenType the type of Athenz token to obtain
+     * @param refreshBefore the duration before the token expires to refresh it
+     */
+    public static Function<HttpClient, AthenzClient> newDecorator(ZtsBaseClient ztsBaseClient,
+                                                                  String domainName, List<String> roleNames,
+                                                                  TokenType tokenType, Duration refreshBefore) {
+        requireNonNull(ztsBaseClient, "ztsBaseClient");
+        requireNonNull(domainName, "domainName");
+        requireNonNull(roleNames, "roleNames");
+        final ImmutableList<String> roleNames0 = ImmutableList.copyOf(roleNames);
+        requireNonNull(tokenType, "tokenType");
+        requireNonNull(refreshBefore, "refreshBefore");
+        return delegate -> new AthenzClient(delegate, ztsBaseClient, domainName, roleNames0,
+                                            tokenType, refreshBefore);
+    }
+
+    private final TokenType tokenType;
+    private final TokenClient tokenClient;
+
+    private AthenzClient(HttpClient delegate, ZtsBaseClient ztsBaseClient, String domainName,
+                         List<String> roleNames, TokenType tokenType, Duration refreshBefore) {
+        super(delegate);
+        this.tokenType = tokenType;
+        switch (tokenType) {
+            case ROLE_TOKEN:
+                tokenClient = new RoleTokenClient(ztsBaseClient, domainName, roleNames, refreshBefore);
+                break;
+            case ACCESS_TOKEN:
+                tokenClient = new AccessTokenClient(ztsBaseClient, domainName, roleNames, refreshBefore);
+                break;
+            default:
+                throw new Error("unknown auth type: " + tokenType);
+        }
+    }
+
+    @Override
+    public HttpResponse execute(ClientRequestContext ctx, HttpRequest req) throws Exception {
+        final CompletableFuture<HttpResponse> future = tokenClient.getToken().thenApply(token -> {
+            final HttpRequest newReq = req.mapHeaders(headers -> {
+                final RequestHeadersBuilder builder = headers.toBuilder();
+                if (tokenType == TokenType.ROLE_TOKEN) {
+                    builder.set(YAHOO_ROLE_AUTH, token);
+                } else {
+                    builder.set(HttpHeaderNames.AUTHORIZATION, "Bearer " + token);
+                }
+                return builder.build();
+            });
+            ctx.updateRequest(newReq);
+            try {
+                return unwrap().execute(ctx, newReq);
+            } catch (Exception e) {
+                return Exceptions.throwUnsafely(e);
+            }
+        });
+
+        return HttpResponse.of(future);
+    }
+}