Don't use String as type for private keys and certificates in interfaces

8 files changed, 34 insertions, 45 deletions

diff --git a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/AthenzInstanceProviderService.java b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/AthenzInstanceProviderService.java
index fc1c995f5c1..745aae62d3a 100644
--- a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/AthenzInstanceProviderService.java
+++ b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/AthenzInstanceProviderService.java
@@ -136,8 +136,8 @@ public class AthenzInstanceProviderService extends AbstractComponent {
         public void run() {
             try {
                 log.log(LogLevel.INFO, "Updating Athenz certificate through ZTS");
-                PrivateKey privateKey = Crypto.loadPrivateKey(keyProvider.getPrivateKey(config.keyVersion()));
-                X509Certificate certificate = Crypto.loadX509Certificate(certificateClient.updateCertificate(privateKey, EXPIRY_TIME));
+                PrivateKey privateKey = keyProvider.getPrivateKey(config.keyVersion());
+                X509Certificate certificate = certificateClient.updateCertificate(privateKey, EXPIRY_TIME);
 
                 String dummyPassword = "athenz";
                 KeyStore keyStore = KeyStore.getInstance("JKS");
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 889a63c9e10..031133ade19 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
@@ -3,11 +3,13 @@ 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.vespa.hosted.athenz.instanceproviderservice.config.AthenzProviderServiceConfig;
 
 import java.security.PrivateKey;
+import java.security.cert.X509Certificate;
 import java.time.temporal.ChronoUnit;
 import java.time.temporal.TemporalAmount;
 import java.util.concurrent.TimeUnit;
@@ -26,7 +28,7 @@ public class AthenzCertificateClient implements CertificateClient {
     }
 
     @Override
-    public String updateCertificate(PrivateKey privateKey, TemporalAmount expiryTime) {
+    public X509Certificate updateCertificate(PrivateKey privateKey, TemporalAmount expiryTime) {
         SimpleServiceIdentityProvider identityProvider = new SimpleServiceIdentityProvider(
                 authority, config.domain(), config.serviceName(),
                 privateKey, Integer.toString(config.keyVersion()), TimeUnit.MINUTES.toSeconds(10));
@@ -36,8 +38,9 @@ public class AthenzCertificateClient implements CertificateClient {
                 ZTSClient.generateInstanceRefreshRequest(
                         config.domain(), config.serviceName(), privateKey,
                         config.certDnsSuffix(), (int)expiryTime.get(ChronoUnit.SECONDS));
-        return ztsClient.postInstanceRefreshRequest(config.domain(), config.serviceName(), req)
+        String pemEncoded = ztsClient.postInstanceRefreshRequest(config.domain(), config.serviceName(), req)
                 .getCertificate();
+        return Crypto.loadX509Certificate(pemEncoded);
     }
 
     private static class AthenzPrincipalAuthority extends PrincipalAuthority {
diff --git a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/CertificateClient.java b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/CertificateClient.java
index 30d918b4235..6465873e092 100644
--- a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/CertificateClient.java
+++ b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/CertificateClient.java
@@ -2,6 +2,7 @@
 package com.yahoo.vespa.hosted.athenz.instanceproviderservice.impl;
 
 import java.security.PrivateKey;
+import java.security.cert.X509Certificate;
 import java.time.temporal.TemporalAmount;
 
 /**
@@ -9,5 +10,5 @@ import java.time.temporal.TemporalAmount;
  */
 @FunctionalInterface
 public interface CertificateClient {
-    String updateCertificate(PrivateKey privateKey, TemporalAmount expiryTime);
+    X509Certificate updateCertificate(PrivateKey privateKey, TemporalAmount expiryTime);
 }
diff --git a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/FileBackedKeyProvider.java b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/FileBackedKeyProvider.java
index f03f8415586..40a2a1dbcc9 100644
--- a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/FileBackedKeyProvider.java
+++ b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/FileBackedKeyProvider.java
@@ -1,10 +1,14 @@
 // 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.util.Crypto;
+
 import java.io.File;
 import java.io.IOException;
 import java.io.UncheckedIOException;
 import java.nio.file.Files;
+import java.security.PrivateKey;
+import java.security.PublicKey;
 
 /**
  * @author bjorncs
@@ -18,16 +22,16 @@ public class FileBackedKeyProvider implements KeyProvider {
     }
 
     @Override
-    public String getPrivateKey(int version) {
-        return loadKey(new File(keyPathPrefix + ".priv." + version));
+    public PrivateKey getPrivateKey(int version) {
+        return Crypto.loadPrivateKey(readPemStringFromFile(new File(keyPathPrefix + ".priv." + version)));
     }
 
     @Override
-    public String getPublicKey(int version) {
-        return loadKey(new File(keyPathPrefix + ".pub." + version));
+    public PublicKey getPublicKey(int version) {
+        return Crypto.loadPublicKey(readPemStringFromFile(new File(keyPathPrefix + ".pub." + version)));
     }
 
-    private static String loadKey(File file) {
+    private static String readPemStringFromFile(File file) {
         try {
             if (!file.exists() || !file.isFile()) {
                 throw new IllegalArgumentException("Key missing: " + file.getAbsolutePath());
diff --git a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/IdentityDocumentGenerator.java b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/IdentityDocumentGenerator.java
index 4669563d8df..644fc929884 100644
--- a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/IdentityDocumentGenerator.java
+++ b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/IdentityDocumentGenerator.java
@@ -48,7 +48,7 @@ public class IdentityDocumentGenerator {
             Signature sigGenerator = Signature.getInstance("SHA512withRSA");
 
             // TODO: Get the correct version 0 ok for now
-            PrivateKey privateKey = Crypto.loadPrivateKey(keyProvider.getPrivateKey(0));
+            PrivateKey privateKey = keyProvider.getPrivateKey(0);
             sigGenerator.initSign(privateKey);
             sigGenerator.update(encodedIdentityDocument.getBytes());
             String signature = Base64.getEncoder().encodeToString(sigGenerator.sign());
diff --git a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/InstanceValidator.java b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/InstanceValidator.java
index f5c2c319041..8d76300c2bb 100644
--- a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/InstanceValidator.java
+++ b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/InstanceValidator.java
@@ -1,7 +1,6 @@
 // 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.util.Crypto;
 import com.yahoo.log.LogLevel;
 import com.yahoo.vespa.hosted.athenz.instanceproviderservice.impl.model.InstanceConfirmation;
 import com.yahoo.vespa.hosted.athenz.instanceproviderservice.impl.model.ProviderUniqueId;
@@ -34,7 +33,7 @@ public class InstanceValidator {
         SignedIdentityDocument signedIdentityDocument = instanceConfirmation.signedIdentityDocument;
         ProviderUniqueId providerUniqueId = signedIdentityDocument.identityDocument.providerUniqueId;
         log.log(LogLevel.INFO, () -> String.format("Validating instance %s.", providerUniqueId));
-        PublicKey publicKey = Crypto.loadPublicKey(keyProvider.getPublicKey(signedIdentityDocument.signingKeyVersion));
+        PublicKey publicKey = keyProvider.getPublicKey(signedIdentityDocument.signingKeyVersion);
         if (isSignatureValid(publicKey, signedIdentityDocument.rawIdentityDocument, signedIdentityDocument.signature)) {
             log.log(LogLevel.INFO, () -> String.format("Instance %s is valid.", providerUniqueId));
             return true;
diff --git a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/KeyProvider.java b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/KeyProvider.java
index 8c807405693..5a1d7e3c1ff 100644
--- a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/KeyProvider.java
+++ b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/KeyProvider.java
@@ -1,11 +1,14 @@
 // 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 java.security.PrivateKey;
+import java.security.PublicKey;
+
 /**
  * @author bjorncs
  */
 public interface KeyProvider {
-    String getPrivateKey(int version);
+    PrivateKey getPrivateKey(int version);
 
-    String getPublicKey(int version);
+    PublicKey getPublicKey(int version);
 }
diff --git a/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/AthenzInstanceProviderServiceTest.java b/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/AthenzInstanceProviderServiceTest.java
index 0fe31e4ff06..a15d681de39 100644
--- a/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/AthenzInstanceProviderServiceTest.java
+++ b/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/AthenzInstanceProviderServiceTest.java
@@ -4,7 +4,6 @@ package com.yahoo.vespa.hosted.athenz.instanceproviderservice;
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.collect.ImmutableSet;
-import com.yahoo.athenz.auth.util.Crypto;
 import com.yahoo.component.Version;
 import com.yahoo.config.provision.ApplicationId;
 import com.yahoo.config.provision.ApplicationName;
@@ -50,7 +49,6 @@ import org.bouncycastle.cert.CertIOException;
 import org.bouncycastle.cert.jcajce.JcaX509CertificateConverter;
 import org.bouncycastle.cert.jcajce.JcaX509v3CertificateBuilder;
 import org.bouncycastle.jce.provider.BouncyCastleProvider;
-import org.bouncycastle.openssl.jcajce.JcaPEMWriter;
 import org.bouncycastle.operator.ContentSigner;
 import org.bouncycastle.operator.OperatorCreationException;
 import org.bouncycastle.operator.jcajce.JcaContentSignerBuilder;
@@ -58,8 +56,6 @@ import org.junit.Test;
 
 import javax.net.ssl.SSLContext;
 import java.io.IOException;
-import java.io.StringWriter;
-import java.io.UncheckedIOException;
 import java.io.UnsupportedEncodingException;
 import java.math.BigInteger;
 import java.security.InvalidKeyException;
@@ -69,6 +65,7 @@ import java.security.KeyPairGenerator;
 import java.security.KeyStoreException;
 import java.security.NoSuchAlgorithmException;
 import java.security.PrivateKey;
+import java.security.PublicKey;
 import java.security.Signature;
 import java.security.SignatureException;
 import java.security.cert.CertificateException;
@@ -105,7 +102,7 @@ public class AthenzInstanceProviderServiceTest {
         String domain = "domain";
         String service = "service";
         AutoGeneratedKeyProvider keyProvider = new AutoGeneratedKeyProvider();
-        PrivateKey privateKey = Crypto.loadPrivateKey(keyProvider.getPrivateKey(0));
+        PrivateKey privateKey = keyProvider.getPrivateKey(0);
         AthenzProviderServiceConfig config = getAthenzProviderConfig(domain, service);
         ScheduledExecutorServiceMock executor = new ScheduledExecutorServiceMock();
 
@@ -173,7 +170,7 @@ public class AthenzInstanceProviderServiceTest {
         assertEquals(expectedProviderUniqueId, signedIdentityDocument.identityDocument.providerUniqueId);
 
         // Validate signature
-        assertTrue("Message", InstanceValidator.isSignatureValid(Crypto.loadPublicKey(keyProvider.getPublicKey(0)),
+        assertTrue("Message", InstanceValidator.isSignatureValid(keyProvider.getPublicKey(0),
                                                                  signedIdentityDocument.rawIdentityDocument,
                                                                  signedIdentityDocument.signature));
     }
@@ -258,25 +255,21 @@ public class AthenzInstanceProviderServiceTest {
     private static class AutoGeneratedKeyProvider implements KeyProvider {
 
         private final KeyPair keyPair;
-        private final String publicKey;
-        private final String privateKey;
 
         public AutoGeneratedKeyProvider() throws IOException, NoSuchAlgorithmException {
             KeyPairGenerator rsa = KeyPairGenerator.getInstance("RSA");
             rsa.initialize(2048);
             keyPair = rsa.genKeyPair();
-            publicKey = toPemString(keyPair.getPublic());
-            privateKey = toPemString(keyPair.getPrivate());
         }
 
         @Override
-        public String getPrivateKey(int version) {
-            return privateKey;
+        public PrivateKey getPrivateKey(int version) {
+            return keyPair.getPrivate();
         }
 
         @Override
-        public String getPublicKey(int version) {
-            return publicKey;
+        public PublicKey getPublicKey(int version) {
+            return keyPair.getPublic();
         }
 
         public KeyPair getKeyPair() {
@@ -295,7 +288,7 @@ public class AthenzInstanceProviderServiceTest {
         }
 
         @Override
-        public String updateCertificate(PrivateKey privateKey, TemporalAmount expiryTime) {
+        public X509Certificate updateCertificate(PrivateKey privateKey, TemporalAmount expiryTime) {
             try {
                 ContentSigner contentSigner = new JcaContentSignerBuilder("SHA512WithRSA").build(keyPair.getPrivate());
                 X500Name dnName = new X500Name("CN=" + config.domain() + "." + config.serviceName());
@@ -306,26 +299,12 @@ public class AthenzInstanceProviderServiceTest {
                         dnName, BigInteger.ONE, new Date(), endDate, dnName, keyPair.getPublic());
                 certBuilder.addExtension(new ASN1ObjectIdentifier("2.5.29.19"), true, new BasicConstraints(true));
 
-                X509Certificate certificate = new JcaX509CertificateConverter()
+                return new JcaX509CertificateConverter()
                         .setProvider(new BouncyCastleProvider())
                         .getCertificate(certBuilder.build(contentSigner));
-                return toPemString(certificate);
             } catch (CertificateException | CertIOException | OperatorCreationException e) {
                 throw new RuntimeException(e);
             }
         }
     }
-
-    private static String toPemString(Object keyOrCertificate) {
-        try (StringWriter stringWriter = new StringWriter();
-             JcaPEMWriter pemWriter = new JcaPEMWriter(stringWriter)) {
-            pemWriter.writeObject(keyOrCertificate);
-            pemWriter.flush();
-            return stringWriter.toString();
-        } catch (CertIOException e) {
-            throw new RuntimeException(e);
-        } catch (IOException e) {
-            throw new UncheckedIOException(e);
-        }
-    }
 }