Merge pull request #21714 from vespa-engine/hakonhall/add-ntokengenerator

Add NTokenGenerator

3 files changed, 176 insertions, 5 deletions

diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/api/NTokenGenerator.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/api/NTokenGenerator.java
new file mode 100644
index 00000000000..93d2f9bf3b5
--- /dev/null
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/api/NTokenGenerator.java
@@ -0,0 +1,114 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.athenz.api;
+
+import com.yahoo.athenz.auth.util.Crypto;
+import com.yahoo.athenz.auth.util.CryptoException;
+
+import java.security.PrivateKey;
+import java.security.SecureRandom;
+import java.time.Clock;
+import java.time.Duration;
+import java.util.function.LongSupplier;
+
+import static java.lang.String.format;
+import static java.util.Objects.requireNonNull;
+
+/**
+ * @author hakonhall
+ */
+public class NTokenGenerator {
+    private final Signer signer;
+    private final Clock clock;
+
+    private String domain = null;
+    private String name = null;
+    private String keyVersion = null;
+    private String keyService = null;
+    private String hostname = null;
+    private String ip = null;
+
+    private final StringBuilder token = new StringBuilder();
+    private final LongSupplier randomGenerator;
+
+    @FunctionalInterface
+    public interface Signer {
+        /** @see Crypto#sign(String, PrivateKey) */
+        String sign(String message, PrivateKey key) throws CryptoException;
+    }
+
+    public NTokenGenerator() { this(Crypto::sign, Clock.systemUTC(), new SecureRandom()::nextLong); }
+
+    /** For testing. */
+    NTokenGenerator(Signer signer, Clock clock, LongSupplier randomGenerator) {
+        this.signer = signer;
+        this.clock = clock;
+        this.randomGenerator = randomGenerator;
+    }
+
+    /** Required. */
+    public NTokenGenerator setIdentity(AthenzIdentity identity) {
+        this.domain = identity.getDomainName();
+        this.name = identity.getName();
+        return this;
+    }
+
+    /** Required. */
+    public NTokenGenerator setKeyVersion(String keyVersion) {
+        this.keyVersion = requireNonNull(keyVersion);
+        return this;
+    }
+
+    public NTokenGenerator setKeyService(String keyService) {
+        this.keyService = requireNonNull(keyService);
+        return this;
+    }
+
+    public NTokenGenerator setHostname(String hostname) {
+        this.hostname = requireNonNull(hostname);
+        return this;
+    }
+
+    public NTokenGenerator setIp(String ip) {
+        this.ip = requireNonNull(ip);
+        return this;
+    }
+
+    public NToken sign(PrivateKey privateKey) {
+        // See https://github.com/AthenZ/athenz/blob/master/libs/go/zmssvctoken/token.go
+
+        var generationTime = clock.instant();
+
+        token.setLength(0);
+        append('v', "S1");
+        append('d', domain);
+        append('n', name);
+        append('k', keyVersion);
+        append('z', keyService, false);
+        append('h', hostname, false);
+        append('i', ip, false);
+        append('a', format("%x", randomGenerator.getAsLong()));
+        append('t', format("%d", generationTime.getEpochSecond()));
+        append('e', format("%d", generationTime.plus(Duration.ofMinutes(10)).getEpochSecond()));
+        append('s', signer.sign(token.toString(), privateKey));
+
+        return new NToken(token.toString());
+    }
+
+    private void append(char name, String value) { append(name, value, true); }
+
+    private void append(char name, String value, boolean required) {
+        if (value == null) {
+            if (required) {
+                throw new IllegalStateException("Missing value for NToken key " + name);
+            } else {
+                return;
+            }
+        }
+
+        if (token.length() > 0) {
+            token.append(';');
+        }
+
+        token.append(name).append('=').append(value);
+    }
+}
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zts/DefaultZtsClient.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zts/DefaultZtsClient.java
index 2eea5d3151a..0c73891bdae 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zts/DefaultZtsClient.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zts/DefaultZtsClient.java
@@ -94,12 +94,17 @@ public class DefaultZtsClient extends ClientBase implements ZtsClient {
 
     @Override
     public Identity getServiceIdentity(AthenzIdentity identity, String keyId, Pkcs10Csr csr) {
+        return getServiceIdentity(identity, keyId, csr, Optional.empty());
+    }
+
+    public Identity getServiceIdentity(AthenzIdentity identity, String keyId, Pkcs10Csr csr, Optional<NToken> nToken) {
         URI uri = ztsUrl.resolve(String.format("instance/%s/%s/refresh", identity.getDomainName(), identity.getName()));
-        HttpUriRequest request = RequestBuilder.post()
-                .setUri(uri)
-                .setEntity(toJsonStringEntity(new IdentityRefreshRequestEntity(csr, keyId)))
-                .build();
-        return execute(request, response -> {
+        RequestBuilder builder = RequestBuilder.post()
+                                               .setUri(uri)
+                                               .setEntity(toJsonStringEntity(new IdentityRefreshRequestEntity(csr, keyId)));
+        nToken.ifPresent(n -> builder.setHeader("Athenz-Principal-Auth", n.getRawToken()));
+
+        return execute(builder.build(), response -> {
             IdentityResponseEntity entity = readEntity(response, IdentityResponseEntity.class);
             return new Identity(entity.certificate(), entity.caCertificateBundle());
         });
diff --git a/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/api/NTokenGeneratorTest.java b/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/api/NTokenGeneratorTest.java
new file mode 100644
index 00000000000..7c4ee36926f
--- /dev/null
+++ b/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/api/NTokenGeneratorTest.java
@@ -0,0 +1,52 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.athenz.api;
+
+import com.yahoo.athenz.auth.util.CryptoException;
+import com.yahoo.test.ManualClock;
+import org.junit.Test;
+
+import java.security.PrivateKey;
+import java.time.Instant;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertSame;
+import static org.mockito.Mockito.mock;
+
+/**
+ * @author hakonhall
+ */
+public class NTokenGeneratorTest {
+    private final PrivateKey key = mock(PrivateKey.class);
+
+    @Test
+    public void ntoken() {
+        var signer = new Signer("signature");
+        NTokenGenerator generator = new NTokenGenerator(signer, new ManualClock(Instant.ofEpochSecond(12L)), () -> 3L);
+        AthenzIdentity identity = new AthenzService("domain", "service");
+
+        NToken token = generator.setIdentity(identity)
+                                .setKeyVersion("0")
+                                .sign(key);
+
+        assertEquals("v=S1;d=domain;n=service;k=0;a=3;t=12;e=612", signer.message);
+        assertSame(key, signer.key);
+        assertEquals("v=S1;d=domain;n=service;k=0;a=3;t=12;e=612;s=signature", token.getRawToken());
+    }
+
+    private static class Signer implements NTokenGenerator.Signer {
+        private final String signature;
+        public String message = null;
+        public PrivateKey key = null;
+
+        public Signer(String signature) {
+            this.signature = signature;
+        }
+
+        @Override
+        public String sign(String message, PrivateKey key) throws CryptoException {
+            this.message = message;
+            this.key = key;
+            return signature;
+        }
+    }
+}
\ No newline at end of file