Handle zms keys in addition to zts keys

4 files changed, 42 insertions, 23 deletions

diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/utils/ntoken/AthenzConfTruststore.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/utils/ntoken/AthenzConfTruststore.java
index 64cdb13d66c..4cb3470635e 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/utils/ntoken/AthenzConfTruststore.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/utils/ntoken/AthenzConfTruststore.java
@@ -21,32 +21,37 @@ import java.util.Optional;
  */
 public class AthenzConfTruststore implements AthenzTruststore {
 
-    private final Map<String, PublicKey> publicKeys;
+    private final Map<String, PublicKey> zmsPublicKeys;
+    private final Map<String, PublicKey> ztsPublicKeys;
 
     public AthenzConfTruststore(Path athenzConfFile) {
-        this.publicKeys = loadPublicKeys(athenzConfFile);
+        try {
+            JsonNode root = new ObjectMapper().readTree(athenzConfFile.toFile());
+            this.zmsPublicKeys = loadPublicKeys((ArrayNode) root.get("zmsPublicKeys"));
+            this.ztsPublicKeys = loadPublicKeys((ArrayNode) root.get("ztsPublicKeys"));
+        } catch (IOException e) {
+            throw new UncheckedIOException(e);
+        }
+    }
+
+    private static Map<String, PublicKey> loadPublicKeys(ArrayNode keysArray) {
+        Map<String, PublicKey> publicKeys = new HashMap<>();
+        for (JsonNode keyEntry : keysArray) {
+            String keyId = keyEntry.get("id").textValue();
+            String encodedPublicKey = keyEntry.get("key").textValue();
+            PublicKey publicKey = Crypto.loadPublicKey(Crypto.ybase64DecodeString(encodedPublicKey));
+            publicKeys.put(keyId, publicKey);
+        }
+        return publicKeys;
     }
 
     @Override
-    public Optional<PublicKey> getPublicKey(String keyId) {
-        return Optional.ofNullable(publicKeys.get(keyId));
+    public Optional<PublicKey> getZmsPublicKey(String keyId) {
+        return Optional.ofNullable(zmsPublicKeys.get(keyId));
     }
 
-    private static Map<String, PublicKey> loadPublicKeys(Path athenzConfFile) {
-        try {
-            Map<String, PublicKey> publicKeys = new HashMap<>();
-            ObjectMapper mapper = new ObjectMapper();
-            JsonNode root = mapper.readTree(athenzConfFile.toFile());
-            ArrayNode keysArray = (ArrayNode) root.get("ztsPublicKeys");
-            for (JsonNode keyEntry : keysArray) {
-                String keyId = keyEntry.get("id").textValue();
-                String encodedPublicKey = keyEntry.get("key").textValue();
-                PublicKey publicKey = Crypto.loadPublicKey(Crypto.ybase64DecodeString(encodedPublicKey));
-                publicKeys.put(keyId, publicKey);
-            }
-            return publicKeys;
-        } catch (IOException e) {
-            throw new UncheckedIOException(e);
-        }
+    @Override
+    public Optional<PublicKey> getZtsPublicKey(String keyId) {
+        return Optional.ofNullable(ztsPublicKeys.get(keyId));
     }
 }
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/utils/ntoken/AthenzTruststore.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/utils/ntoken/AthenzTruststore.java
index 3139e6b847f..83afa288cf0 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/utils/ntoken/AthenzTruststore.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/utils/ntoken/AthenzTruststore.java
@@ -10,5 +10,6 @@ import java.util.Optional;
  * @author bjorncs
  */
 public interface AthenzTruststore {
-   Optional<PublicKey> getPublicKey(String keyId);
+   Optional<PublicKey> getZmsPublicKey(String keyId);
+   Optional<PublicKey> getZtsPublicKey(String keyId);
 }
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/utils/ntoken/NTokenValidator.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/utils/ntoken/NTokenValidator.java
index 85c62bc07ff..f4ec0b168d7 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/utils/ntoken/NTokenValidator.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/utils/ntoken/NTokenValidator.java
@@ -39,7 +39,8 @@ public class NTokenValidator {
     public AthenzPrincipal validate(NToken token) throws InvalidTokenException {
         PrincipalToken principalToken = new PrincipalToken(token.getRawToken());
         String keyId = principalToken.getKeyId();
-        PublicKey zmsPublicKey = truststore.getPublicKey(keyId)
+        String keyService = principalToken.getKeyService();
+        PublicKey zmsPublicKey = (keyService == null || keyService.equals("zms") ? truststore.getZmsPublicKey(keyId) : truststore.getZtsPublicKey(keyId))
                 .orElseThrow(() -> {
                     String message = "NToken has an unknown keyId: " + keyId;
                     log.log(LogLevel.WARNING, message);
diff --git a/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/utils/ntoken/NTokenValidatorTest.java b/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/utils/ntoken/NTokenValidatorTest.java
index 0e70993792f..22f97ca8b60 100644
--- a/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/utils/ntoken/NTokenValidatorTest.java
+++ b/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/utils/ntoken/NTokenValidatorTest.java
@@ -15,6 +15,7 @@ import org.junit.rules.ExpectedException;
 
 import java.security.KeyPair;
 import java.security.PrivateKey;
+import java.security.PublicKey;
 import java.time.Instant;
 import java.util.Optional;
 
@@ -68,7 +69,17 @@ public class NTokenValidatorTest {
     }
 
     private static AthenzTruststore createTruststore() {
-        return keyId -> keyId.equals("0") ? Optional.of(TRUSTED_KEY.getPublic()) : Optional.empty();
+        return new AthenzTruststore() {
+            @Override
+            public Optional<PublicKey> getZmsPublicKey(String keyId) {
+                return keyId.equals("0") ? Optional.of(TRUSTED_KEY.getPublic()) : Optional.empty();
+            }
+
+            @Override
+            public Optional<PublicKey> getZtsPublicKey(String keyId) {
+                return Optional.empty();
+            }
+        };
     }
 
     private static NToken createNToken(AthenzIdentity identity, Instant issueTime, PrivateKey privateKey, String keyId) {
@@ -77,6 +88,7 @@ public class NTokenValidatorTest {
                 .salt("1234")
                 .host("host")
                 .ip("1.2.3.4")
+                .keyService("zms")
                 .issueTime(issueTime.getEpochSecond())
                 .expirationWindow(1000)
                 .build();