Use getClientCertificateChain() in AthenzPrincipalFilter

3 files changed, 16 insertions, 5 deletions

diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/filter/AthenzPrincipalFilter.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/filter/AthenzPrincipalFilter.java
index 81692e790a9..c5406669f67 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/filter/AthenzPrincipalFilter.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/filter/AthenzPrincipalFilter.java
@@ -81,7 +81,7 @@ public class AthenzPrincipalFilter implements SecurityRequestFilter {
     }
 
     private static Optional<X509Certificate> getClientCertificate(DiscFilterRequest request) {
-        return Optional.ofNullable((X509Certificate[]) request.getAttribute("jdisc.request.X509Certificate"))
+        return request.getClientCertificateChain()
                 .map(chain -> chain[0]);
     }
 
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/athenz/filter/AthenzPrincipalFilterTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/athenz/filter/AthenzPrincipalFilterTest.java
index 697f69d8da3..b0a51ecb16f 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/athenz/filter/AthenzPrincipalFilterTest.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/athenz/filter/AthenzPrincipalFilterTest.java
@@ -36,6 +36,7 @@ import java.time.Duration;
 import java.time.Instant;
 import java.util.Date;
 import java.util.Objects;
+import java.util.Optional;
 
 import static com.yahoo.jdisc.Response.Status.UNAUTHORIZED;
 import static java.util.stream.Collectors.joining;
@@ -69,6 +70,7 @@ public class AthenzPrincipalFilterTest {
         DiscFilterRequest request = mock(DiscFilterRequest.class);
         AthenzPrincipal principal = new AthenzPrincipal(IDENTITY, NTOKEN);
         when(request.getHeader(ATHENZ_PRINCIPAL_HEADER)).thenReturn(NTOKEN.getRawToken());
+        when(request.getClientCertificateChain()).thenReturn(Optional.empty());
         when(validator.validate(NTOKEN)).thenReturn(principal);
 
         AthenzPrincipalFilter filter = new AthenzPrincipalFilter(validator, Runnable::run, ATHENZ_PRINCIPAL_HEADER);
@@ -81,6 +83,7 @@ public class AthenzPrincipalFilterTest {
     public void missing_token_and_certificate_is_unauthorized() {
         DiscFilterRequest request = mock(DiscFilterRequest.class);
         when(request.getHeader(ATHENZ_PRINCIPAL_HEADER)).thenReturn(null);
+        when(request.getClientCertificateChain()).thenReturn(Optional.empty());
 
         ResponseHandlerMock responseHandler = new ResponseHandlerMock();
 
@@ -95,6 +98,7 @@ public class AthenzPrincipalFilterTest {
         DiscFilterRequest request = mock(DiscFilterRequest.class);
         String errorMessage = "Invalid token";
         when(request.getHeader(ATHENZ_PRINCIPAL_HEADER)).thenReturn(NTOKEN.getRawToken());
+        when(request.getClientCertificateChain()).thenReturn(Optional.empty());
         when(validator.validate(NTOKEN)).thenThrow(new InvalidTokenException(errorMessage));
 
         ResponseHandlerMock responseHandler = new ResponseHandlerMock();
@@ -109,7 +113,7 @@ public class AthenzPrincipalFilterTest {
     public void certificate_is_accepted() {
         DiscFilterRequest request = mock(DiscFilterRequest.class);
         when(request.getHeader(ATHENZ_PRINCIPAL_HEADER)).thenReturn(null);
-        when(request.getAttribute("jdisc.request.X509Certificate")).thenReturn(new X509Certificate[]{CERTIFICATE});
+        when(request.getClientCertificateChain()).thenReturn(Optional.of(new X509Certificate[]{CERTIFICATE}));
 
         ResponseHandlerMock responseHandler = new ResponseHandlerMock();
 
@@ -125,7 +129,7 @@ public class AthenzPrincipalFilterTest {
         DiscFilterRequest request = mock(DiscFilterRequest.class);
         AthenzPrincipal principalWithToken = new AthenzPrincipal(IDENTITY, NTOKEN);
         when(request.getHeader(ATHENZ_PRINCIPAL_HEADER)).thenReturn(NTOKEN.getRawToken());
-        when(request.getAttribute("jdisc.request.X509Certificate")).thenReturn(new X509Certificate[]{CERTIFICATE});
+        when(request.getClientCertificateChain()).thenReturn(Optional.of(new X509Certificate[]{CERTIFICATE}));
         when(validator.validate(NTOKEN)).thenReturn(principalWithToken);
 
         ResponseHandlerMock responseHandler = new ResponseHandlerMock();
@@ -141,8 +145,8 @@ public class AthenzPrincipalFilterTest {
         DiscFilterRequest request = mock(DiscFilterRequest.class);
         AthenzUser conflictingIdentity = AthenzUser.fromUserId("mallory");
         when(request.getHeader(ATHENZ_PRINCIPAL_HEADER)).thenReturn(NTOKEN.getRawToken());
-        when(request.getAttribute("jdisc.request.X509Certificate"))
-                .thenReturn(new X509Certificate[]{createSelfSignedCertificate(conflictingIdentity)});
+        when(request.getClientCertificateChain())
+                .thenReturn(Optional.of(new X509Certificate[]{createSelfSignedCertificate(conflictingIdentity)}));
         when(validator.validate(NTOKEN)).thenReturn(new AthenzPrincipal(IDENTITY));
 
         ResponseHandlerMock responseHandler = new ResponseHandlerMock();
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 da9dd1f0786..d5b1b85de5f 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
@@ -11,10 +11,12 @@ import com.yahoo.jdisc.http.servlet.ServletOrJdiscHttpRequest;
 import java.net.SocketAddress;
 import java.net.URI;
 import java.security.Principal;
+import java.security.cert.X509Certificate;
 import java.util.Collections;
 import java.util.Enumeration;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.TimeUnit;
 
 /**
@@ -176,6 +178,11 @@ public class ApplicationRequestToDiscFilterRequestWrapper extends DiscFilterRequ
     }
 
     @Override
+    public Optional<X509Certificate[]> getClientCertificateChain() {
+        return Optional.empty();
+    }
+
+    @Override
     public void clearCookies() {
         throw new UnsupportedOperationException();
     }