Mock retrieving TLS certificate from unit test

1 files changed, 83 insertions, 34 deletions

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 64ae296b0a7..2d057875d24 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
@@ -14,10 +14,12 @@ import com.yahoo.config.provision.Flavor;
 import com.yahoo.config.provision.InstanceName;
 import com.yahoo.config.provision.NodeType;
 import com.yahoo.config.provision.RegionName;
+import com.yahoo.config.provision.SystemName;
 import com.yahoo.config.provision.TenantName;
 import com.yahoo.config.provision.Zone;
 import com.yahoo.log.LogLevel;
 import com.yahoo.vespa.hosted.athenz.instanceproviderservice.config.AthenzProviderServiceConfig;
+import com.yahoo.vespa.hosted.athenz.instanceproviderservice.impl.CertificateClient;
 import com.yahoo.vespa.hosted.athenz.instanceproviderservice.impl.IdentityDocumentGenerator;
 import com.yahoo.vespa.hosted.athenz.instanceproviderservice.impl.InstanceValidator;
 import com.yahoo.vespa.hosted.athenz.instanceproviderservice.impl.KeyProvider;
@@ -42,14 +44,25 @@ import org.apache.http.entity.StringEntity;
 import org.apache.http.impl.client.CloseableHttpClient;
 import org.apache.http.impl.client.HttpClients;
 import org.apache.http.ssl.SSLContextBuilder;
+import org.bouncycastle.asn1.ASN1ObjectIdentifier;
+import org.bouncycastle.asn1.x500.X500Name;
+import org.bouncycastle.asn1.x509.BasicConstraints;
+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.junit.Ignore;
+import org.bouncycastle.operator.ContentSigner;
+import org.bouncycastle.operator.OperatorCreationException;
+import org.bouncycastle.operator.jcajce.JcaContentSignerBuilder;
 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;
 import java.security.Key;
 import java.security.KeyManagementException;
@@ -60,8 +73,13 @@ import java.security.NoSuchAlgorithmException;
 import java.security.PrivateKey;
 import java.security.Signature;
 import java.security.SignatureException;
+import java.security.cert.CertificateException;
+import java.security.cert.X509Certificate;
 import java.time.Instant;
+import java.time.temporal.TemporalAmount;
 import java.util.Base64;
+import java.util.Calendar;
+import java.util.Date;
 import java.util.HashSet;
 import java.util.Optional;
 import java.util.logging.Logger;
