Revert "Move NTokenValidator to vespa-athenz + load pub keys from file"

10 files changed, 398 insertions, 22 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 b7ede7635c6..5166f53c6d2 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
@@ -9,16 +9,15 @@ import com.yahoo.jdisc.http.filter.security.cors.CorsRequestFilterBase;
 import com.yahoo.vespa.athenz.api.AthenzPrincipal;
 import com.yahoo.vespa.athenz.api.NToken;
 import com.yahoo.vespa.athenz.utils.AthenzIdentities;
-import com.yahoo.vespa.athenz.utils.ntoken.AthenzConfTruststore;
-import com.yahoo.vespa.athenz.utils.ntoken.NTokenValidator;
+import com.yahoo.vespa.hosted.controller.api.integration.athenz.ZmsKeystore;
 import com.yahoo.vespa.hosted.controller.athenz.config.AthenzConfig;
 
-import java.nio.file.Paths;
 import java.security.cert.X509Certificate;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.Executor;
 
 
 /**
@@ -31,23 +30,31 @@ import java.util.Set;
  *
  * @author bjorncs
  */
-// TODO bjorncs: Move this class to jdisc-security-filters bundle
+// TODO bjorncs: Move this class to vespa-athenz bundle
 public class AthenzPrincipalFilter extends CorsRequestFilterBase {
 
     private final NTokenValidator validator;
     private final String principalTokenHeader;
 
+    /**
+     * @param executor to preload the ZMS public keys with
+     */
     @Inject
-    public AthenzPrincipalFilter(AthenzConfig athenzConfig, CorsFilterConfig corsConfig) {
-        this(new NTokenValidator(Paths.get(athenzConfig.athenzConfFile())), athenzConfig.principalHeaderName(), new HashSet<>(corsConfig.allowedUrls()));
+    public AthenzPrincipalFilter(ZmsKeystore zmsKeystore,
+                                 Executor executor,
+                                 AthenzConfig athenzConfig,
+                                 CorsFilterConfig corsConfig) {
+        this(new NTokenValidator(zmsKeystore), executor, athenzConfig.principalHeaderName(), new HashSet<>(corsConfig.allowedUrls()));
     }
 
     AthenzPrincipalFilter(NTokenValidator validator,
+                          Executor executor,
                           String principalTokenHeader,
                           Set<String> corsAllowedUrls) {
         super(corsAllowedUrls);
         this.validator = validator;
         this.principalTokenHeader = principalTokenHeader;
+        executor.execute(validator::preloadPublicKeys);
     }
 
     @Override
diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/filter/NTokenValidator.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/filter/NTokenValidator.java
new file mode 100644
index 00000000000..4dcca519058
--- /dev/null
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/filter/NTokenValidator.java
@@ -0,0 +1,81 @@
+// Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.hosted.controller.athenz.filter;
+
+import com.yahoo.athenz.auth.token.PrincipalToken;
+import com.yahoo.log.LogLevel;
+import com.yahoo.vespa.athenz.api.AthenzDomain;
+import com.yahoo.vespa.athenz.api.AthenzPrincipal;
+import com.yahoo.vespa.athenz.api.NToken;
+import com.yahoo.vespa.athenz.utils.AthenzIdentities;
+import com.yahoo.vespa.hosted.controller.api.integration.athenz.InvalidTokenException;
+import com.yahoo.vespa.hosted.controller.api.integration.athenz.ZmsKeystore;
+
+import java.security.PublicKey;
+import java.time.Duration;
+import java.util.Optional;
+import java.util.logging.Logger;
+
+import static com.yahoo.vespa.athenz.utils.AthenzIdentities.ZMS_ATHENZ_SERVICE;
+
+
+/**
+ * Validates the content of an NToken:
+ *   1) Verifies that the token is signed by the sys.auth.zms service (by validating the signature)
+ *   2) Verifies that the token is not expired
+ *
+ * @author bjorncs
+ */
+// TODO Move to vespa-athenz
+class NTokenValidator {
+
+        // Max allowed skew in token timestamp (only for creation, not expiry timestamp)
+    private static final long ALLOWED_TIMESTAMP_OFFSET = Duration.ofMinutes(5).getSeconds();
+
+    private static final Logger log = Logger.getLogger(NTokenValidator.class.getName());
+
+    private final ZmsKeystore keystore;
+
+    NTokenValidator(ZmsKeystore keystore) {
+        this.keystore = keystore;
+    }
+
+    void preloadPublicKeys() {
+        keystore.preloadKeys(ZMS_ATHENZ_SERVICE);
+    }
+
+    AthenzPrincipal validate(NToken token) throws InvalidTokenException {
+        PrincipalToken principalToken = new PrincipalToken(token.getRawToken());
+        PublicKey zmsPublicKey = getPublicKey(principalToken.getKeyId())
+                .orElseThrow(() -> new InvalidTokenException("NToken has an unknown keyId"));
+        validateSignatureAndExpiration(principalToken, zmsPublicKey);
+        return new AthenzPrincipal(
+                AthenzIdentities.from(
+                        new AthenzDomain(principalToken.getDomain()),
+                        principalToken.getName()),
+                token);
+    }
+
+    private Optional<PublicKey> getPublicKey(String keyId) throws InvalidTokenException {
+        try {
+            return keystore.getPublicKey(ZMS_ATHENZ_SERVICE, keyId);
+        } catch (Exception e) {
+            logDebug(e.getMessage());
+            throw new InvalidTokenException("Failed to retrieve public key");
+        }
+    }
+
+    private static void validateSignatureAndExpiration(PrincipalToken token,
+                                                       PublicKey zmsPublicKey) throws InvalidTokenException {
+        StringBuilder errorMessageBuilder = new StringBuilder();
+        if (!token.validate(zmsPublicKey, (int) ALLOWED_TIMESTAMP_OFFSET, true, errorMessageBuilder)) {
+            String message = "NToken is expired or has invalid signature: " + errorMessageBuilder.toString();
+            logDebug(message);
+            throw new InvalidTokenException(message);
+        }
+    }
+
+    private static void logDebug(String message) {
+        log.log(LogLevel.DEBUG, "Failed to validate NToken: " + message);
+    }
+
+}
diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/filter/UserAuthWithAthenzPrincipalFilter.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/filter/UserAuthWithAthenzPrincipalFilter.java
index 0aa5c89c971..b801c038bd8 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/filter/UserAuthWithAthenzPrincipalFilter.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/filter/UserAuthWithAthenzPrincipalFilter.java
@@ -10,11 +10,13 @@ import com.yahoo.vespa.athenz.api.AthenzPrincipal;
 import com.yahoo.vespa.athenz.api.AthenzUser;
 import com.yahoo.vespa.athenz.api.NToken;
 import com.yahoo.vespa.hosted.controller.api.identifiers.UserId;
+import com.yahoo.vespa.hosted.controller.api.integration.athenz.ZmsKeystore;
 import com.yahoo.vespa.hosted.controller.athenz.config.AthenzConfig;
 import com.yahoo.yolean.chain.After;
 
 import java.security.Principal;
 import java.util.Optional;
+import java.util.concurrent.Executor;
 import java.util.logging.Logger;
 import java.util.stream.Stream;
 
@@ -36,8 +38,11 @@ public class UserAuthWithAthenzPrincipalFilter extends AthenzPrincipalFilter {
     private final String principalHeaderName;
 
     @Inject
-    public UserAuthWithAthenzPrincipalFilter(AthenzConfig athenzConfig, CorsFilterConfig corsConfig) {
-        super(athenzConfig, corsConfig);
+    public UserAuthWithAthenzPrincipalFilter(ZmsKeystore zmsKeystore,
+                                             Executor executor,
+                                             AthenzConfig athenzConfig,
+                                             CorsFilterConfig corsConfig) {
+        super(zmsKeystore, executor, athenzConfig, corsConfig);
         this.userAuthenticationPassThruAttribute = athenzConfig.userAuthenticationPassThruAttribute();
         this.principalHeaderName = athenzConfig.principalHeaderName();
     }
diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/impl/ZmsClientImpl.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/impl/ZmsClientImpl.java
index 6179d9891fd..67191d4c09d 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/impl/ZmsClientImpl.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/impl/ZmsClientImpl.java
@@ -1,18 +1,22 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.hosted.controller.athenz.impl;
 
+import com.yahoo.athenz.auth.util.Crypto;
 import com.yahoo.athenz.zms.DomainList;
 import com.yahoo.athenz.zms.ProviderResourceGroupRoles;
+import com.yahoo.athenz.zms.PublicKeyEntry;
+import com.yahoo.athenz.zms.ServiceIdentity;
 import com.yahoo.athenz.zms.Tenancy;
 import com.yahoo.athenz.zms.TenantRoleAction;
 import com.yahoo.athenz.zms.ZMSClient;
 import com.yahoo.athenz.zms.ZMSClientException;
 import com.yahoo.log.LogLevel;
+import com.yahoo.vespa.hosted.controller.api.identifiers.ApplicationId;
 import com.yahoo.vespa.athenz.api.AthenzDomain;
+import com.yahoo.vespa.hosted.controller.api.integration.athenz.ApplicationAction;
 import com.yahoo.vespa.athenz.api.AthenzIdentity;
+import com.yahoo.vespa.athenz.api.AthenzPublicKey;
 import com.yahoo.vespa.athenz.api.AthenzService;
-import com.yahoo.vespa.hosted.controller.api.identifiers.ApplicationId;
-import com.yahoo.vespa.hosted.controller.api.integration.athenz.ApplicationAction;
 import com.yahoo.vespa.hosted.controller.api.integration.athenz.ZmsClient;
 import com.yahoo.vespa.hosted.controller.api.integration.athenz.ZmsException;
 import com.yahoo.vespa.hosted.controller.athenz.config.AthenzConfig;
@@ -126,6 +130,28 @@ public class ZmsClientImpl implements ZmsClient {
                 });
     }
 
+    @Override
+    public AthenzPublicKey getPublicKey(AthenzService service, String keyId) {
+        log("getPublicKeyEntry(domain=%s, service=%s, keyId=%s)", service.getDomain().getName(), service.getName(), keyId);
+        return getOrThrow(() -> {
+            PublicKeyEntry entry = zmsClient.getPublicKeyEntry(service.getDomain().getName(), service.getName(), keyId);
+            return fromYbase64EncodedKey(entry.getKey(), keyId);
+        });
+    }
+
+    @Override
+    public List<AthenzPublicKey> getPublicKeys(AthenzService service) {
+        log("getServiceIdentity(domain=%s, service=%s)", service.getDomain().getName(), service.getName());
+        return getOrThrow(() -> {
+            ServiceIdentity serviceIdentity = zmsClient.getServiceIdentity(service.getDomain().getName(), service.getName());
+            return toAthenzPublicKeys(serviceIdentity.getPublicKeys());
+        });
+    }
+
+    private static AthenzPublicKey fromYbase64EncodedKey(String encodedKey, String keyId) {
+        return new AthenzPublicKey(Crypto.loadPublicKey(Crypto.ybase64DecodeString(encodedKey)), keyId);
+    }
+
     private static List<TenantRoleAction> createTenantRoleActions() {
         return Arrays.stream(ApplicationAction.values())
                 .map(action -> new TenantRoleAction().setAction(action.name()).setRole(action.roleName))
@@ -136,6 +162,12 @@ public class ZmsClientImpl implements ZmsClient {
         return domains.stream().map(AthenzDomain::new).collect(toList());
     }
 
+    private static List<AthenzPublicKey> toAthenzPublicKeys(List<PublicKeyEntry> publicKeys) {
+        return publicKeys.stream()
+                .map(entry -> fromYbase64EncodedKey(entry.getKey(), entry.getId()))
+                .collect(toList());
+    }
+
     private boolean hasAccess(String action, String resource, AthenzIdentity identity) {
         log("getAccess(action=%s, resource=%s, principal=%s)", action, resource, identity);
         return getOrThrow(
diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/impl/ZmsKeystoreImpl.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/impl/ZmsKeystoreImpl.java
new file mode 100644
index 00000000000..4b194651439
--- /dev/null
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/impl/ZmsKeystoreImpl.java
@@ -0,0 +1,120 @@
+// Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.hosted.controller.athenz.impl;
+
+import com.google.inject.Inject;
+import com.yahoo.log.LogLevel;
+import com.yahoo.vespa.athenz.api.AthenzPublicKey;
+import com.yahoo.vespa.athenz.api.AthenzService;
+import com.yahoo.vespa.hosted.controller.api.integration.athenz.AthenzClientFactory;
+import com.yahoo.vespa.hosted.controller.api.integration.athenz.ZmsException;
+import com.yahoo.vespa.hosted.controller.api.integration.athenz.ZmsKeystore;
+
+import java.security.PublicKey;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.logging.Logger;
+
+/**
+ *  Downloads and caches public keys for Athens services.
+ *
+ * @author bjorncs
+ */
+public class ZmsKeystoreImpl implements ZmsKeystore {
+    private static final Logger log = Logger.getLogger(ZmsKeystoreImpl.class.getName());
+
+    private final Map<FullKeyId, PublicKey> cachedKeys = new ConcurrentHashMap<>();
+    private final AthenzClientFactory athenzClientFactory;
+
+    @Inject
+    public ZmsKeystoreImpl(AthenzClientFactory factory) {
+        this.athenzClientFactory = factory;
+    }
+
+    @Override
+    public Optional<PublicKey> getPublicKey(AthenzService service, String keyId)  {
+        FullKeyId fullKeyId = new FullKeyId(service, keyId);
+        PublicKey cachedKey = cachedKeys.get(fullKeyId);
+        if (cachedKey != null) {
+            return Optional.of(cachedKey);
+        }
+        Optional<PublicKey> downloadedKey = downloadPublicKey(fullKeyId);
+        downloadedKey.ifPresent(key -> {
+            log.log(LogLevel.INFO, "Adding key " + fullKeyId + " to the cache");
+            cachedKeys.put(fullKeyId, key);
+        });
+        return downloadedKey;
+    }
+
+    @Override
+    public void preloadKeys(AthenzService service) {
+        try {
+            log.log(LogLevel.INFO, "Downloading keys for " + service);
+            List<AthenzPublicKey> publicKeys = athenzClientFactory.createZmsClientWithServicePrincipal()
+                    .getPublicKeys(service);
+            for (AthenzPublicKey publicKey : publicKeys) {
+                FullKeyId fullKeyId = new FullKeyId(service, publicKey.getKeyId());
+                log.log(LogLevel.DEBUG, "Adding key " + fullKeyId + " to the cache");
+                cachedKeys.put(fullKeyId, publicKey.getPublicKey());
+            }
+            log.log(LogLevel.INFO, "Successfully downloaded keys for " + service);
+        } catch (ZmsException e) {
+            log.log(LogLevel.WARNING, "Failed to download keys for " + service + ": " + e.getMessage());
+        }
+    }
+
+    private Optional<PublicKey> downloadPublicKey(FullKeyId fullKeyId) {
+        try {
+            log.log(LogLevel.INFO, "Downloading key " + fullKeyId);
+            AthenzPublicKey publicKey = athenzClientFactory.createZmsClientWithServicePrincipal()
+                    .getPublicKey(fullKeyId.service, fullKeyId.keyId);
+            return Optional.of(publicKey.getPublicKey());
+        } catch (ZmsException e) {
+            if (e.getCode() == 404) { // Key does not exist
+                log.log(LogLevel.INFO, "Key " + fullKeyId + " not found");
+                return Optional.empty();
+            }
+            String msg = String.format("Unable to retrieve public key from Athens (%s): %s", fullKeyId, e.getMessage());
+            throw createException(msg, e);
+        }
+    }
+
+    private static RuntimeException createException(String message, Exception cause) {
+        log.log(LogLevel.ERROR, message);
+        return new RuntimeException(message, cause);
+    }
+
+    private static class FullKeyId {
+        private final AthenzService service;
+        private final String keyId;
+
+        private FullKeyId(AthenzService service, String keyId) {
+            this.service = service;
+            this.keyId = keyId;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (o == null || getClass() != o.getClass()) return false;
+            FullKeyId fullKeyId1 = (FullKeyId) o;
+            return Objects.equals(service, fullKeyId1.service) &&
+                    Objects.equals(keyId, fullKeyId1.keyId);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(service, keyId);
+        }
+
+        @Override
+        public String toString() {
+            return "FullKeyId{" +
+                    "service=" + service +
+                    ", keyId='" + keyId + '\'' +
+                    '}';
+        }
+    }
+}
diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/mock/ZmsClientMock.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/mock/ZmsClientMock.java
index 5e8674ce637..3ee2655108a 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/mock/ZmsClientMock.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/mock/ZmsClientMock.java
@@ -5,6 +5,8 @@ import com.yahoo.vespa.hosted.controller.api.identifiers.ApplicationId;
 import com.yahoo.vespa.athenz.api.AthenzDomain;
 import com.yahoo.vespa.hosted.controller.api.integration.athenz.ApplicationAction;
 import com.yahoo.vespa.athenz.api.AthenzIdentity;
+import com.yahoo.vespa.athenz.api.AthenzPublicKey;
+import com.yahoo.vespa.athenz.api.AthenzService;
 import com.yahoo.vespa.hosted.controller.api.integration.athenz.ZmsClient;
 import com.yahoo.vespa.hosted.controller.api.integration.athenz.ZmsException;
 
@@ -94,6 +96,16 @@ public class ZmsClientMock implements ZmsClient {
         return new ArrayList<>(athenz.domains.keySet());
     }
 
+    @Override
+    public AthenzPublicKey getPublicKey(AthenzService service, String keyId) {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public List<AthenzPublicKey> getPublicKeys(AthenzService service) {
+        throw new UnsupportedOperationException();
+    }
+
     private AthenzDbMock.Domain getDomainOrThrow(AthenzDomain domainName, boolean verifyVespaTenant) {
         AthenzDbMock.Domain domain = Optional.ofNullable(athenz.domains.get(domainName))
                 .orElseThrow(() -> zmsException(400, "Domain '%s' not found", domainName));
diff --git a/controller-server/src/main/resources/configdefinitions/athenz.def b/controller-server/src/main/resources/configdefinitions/athenz.def
index 8026c0d7f44..f8d65c25e47 100644
--- a/controller-server/src/main/resources/configdefinitions/athenz.def
+++ b/controller-server/src/main/resources/configdefinitions/athenz.def
@@ -42,6 +42,3 @@ service.privateKeySecretName    string
 
 # Expiry of service principal token and certificate
 service.credentialsExpiryMinutes int   default=43200  # 30 days
-
-# Path to athenz.conf file
-athenzConfFile                  string
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 9fe582b829f..301fc461b6f 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
@@ -14,7 +14,7 @@ import com.yahoo.vespa.athenz.api.NToken;
 import com.yahoo.vespa.athenz.tls.KeyAlgorithm;
 import com.yahoo.vespa.athenz.tls.KeyUtils;
 import com.yahoo.vespa.athenz.tls.X509CertificateBuilder;
-import com.yahoo.vespa.athenz.utils.ntoken.NTokenValidator;
+import com.yahoo.vespa.hosted.controller.api.integration.athenz.InvalidTokenException;
 import com.yahoo.vespa.hosted.controller.restapi.ApplicationRequestToDiscFilterRequestWrapper;
 import org.junit.Before;
 import org.junit.Test;
@@ -71,7 +71,7 @@ public class AthenzPrincipalFilterTest {
         AthenzPrincipal principal = new AthenzPrincipal(IDENTITY, NTOKEN);
         validator.add(NTOKEN, principal);
 
-        AthenzPrincipalFilter filter = new AthenzPrincipalFilter(validator, ATHENZ_PRINCIPAL_HEADER, CORS_ALLOWED_URLS);
+        AthenzPrincipalFilter filter = new AthenzPrincipalFilter(validator, Runnable::run, ATHENZ_PRINCIPAL_HEADER, CORS_ALLOWED_URLS);
         DiscFilterRequest filterRequest = new ApplicationRequestToDiscFilterRequestWrapper(request);
         filter.filter(filterRequest, new ResponseHandlerMock());
 
@@ -80,7 +80,7 @@ public class AthenzPrincipalFilterTest {
 
     @Test
     public void missing_token_and_certificate_is_unauthorized() {
-        AthenzPrincipalFilter filter = new AthenzPrincipalFilter(validator, ATHENZ_PRINCIPAL_HEADER, CORS_ALLOWED_URLS);
+        AthenzPrincipalFilter filter = new AthenzPrincipalFilter(validator, Runnable::run, ATHENZ_PRINCIPAL_HEADER, CORS_ALLOWED_URLS);
         DiscFilterRequest filterRequest = new ApplicationRequestToDiscFilterRequestWrapper(new Request("/"));
         filter.filter(filterRequest, responseHandler);
 
@@ -91,7 +91,7 @@ public class AthenzPrincipalFilterTest {
     public void invalid_token_is_unauthorized() {
         Request request = defaultRequest();
 
-        AthenzPrincipalFilter filter = new AthenzPrincipalFilter(validator, ATHENZ_PRINCIPAL_HEADER, CORS_ALLOWED_URLS);
+        AthenzPrincipalFilter filter = new AthenzPrincipalFilter(validator, Runnable::run, ATHENZ_PRINCIPAL_HEADER, CORS_ALLOWED_URLS);
         DiscFilterRequest filterRequest = new ApplicationRequestToDiscFilterRequestWrapper(request);
         filter.filter(filterRequest, responseHandler);
 
@@ -101,7 +101,7 @@ public class AthenzPrincipalFilterTest {
 
     @Test
     public void certificate_is_accepted() {
-        AthenzPrincipalFilter filter = new AthenzPrincipalFilter(validator, ATHENZ_PRINCIPAL_HEADER, CORS_ALLOWED_URLS);
+        AthenzPrincipalFilter filter = new AthenzPrincipalFilter(validator, Runnable::run, ATHENZ_PRINCIPAL_HEADER, CORS_ALLOWED_URLS);
         DiscFilterRequest filterRequest = new ApplicationRequestToDiscFilterRequestWrapper(new Request("/"), singletonList(CERTIFICATE));
         filter.filter(filterRequest, responseHandler);
 
@@ -116,7 +116,7 @@ public class AthenzPrincipalFilterTest {
         AthenzPrincipal principalWithToken = new AthenzPrincipal(IDENTITY, NTOKEN);
         validator.add(NTOKEN, principalWithToken);
 
-        AthenzPrincipalFilter filter = new AthenzPrincipalFilter(validator, ATHENZ_PRINCIPAL_HEADER, CORS_ALLOWED_URLS);
+        AthenzPrincipalFilter filter = new AthenzPrincipalFilter(validator, Runnable::run, ATHENZ_PRINCIPAL_HEADER, CORS_ALLOWED_URLS);
         DiscFilterRequest filterRequest = new ApplicationRequestToDiscFilterRequestWrapper(request, singletonList(CERTIFICATE));
         filter.filter(filterRequest, responseHandler);
 
@@ -130,7 +130,7 @@ public class AthenzPrincipalFilterTest {
 
         AthenzUser conflictingIdentity = AthenzUser.fromUserId("mallory");
         DiscFilterRequest filterRequest = new ApplicationRequestToDiscFilterRequestWrapper(request, singletonList(createSelfSignedCertificate(conflictingIdentity)));
-        AthenzPrincipalFilter filter = new AthenzPrincipalFilter(validator, ATHENZ_PRINCIPAL_HEADER, CORS_ALLOWED_URLS);
+        AthenzPrincipalFilter filter = new AthenzPrincipalFilter(validator, Runnable::run, ATHENZ_PRINCIPAL_HEADER, CORS_ALLOWED_URLS);
         filter.filter(filterRequest, responseHandler);
 
         assertUnauthorized(responseHandler, "Identity in principal token does not match x509 CN");
@@ -176,7 +176,7 @@ public class AthenzPrincipalFilterTest {
         private final Map<NToken, AthenzPrincipal> validTokens = new HashMap<>();
 
         NTokenValidatorMock() {
-            super(keyId -> Optional.empty());
+            super((service, keyId) -> Optional.empty());
         }
 
         public NTokenValidatorMock add(NToken token, AthenzPrincipal principal) {
@@ -185,7 +185,7 @@ public class AthenzPrincipalFilterTest {
         }
 
         @Override
-        public AthenzPrincipal validate(NToken token) throws InvalidTokenException {
+        AthenzPrincipal validate(NToken token) throws InvalidTokenException {
             if (!validTokens.containsKey(token)) {
                 throw new InvalidTokenException("Invalid token");
             }
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/athenz/filter/AthenzTestUtils.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/athenz/filter/AthenzTestUtils.java
new file mode 100644
index 00000000000..40b38254dda
--- /dev/null
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/athenz/filter/AthenzTestUtils.java
@@ -0,0 +1,22 @@
+// Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.hosted.controller.athenz.filter;
+
+import java.security.KeyPair;
+import java.security.KeyPairGenerator;
+import java.security.NoSuchAlgorithmException;
+
+/**
+ * @author bjorncs
+ */
+public class AthenzTestUtils {
+    public static KeyPair generateRsaKeypair() {
+        try {
+            KeyPairGenerator keyGen = KeyPairGenerator.getInstance("RSA");
+            keyGen.initialize(512);
+            return keyGen.genKeyPair();
+        } catch (NoSuchAlgorithmException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+}
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/athenz/filter/NTokenValidatorTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/athenz/filter/NTokenValidatorTest.java
new file mode 100644
index 00000000000..510c806383c
--- /dev/null
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/athenz/filter/NTokenValidatorTest.java
@@ -0,0 +1,100 @@
+// Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.hosted.controller.athenz.filter;
+
+import com.yahoo.athenz.auth.token.PrincipalToken;
+import com.yahoo.vespa.athenz.api.AthenzIdentity;
+import com.yahoo.vespa.athenz.api.AthenzPrincipal;
+import com.yahoo.vespa.athenz.api.AthenzUser;
+import com.yahoo.vespa.athenz.api.NToken;
+import com.yahoo.vespa.hosted.controller.api.integration.athenz.InvalidTokenException;
+import com.yahoo.vespa.hosted.controller.api.integration.athenz.ZmsKeystore;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
+
+import java.security.KeyPair;
+import java.security.PrivateKey;
+import java.time.Instant;
+import java.util.Optional;
+
+import static com.yahoo.vespa.athenz.utils.AthenzIdentities.ZMS_ATHENZ_SERVICE;
+import static org.junit.Assert.assertEquals;
+
+/**
+ * @author bjorncs
+ */
+public class NTokenValidatorTest {
+
+    private static final KeyPair TRUSTED_KEY = AthenzTestUtils.generateRsaKeypair();
+    private static final KeyPair UNKNOWN_KEY = AthenzTestUtils.generateRsaKeypair();
+    private static final AthenzIdentity IDENTITY = AthenzUser.fromUserId("myuser");
+
+    @Rule
+    public ExpectedException exceptionRule = ExpectedException.none();
+
+    @Test
+    public void valid_token_is_accepted() throws InvalidTokenException {
+        NTokenValidator validator = new NTokenValidator(createKeystore());
+        NToken token = createNToken(IDENTITY, Instant.now(), TRUSTED_KEY.getPrivate(), "0");
+        AthenzPrincipal principal = validator.validate(token);
+        assertEquals("user.myuser", principal.getIdentity().getFullName());
+    }
+
+    @Test
+    public void invalid_signature_is_not_accepted() throws InvalidTokenException {
+        NTokenValidator validator = new NTokenValidator(createKeystore());
+        NToken token = createNToken(IDENTITY, Instant.now(), UNKNOWN_KEY.getPrivate(), "0");
+        exceptionRule.expect(InvalidTokenException.class);
+        exceptionRule.expectMessage("NToken is expired or has invalid signature");
+        validator.validate(token);
+    }
+
+    @Test
+    public void expired_token_is_not_accepted() throws InvalidTokenException {
+        NTokenValidator validator = new NTokenValidator(createKeystore());
+        NToken token = createNToken(IDENTITY, Instant.ofEpochMilli(1234) /*long time ago*/, TRUSTED_KEY.getPrivate(), "0");
+        exceptionRule.expect(InvalidTokenException.class);
+        exceptionRule.expectMessage("NToken is expired or has invalid signature");
+        validator.validate(token);
+    }
+
+    @Test
+    public void unknown_keyId_is_not_accepted() throws InvalidTokenException {
+        NTokenValidator validator = new NTokenValidator(createKeystore());
+        NToken token = createNToken(IDENTITY, Instant.now(), TRUSTED_KEY.getPrivate(), "unknown-key-id");
+        exceptionRule.expect(InvalidTokenException.class);
+        exceptionRule.expectMessage("NToken has an unknown keyId");
+        validator.validate(token);
+    }
+
+    @Test
+    public void failing_to_find_key_should_throw_exception() throws InvalidTokenException {
+        ZmsKeystore keystore = (athensService, keyId) -> { throw new  RuntimeException(); };
+        NTokenValidator validator = new NTokenValidator(keystore);
+        NToken token = createNToken(IDENTITY, Instant.now(), TRUSTED_KEY.getPrivate(), "0");
+        exceptionRule.expect(InvalidTokenException.class);
+        exceptionRule.expectMessage("Failed to retrieve public key");
+        validator.validate(token);
+    }
+
+    private static ZmsKeystore createKeystore() {
+        return (athensService, keyId) ->
+            athensService.equals(ZMS_ATHENZ_SERVICE) && keyId.equals("0")
+                    ? Optional.of(TRUSTED_KEY.getPublic())
+                    : Optional.empty();
+    }
+
+    private static NToken createNToken(AthenzIdentity identity, Instant issueTime, PrivateKey privateKey, String keyId) {
+        PrincipalToken token = new PrincipalToken.Builder("U1", identity.getDomain().getName(), identity.getName())
+                .keyId(keyId)
+                .salt("1234")
+                .host("host")
+                .ip("1.2.3.4")
+                .issueTime(issueTime.getEpochSecond())
+                .expirationWindow(1000)
+                .build();
+        token.sign(privateKey);
+        return new NToken(token.getSignedToken());
+    }
+
+}