1 files changed, 28 insertions, 11 deletions

diff --git a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/AthenzCertificateClient.java b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/AthenzCertificateClient.java
index 62c7038a265..c849de481dc 100644
--- a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/AthenzCertificateClient.java
+++ b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/AthenzCertificateClient.java
@@ -1,15 +1,18 @@
 // 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.athenz.instanceproviderservice.impl;
 
+import com.yahoo.athenz.auth.impl.PrincipalAuthority;
+import com.yahoo.athenz.auth.impl.SimpleServiceIdentityProvider;
 import com.yahoo.athenz.auth.util.Crypto;
 import com.yahoo.athenz.zts.InstanceRefreshRequest;
 import com.yahoo.athenz.zts.ZTSClient;
-import com.yahoo.container.jdisc.athenz.AthenzIdentityProvider;
 import com.yahoo.vespa.hosted.athenz.instanceproviderservice.config.AthenzProviderServiceConfig;
 
-import javax.net.ssl.SSLContext;
 import java.security.PrivateKey;
 import java.security.cert.X509Certificate;
+import java.time.temporal.ChronoUnit;
+import java.time.temporal.TemporalAmount;
+import java.util.concurrent.TimeUnit;
 
 /**
  * @author bjorncs
@@ -17,27 +20,41 @@ import java.security.cert.X509Certificate;
 public class AthenzCertificateClient {
 
     private final AthenzProviderServiceConfig config;
+    private final AthenzPrincipalAuthority authority;
     private final AthenzProviderServiceConfig.Zones zoneConfig;
-    private final AthenzIdentityProvider bootstrapIdentity;
 
-    public AthenzCertificateClient(AthenzIdentityProvider bootstrapIdentity,
-                                   AthenzProviderServiceConfig config,
-                                   AthenzProviderServiceConfig.Zones zoneConfig) {
-        this.bootstrapIdentity = bootstrapIdentity;
+    public AthenzCertificateClient(AthenzProviderServiceConfig config, AthenzProviderServiceConfig.Zones zoneConfig) {
         this.config = config;
+        this.authority = new AthenzPrincipalAuthority(config.athenzPrincipalHeaderName());
         this.zoneConfig = zoneConfig;
     }
 
     public X509Certificate updateCertificate(PrivateKey privateKey) {
-        SSLContext bootstrapSslContext = bootstrapIdentity.getIdentitySslContext();
-        ZTSClient ztsClient = new ZTSClient(config.ztsUrl(), bootstrapSslContext);
+        SimpleServiceIdentityProvider identityProvider = new SimpleServiceIdentityProvider(
+                authority, zoneConfig.domain(), zoneConfig.serviceName(),
+                privateKey, Integer.toString(zoneConfig.secretVersion()), TimeUnit.MINUTES.toSeconds(10));
+        ZTSClient ztsClient = new ZTSClient(
+                config.ztsUrl(), zoneConfig.domain(), zoneConfig.serviceName(), identityProvider);
         InstanceRefreshRequest req =
                 ZTSClient.generateInstanceRefreshRequest(
-                        zoneConfig.domain(), zoneConfig.serviceName(), privateKey, zoneConfig.certDnsSuffix(), /*expiryTime*/0);
-        req.setKeyId(Integer.toString(zoneConfig.secretVersion()));
+                        zoneConfig.domain(), zoneConfig.serviceName(), privateKey,
+                        config.certDnsSuffix(), /*expiryTime*/0);
         String pemEncoded = ztsClient.postInstanceRefreshRequest(zoneConfig.domain(), zoneConfig.serviceName(), req)
                 .getCertificate();
         return Crypto.loadX509Certificate(pemEncoded);
     }
 
+    private static class AthenzPrincipalAuthority extends PrincipalAuthority {
+        private final String headerName;
+
+        public AthenzPrincipalAuthority(String headerName) {
+            this.headerName = headerName;
+        }
+
+        @Override
+        public String getHeader() {
+            return headerName;
+        }
+    }
+
 }