@@ -84,19 +102,21 @@ public class AthenzInstanceProviderServiceTest {
     private static final int PORT = 12345;
 
     @Test
-    @Ignore("Requires private key for Athenz service")
     public void provider_service_hosts_endpoint_secured_with_tls() throws Exception {
-        String domain = "INSERT DOMAIN HERE";
-        String service = "INSERT SERVICE NAME HERE";
-        DummyKeyProvider keyProvider = new DummyKeyProvider();
+        String domain = "domain";
+        String service = "service";
+        AutoGeneratedKeyProvider keyProvider = new AutoGeneratedKeyProvider();
         PrivateKey privateKey = Crypto.loadPrivateKey(keyProvider.getPrivateKey(0));
-
-        AthenzProviderServiceConfig config = getAthenzProviderConfig(domain, service, "INSERT ZTS URL HERE", "INSERT DNS SUFFIX HERE");
-
+        AthenzProviderServiceConfig config = getAthenzProviderConfig(domain, service);
         ScheduledExecutorServiceMock executor = new ScheduledExecutorServiceMock();
-        NodeRepository nodeRepository = mock(NodeRepository.class);
-        Zone zone = new Zone(Environment.dev, RegionName.from("us-north-1"));
-        AthenzInstanceProviderService athenzInstanceProviderService = new AthenzInstanceProviderService(config, keyProvider, executor, nodeRepository, zone);
+
+        AthenzInstanceProviderService athenzInstanceProviderService =
+                new AthenzInstanceProviderService(config,
+                                                  keyProvider,
+                                                  executor,
+                                                  mock(NodeRepository.class),
+                                                  new Zone(SystemName.cd, Environment.dev, RegionName.from("us-north-1")),
+                                                  new SelfSignedCertificateClient(keyProvider.getKeyPair(), config));
 
         try (CloseableHttpClient client = createHttpClient(domain, service)) {
             Runnable certificateRefreshCommand = executor.getCommand().orElseThrow(() -> new AssertionError("Command not present"));
@@ -116,7 +136,7 @@ public class AthenzInstanceProviderServiceTest {
     public void generates_valid_identity_document() throws IOException, NoSuchAlgorithmException, KeyStoreException, KeyManagementException {
         String hostname = "x.y.com";
         AutoGeneratedKeyProvider keyProvider = new AutoGeneratedKeyProvider();
-        AthenzProviderServiceConfig config = getAthenzProviderConfig("domain", "service", "localhost/zts", "dnsSuffix");
+        AthenzProviderServiceConfig config = getAthenzProviderConfig("domain", "service");
 
         NodeRepository nodeRepository = mock(NodeRepository.class);
         MockNodeFlavors nodeFlavors = new MockNodeFlavors();
@@ -142,16 +162,16 @@ public class AthenzInstanceProviderServiceTest {
 
     }
 
-    private AthenzProviderServiceConfig getAthenzProviderConfig(String domain, String service, String ztsUrl, String dnsSuffix) {
+    private static AthenzProviderServiceConfig getAthenzProviderConfig(String domain, String service) {
         return new AthenzProviderServiceConfig(
                         new AthenzProviderServiceConfig.Builder()
                                 .domain(domain)
                                 .serviceName(service)
                                 .port(PORT)
                                 .keyPathPrefix("dummy-path")
-                                .certDnsSuffix(dnsSuffix)
-                                .ztsUrl(ztsUrl)
-                                .athenzPrincipalHeaderName("INSERT PRINCIPAL HEADER NAME HERE")
+                                .certDnsSuffix("dnsSuffix")
+                                .ztsUrl("localhost/zts")
+                                .athenzPrincipalHeaderName("Yahoo-Principal-Auth")
                                 .apiPath("/"));
 
     }
@@ -212,39 +232,26 @@ public class AthenzInstanceProviderServiceTest {
         }
     }
 
-    private static class DummyKeyProvider implements KeyProvider {
-
-        @Override
-        public String getPrivateKey(int version) {
-            return "INSERT PRIV KEY";
-        }
-
-        @Override
-        public String getPublicKey(int version) {
-            return "INSERT PUB KEY";
-        }
-    }
-
     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 keyPair = rsa.genKeyPair();
-            publicKey = pemEncode("RSA PUBLIC KEY", keyPair.getPublic());
-            privateKey = pemEncode("RSA PRIVATE KEY", keyPair.getPrivate());
+            keyPair = rsa.genKeyPair();
+            publicKey = pemEncode(keyPair.getPublic());
+            privateKey = pemEncode(keyPair.getPrivate());
         }
 
-        private String pemEncode(String description, Key key) throws IOException {
+        private static String pemEncode(Key key) throws IOException {
             StringWriter stringWriter = new StringWriter();
             JcaPEMWriter pemWriter = new JcaPEMWriter(stringWriter);
             pemWriter.writeObject(key);
             pemWriter.flush();
             return stringWriter.toString();
-
         }
 
         @Override
@@ -256,5 +263,47 @@ public class AthenzInstanceProviderServiceTest {
         public String getPublicKey(int version) {
             return publicKey;
         }
+
+        public KeyPair getKeyPair() {
+            return keyPair;
+        }
+    }
+
+    private static class SelfSignedCertificateClient implements CertificateClient {
+
+        private final KeyPair keyPair;
+        private final AthenzProviderServiceConfig config;
+
+        private SelfSignedCertificateClient(KeyPair keyPair, AthenzProviderServiceConfig config) {
+            this.keyPair = keyPair;
+            this.config = config;
+        }
+
+        @Override
+        public String updateCertificate(PrivateKey privateKey, TemporalAmount expiryTime) {
+            try {
+                ContentSigner contentSigner = new JcaContentSignerBuilder("SHA512WithRSA").build(keyPair.getPrivate());
+                X500Name dnName = new X500Name("CN=" + config.domain() + "." + config.serviceName());
+                Calendar calendar = Calendar.getInstance();
+                calendar.add(Calendar.HOUR, 1);
+                Date endDate = calendar.getTime();
+                JcaX509v3CertificateBuilder certBuilder = new JcaX509v3CertificateBuilder(
+                        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().setProvider(new BouncyCastleProvider())
+                        .getCertificate(certBuilder.build(contentSigner));
+                try (StringWriter stringWriter = new StringWriter();
+                     JcaPEMWriter pemWriter = new JcaPEMWriter(stringWriter)) {
+                    pemWriter.writeObject(certificate);
+                    pemWriter.flush();
+                    return stringWriter.toString();
+                }
+            } catch (CertificateException | CertIOException | OperatorCreationException e) {
+                throw new RuntimeException(e);
+            } catch (IOException e) {
+                throw new UncheckedIOException(e);
+            }
+        }
     }
 }