Get access token from request attribute

5 files changed, 20 insertions, 5 deletions

diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/restapi/application/ApplicationApiHandler.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/restapi/application/ApplicationApiHandler.java
index b3cba4809f0..c5db553219e 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/restapi/application/ApplicationApiHandler.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/restapi/application/ApplicationApiHandler.java
@@ -1238,8 +1238,8 @@ public class ApplicationApiHandler extends LoggingRequestHandler {
     }
 
     private static Optional<OktaAccessToken> getOktaAccessToken(HttpRequest request) {
-        return Optional.ofNullable(request.getHeader(OktaAccessToken.HTTP_HEADER_NAME))
-                .map(OktaAccessToken::new);
+        return Optional.ofNullable(request.getJDiscRequest().context().get("okta.access-token"))
+                .map(attribute -> new OktaAccessToken((String) attribute));
     }
 
     private static ApplicationId appIdFromPath(Path path) {
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/integration/AthenzFilterMock.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/integration/AthenzFilterMock.java
index 72027234b28..09515b8905e 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/integration/AthenzFilterMock.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/integration/AthenzFilterMock.java
@@ -16,6 +16,8 @@ import com.yahoo.vespa.athenz.api.AthenzPrincipal;
 import com.yahoo.vespa.athenz.utils.AthenzIdentities;
 import com.yahoo.yolean.chain.Before;
 
+import java.util.Optional;
+
 /**
  * @author bjorncs
  */
@@ -23,6 +25,7 @@ import com.yahoo.yolean.chain.Before;
 public class AthenzFilterMock implements SecurityRequestFilter {
 
     public static final String IDENTITY_HEADER_NAME = "Athenz-Identity";
+    public static final String OKTA_ACCESS_TOKEN_HEADER_NAME = "Okta-Access-Token";
 
     private static final ObjectMapper mapper = new ObjectMapper();
 
@@ -45,6 +48,8 @@ public class AthenzFilterMock implements SecurityRequestFilter {
             AthenzPrincipal principal = new AthenzPrincipal(identity);
             request.setUserPrincipal(principal);
         }
+        Optional.ofNullable(request.getHeader(OKTA_ACCESS_TOKEN_HEADER_NAME))
+                .ifPresent(header -> request.setAttribute("okta.access-token", header));
     }
 
 }
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/ApplicationRequestToDiscFilterRequestWrapper.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/ApplicationRequestToDiscFilterRequestWrapper.java
index 4883bde99b1..f2dc22c9c60 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/ApplicationRequestToDiscFilterRequestWrapper.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/ApplicationRequestToDiscFilterRequestWrapper.java
@@ -104,6 +104,7 @@ public class ApplicationRequestToDiscFilterRequestWrapper extends DiscFilterRequ
     public Request getUpdatedRequest() {
         Request updatedRequest = new Request(this.request.getUri(), this.request.getBody(), this.request.getMethod(), this.userPrincipal);
         this.request.getHeaders().forEach(updatedRequest.getHeaders()::put);
+        updatedRequest.getAttributes().putAll(this.request.getAttributes());
         return updatedRequest;
     }
 
@@ -191,4 +192,14 @@ public class ApplicationRequestToDiscFilterRequestWrapper extends DiscFilterRequ
     public void clearCookies() {
         throw new UnsupportedOperationException();
     }
+
+    @Override
+    public Object getAttribute(String name) {
+        return request.getAttributes().get(name);
+    }
+
+    @Override
+    public void setAttribute(String name, Object value) {
+        request.getAttributes().put(name, value);
+    }
 }
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/ControllerContainerTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/ControllerContainerTest.java
index e908777a8b0..ce69f32a21e 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/ControllerContainerTest.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/ControllerContainerTest.java
@@ -16,6 +16,7 @@ import java.io.UncheckedIOException;
 import java.nio.charset.CharacterCodingException;
 
 import static com.yahoo.vespa.hosted.controller.integration.AthenzFilterMock.IDENTITY_HEADER_NAME;
+import static com.yahoo.vespa.hosted.controller.integration.AthenzFilterMock.OKTA_ACCESS_TOKEN_HEADER_NAME;
 import static org.junit.Assert.assertEquals;
 
 /**
@@ -156,7 +157,7 @@ public class ControllerContainerTest {
     }
 
     protected static Request addOktaAccessToken(Request request, OktaAccessToken token) {
-        request.getHeaders().put(OktaAccessToken.HTTP_HEADER_NAME, token.token());
+        request.getHeaders().put(OKTA_ACCESS_TOKEN_HEADER_NAME, token.token());
         return request;
     }
 
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/api/OktaAccessToken.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/api/OktaAccessToken.java
index 8b19f7abdd5..b6d34bce4e4 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/api/OktaAccessToken.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/api/OktaAccessToken.java
@@ -8,8 +8,6 @@ import java.util.Objects;
  */
 public class OktaAccessToken {
 
-    public static final String HTTP_HEADER_NAME = "Okta-Access-Token";
-
     private final String token;
 
     public OktaAccessToken(String token) {