Support Athenz-Role-Auth header in Athenz module (#6341)

ikhoon · web-flow · commit 5563de0edea2 · 2025-08-06T14:35:03.000+09:00
Motivation:

I realized that Athenz has replaced the `Yahoo-Role-Auth` header with
`Athenz-Role-Auth` header for sending role tokens.

Modifications:

- Raname the orignal `ROLE_TOKEN` to `YAHOO_ROLE_TOKEN`
- Add `ATHENZ_ROLE_TOKEN`

Result:

Users can use `Athenz-Role-Auth` header for sending role tokens.
diff --git a/athenz/src/main/java/com/linecorp/armeria/client/athenz/AthenzClient.java b/athenz/src/main/java/com/linecorp/armeria/client/athenz/AthenzClient.java
@@ -16,7 +16,6 @@
 
 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;
@@ -29,7 +28,6 @@
 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;
@@ -39,7 +37,7 @@
 
 /**
  * An {@link HttpClient} that adds an Athenz token to the request headers.
- * {@link TokenType#ACCESS_TOKEN} and {@link TokenType#ROLE_TOKEN} are supported.
+ * {@link TokenType#ACCESS_TOKEN} and {@link TokenType#YAHOO_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.
@@ -141,15 +139,10 @@ private AthenzClient(HttpClient delegate, ZtsBaseClient ztsBaseClient, String do
                          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);
+        if (tokenType.isRoleToken()) {
+            tokenClient = new RoleTokenClient(ztsBaseClient, domainName, roleNames, refreshBefore);
+        } else {
+            tokenClient = new AccessTokenClient(ztsBaseClient, domainName, roleNames, refreshBefore);
         }
     }
 
@@ -158,11 +151,11 @@ public HttpResponse execute(ClientRequestContext ctx, HttpRequest req) throws Ex
         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);
+                String token0 = token;
+                if (tokenType.authScheme() != null) {
+                    token0 = tokenType.authScheme() + ' ' + token0;
                 }
+                builder.set(tokenType.headerName(), token0);
                 return builder.build();
             });
             ctx.updateRequest(newReq);
diff --git a/athenz/src/main/java/com/linecorp/armeria/common/athenz/TokenType.java b/athenz/src/main/java/com/linecorp/armeria/common/athenz/TokenType.java
@@ -17,6 +17,7 @@
 package com.linecorp.armeria.common.athenz;
 
 import com.linecorp.armeria.common.HttpHeaderNames;
+import com.linecorp.armeria.common.annotation.Nullable;
 import com.linecorp.armeria.common.annotation.UnstableApi;
 import com.linecorp.armeria.internal.common.athenz.AthenzHeaderNames;
 
@@ -28,24 +29,52 @@
 @UnstableApi
 public enum TokenType {
     /**
-     * Athenz role token.
+     * Athenz access token.
      */
-    ROLE_TOKEN(AthenzHeaderNames.YAHOO_ROLE_AUTH),
+    ACCESS_TOKEN(HttpHeaderNames.AUTHORIZATION, false, "Bearer"),
     /**
-     * Athenz access token.
+     * Athenz role token. {@code "Athenz-Role-Auth"} is used as the header name for this token type.
+     */
+    ATHENZ_ROLE_TOKEN(AthenzHeaderNames.ATHENZ_ROLE_AUTH, true, null),
+    /**
+     * The legacy Athenz role token used by Yahoo. {@code "Yahoo-Role-Auth"} is used as the
+     * header name for this token type.
      */
-    ACCESS_TOKEN(HttpHeaderNames.AUTHORIZATION);
+    YAHOO_ROLE_TOKEN(AthenzHeaderNames.YAHOO_ROLE_AUTH, true, null);
 
-    TokenType(AsciiString headerName) {
+    TokenType(AsciiString headerName, boolean isRoleToken, @Nullable String authScheme) {
         this.headerName = headerName;
+        this.isRoleToken = isRoleToken;
+        this.authScheme = authScheme;
     }
 
     private final AsciiString headerName;
+    private final boolean isRoleToken;
+    @Nullable
+    private final String authScheme;
 
     /**
      * Returns the header name used to pass the token.
      */
     public AsciiString headerName() {
         return headerName;
     }
+
+    /**
+     * Returns whether this token type is a role token.
+     */
+    public boolean isRoleToken() {
+        return isRoleToken;
+    }
+
+    /**
+     * Returns the authentication scheme used for this token type, or {@code null} if not applicable.
+     * For example, {@link TokenType#ACCESS_TOKEN} uses "Bearer" as the authentication scheme,
+     * while {@link TokenType#YAHOO_ROLE_TOKEN} and {@link TokenType#ATHENZ_ROLE_TOKEN} do not use
+     * any authentication scheme.
+     */
+    @Nullable
+    public String authScheme() {
+        return authScheme;
+    }
 }
diff --git a/athenz/src/main/java/com/linecorp/armeria/internal/common/athenz/AthenzHeaderNames.java b/athenz/src/main/java/com/linecorp/armeria/internal/common/athenz/AthenzHeaderNames.java
@@ -22,7 +22,9 @@
 
 public final class AthenzHeaderNames {
 
-    public static final AsciiString YAHOO_ROLE_AUTH = HttpHeaderNames.of("yahoo-role-auth");
+    public static final AsciiString YAHOO_ROLE_AUTH = HttpHeaderNames.of("Yahoo-Role-Auth");
+
+    public static final AsciiString ATHENZ_ROLE_AUTH = HttpHeaderNames.of("Athenz-Role-Auth");
 
     private AthenzHeaderNames() {}
 }
diff --git a/athenz/src/main/java/com/linecorp/armeria/server/athenz/RequiresAthenzRole.java b/athenz/src/main/java/com/linecorp/armeria/server/athenz/RequiresAthenzRole.java
@@ -81,7 +81,9 @@
     /**
      * The required {@link TokenType}.
      */
-    TokenType[] tokenType() default { TokenType.ROLE_TOKEN, TokenType.ACCESS_TOKEN };
+    TokenType[] tokenType() default {
+            TokenType.ACCESS_TOKEN, TokenType.ATHENZ_ROLE_TOKEN, TokenType.YAHOO_ROLE_TOKEN
+    };
 
     /**
      * A special parameter in order to specify the order of a {@link Decorator}.
diff --git a/athenz/src/test/java/com/linecorp/armeria/common/athenz/TokenTypeTest.java b/athenz/src/test/java/com/linecorp/armeria/common/athenz/TokenTypeTest.java
@@ -0,0 +1,48 @@
+/*
+ * 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.common.athenz;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import org.junit.jupiter.api.Test;
+
+class TokenTypeTest {
+
+    @Test
+    void testAccessToken() {
+        final TokenType tokenType = TokenType.ACCESS_TOKEN;
+        assertThat(tokenType.headerName().toString()).isEqualToIgnoringCase("Authorization");
+        assertThat(tokenType.isRoleToken()).isFalse();
+        assertThat(tokenType.authScheme()).isEqualTo("Bearer");
+    }
+
+    @Test
+    void testYahooRoleToken() {
+        final TokenType tokenType = TokenType.YAHOO_ROLE_TOKEN;
+        assertThat(tokenType.headerName().toString()).isEqualToIgnoringCase("Yahoo-Role-Auth");
+        assertThat(tokenType.isRoleToken()).isTrue();
+        assertThat(tokenType.authScheme()).isNull();
+    }
+
+    @Test
+    void testAthenzRoleToken() {
+        final TokenType tokenType = TokenType.ATHENZ_ROLE_TOKEN;
+        assertThat(tokenType.headerName().toString()).isEqualToIgnoringCase("Athenz-Role-Auth");
+        assertThat(tokenType.isRoleToken()).isTrue();
+        assertThat(tokenType.authScheme()).isNull();
+    }
+}
diff --git a/athenz/src/test/java/com/linecorp/armeria/server/athenz/AthenzAnnotatedServiceTest.java b/athenz/src/test/java/com/linecorp/armeria/server/athenz/AthenzAnnotatedServiceTest.java
@@ -76,9 +76,11 @@ protected void configure(ServerBuilder sb) {
     };
 
     @CsvSource({
-            "foo-service, ROLE_TOKEN",
+            "foo-service, YAHOO_ROLE_TOKEN",
+            "foo-service, ATHENZ_ROLE_TOKEN",
             "foo-service, ACCESS_TOKEN",
-            "test-service, ROLE_TOKEN",
+            "test-service, YAHOO_ROLE_TOKEN",
+            "test-service, ATHENZ_ROLE_TOKEN",
             "test-service, ACCESS_TOKEN"
     })
     @ParameterizedTest
@@ -98,9 +100,11 @@ void testUserRole(String serviceName, TokenType tokenType) {
     }
 
     @CsvSource({
-            "foo-service, ROLE_TOKEN, false",
+            "foo-service, YAHOO_ROLE_TOKEN, false",
+            "foo-service, ATHENZ_ROLE_TOKEN, false",
             "foo-service, ACCESS_TOKEN, false",
-            "test-service, ROLE_TOKEN, true",
+            "test-service, YAHOO_ROLE_TOKEN, true",
+            "test-service, ATHENZ_ROLE_TOKEN, true",
             "test-service, ACCESS_TOKEN, true"
     })
     @ParameterizedTest
@@ -122,7 +126,7 @@ void testAdminRole(String serviceName, TokenType tokenType, boolean shouldSuccee
                         .isInstanceOf(AccessDeniedException.class)
                         .hasMessage("Failed to obtain an Athenz %s token. " +
                                     "(domain: testing, roles: test_role_admin)",
-                                    tokenType == TokenType.ROLE_TOKEN ? "role" : "access");
+                                    tokenType.isRoleToken() ? "role" : "access");
             }
         }
     }
diff --git a/athenz/src/test/java/com/linecorp/armeria/server/athenz/AthenzIntegrationTest.java b/athenz/src/test/java/com/linecorp/armeria/server/athenz/AthenzIntegrationTest.java
@@ -74,6 +74,10 @@ protected void configure(ServerBuilder sb) {
                 if (!authorization.isEmpty()) {
                     return HttpResponse.of("YahooRoleAuth " + authorization);
                 }
+                authorization = req.headers().get(AthenzHeaderNames.ATHENZ_ROLE_AUTH, "");
+                if (!authorization.isEmpty()) {
+                    return HttpResponse.of("AthenzRoleAuth " + authorization);
+                }
                 // Should not reach here.
                 return HttpResponse.of(HttpStatus.INTERNAL_SERVER_ERROR);
             });
@@ -93,6 +97,10 @@ protected void configure(ServerBuilder sb) {
                 if (!authorization.isEmpty()) {
                     return HttpResponse.of("YahooRoleAuth " + authorization);
                 }
+                authorization = req.headers().get(AthenzHeaderNames.ATHENZ_ROLE_AUTH, "");
+                if (!authorization.isEmpty()) {
+                    return HttpResponse.of("AthenzRoleAuth " + authorization);
+                }
                 // Should not reach here.
                 return HttpResponse.of(HttpStatus.INTERNAL_SERVER_ERROR);
             });
@@ -119,9 +127,12 @@ void shouldObtainAdminRole(TokenType tokenType) {
             final AggregatedHttpResponse response = client.get("/admin");
             assertThat(response.status()).isEqualTo(HttpStatus.OK);
             switch (tokenType) {
-                case ROLE_TOKEN:
+                case YAHOO_ROLE_TOKEN:
                     assertThat(response.contentUtf8()).startsWith("YahooRoleAuth ");
                     break;
+                case ATHENZ_ROLE_TOKEN:
+                    assertThat(response.contentUtf8()).startsWith("AthenzRoleAuth ");
+                    break;
                 case ACCESS_TOKEN:
                     assertThat(response.contentUtf8()).startsWith("Authorization ");
                     break;
@@ -144,7 +155,7 @@ void shouldNotObtainAdminRole(TokenType tokenType) {
                     client.get("/admin");
                 }).isInstanceOf(AccessDeniedException.class)
                   .hasMessage("Failed to obtain an Athenz " +
-                              (tokenType == TokenType.ROLE_TOKEN ? "role" : "access") +
+                              (tokenType.isRoleToken() ? "role" : "access") +
                               " token. (domain: testing, roles: test_role_admin)");
                 final ClientRequestContext ctx = captor.get();
                 // Make sure the RequestLog is completed when the request was rejected by the decorator.

Original file line number	Diff line number	Diff line change
`@@ -22,7 +22,9 @@`
`22`	`22`
`23`	`23`	`public final class AthenzHeaderNames {`
`24`	`24`
`25`		`- public static final AsciiString YAHOO_ROLE_AUTH = HttpHeaderNames.of("yahoo-role-auth");`
	`25`	`+ public static final AsciiString YAHOO_ROLE_AUTH = HttpHeaderNames.of("Yahoo-Role-Auth");`
	`26`	`+`
	`27`	`+ public static final AsciiString ATHENZ_ROLE_AUTH = HttpHeaderNames.of("Athenz-Role-Auth");`
`26`	`28`
`27`	`29`	`private AthenzHeaderNames() {}`
`28`	`30`	`}`