Implement ZTS instance refresh request

2 files changed, 100 insertions, 6 deletions

diff --git a/container-disc/src/main/java/com/yahoo/container/jdisc/athenz/impl/AthenzService.java b/container-disc/src/main/java/com/yahoo/container/jdisc/athenz/impl/AthenzService.java
index c259b01876c..dc1f8956def 100644
--- a/container-disc/src/main/java/com/yahoo/container/jdisc/athenz/impl/AthenzService.java
+++ b/container-disc/src/main/java/com/yahoo/container/jdisc/athenz/impl/AthenzService.java
@@ -1,10 +1,12 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.container.jdisc.athenz.impl;
 
+import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import org.apache.http.client.methods.CloseableHttpResponse;
 import org.apache.http.client.methods.HttpUriRequest;
 import org.apache.http.client.methods.RequestBuilder;
+import org.apache.http.conn.ssl.SSLContextBuilder;
 import org.apache.http.entity.ContentType;
 import org.apache.http.entity.StringEntity;
 import org.apache.http.impl.client.CloseableHttpClient;
@@ -12,24 +14,70 @@ import org.apache.http.impl.client.HttpClientBuilder;
 import org.apache.http.util.EntityUtils;
 import org.eclipse.jetty.http.HttpStatus;
 
+import javax.net.ssl.SSLContext;
 import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.security.KeyManagementException;
+import java.security.KeyStore;
+import java.security.KeyStoreException;
+import java.security.NoSuchAlgorithmException;
+import java.security.PrivateKey;
+import java.security.UnrecoverableKeyException;
+import java.security.cert.Certificate;
+import java.security.cert.CertificateException;
+import java.security.cert.X509Certificate;
 
 /**
  * @author mortent
+ * @author bjorncs
  */
 public class AthenzService {
 
+    private static final String INSTANCE_API_PATH = "zts/v1/instance";
+
+    private final ObjectMapper objectMapper = new ObjectMapper();
+
     /**
      * Send instance register request to ZTS, get InstanceIdentity
      */
-     public InstanceIdentity sendInstanceRegisterRequest(InstanceRegisterInformation instanceRegisterInformation, String athenzUrl) {
+     public InstanceIdentity sendInstanceRegisterRequest(InstanceRegisterInformation instanceRegisterInformation,
+                                                         String ztsEndpoint) {
         try(CloseableHttpClient client = HttpClientBuilder.create().build()) {
-            ObjectMapper objectMapper = new ObjectMapper();
             HttpUriRequest postRequest = RequestBuilder.post()
-                    .setUri(athenzUrl + "zts/v1/instance")
-                    .setEntity(new StringEntity(objectMapper.writeValueAsString(instanceRegisterInformation), ContentType.APPLICATION_JSON))
+                    .setUri(ztsEndpoint + INSTANCE_API_PATH)
+                    .setEntity(toJsonStringEntity(instanceRegisterInformation))
                     .build();
-            CloseableHttpResponse response = client.execute(postRequest);
+            return getInstanceIdentity(client, postRequest);
+        } catch (IOException e) {
+            throw new UncheckedIOException(e);
+        }
+    }
+
+    public InstanceIdentity sendInstanceRefreshRequest(String providerService,
+                                                       String instanceDomain,
+                                                       String instanceServiceName,
+                                                       String instanceId,
+                                                       InstanceRefreshInformation instanceRefreshInformation,
+                                                       String ztsEndpoint,
+                                                       X509Certificate certicate,
+                                                       PrivateKey privateKey) {
+        try (CloseableHttpClient client = createHttpClientWithTlsAuth(certicate, privateKey)) {
+            String uri = String.format("%s/%s/%s/%s/%s",
+                                       ztsEndpoint + INSTANCE_API_PATH,
+                                       providerService, instanceDomain, instanceServiceName, instanceId);
+            HttpUriRequest postRequest = RequestBuilder.post()
+                    .setUri(uri)
+                    .setEntity(toJsonStringEntity(instanceRefreshInformation))
+                    .build();
+            return getInstanceIdentity(client, postRequest);
+        } catch (IOException e) {
+            throw new UncheckedIOException(e);
+        }
+    }
+
+    private InstanceIdentity getInstanceIdentity(CloseableHttpClient client, HttpUriRequest postRequest)
+            throws IOException {
+        try (CloseableHttpResponse response = client.execute(postRequest)) {
             if(HttpStatus.isSuccess(response.getStatusLine().getStatusCode())) {
                 return objectMapper.readValue(response.getEntity().getContent(), InstanceIdentity.class);
             } else {
@@ -37,8 +85,30 @@ public class AthenzService {
                 throw new RuntimeException(String.format("Unable to get identity. http code/message: %d/%s",
                                                          response.getStatusLine().getStatusCode(), message));
             }
-        } catch (IOException e) {
+        }
+    }
+
+    private StringEntity toJsonStringEntity(Object value) throws JsonProcessingException {
+        return new StringEntity(objectMapper.writeValueAsString(value), ContentType.APPLICATION_JSON);
+    }
+
+    private static CloseableHttpClient createHttpClientWithTlsAuth(X509Certificate certificate, PrivateKey privateKey) {
+        try {
+            String dummyPassword = "athenz";
+            KeyStore keyStore = KeyStore.getInstance("JKS");
+            keyStore.load(null);
+            keyStore.setKeyEntry("athenz", privateKey, dummyPassword.toCharArray(), new Certificate[]{certificate});
+            SSLContext sslContext = new SSLContextBuilder()
+                    .loadKeyMaterial(keyStore, dummyPassword.toCharArray())
+                    .build();
+            return HttpClientBuilder.create()
+                    .setSslcontext(sslContext)
+                    .build();
+        } catch (KeyStoreException | UnrecoverableKeyException | NoSuchAlgorithmException |
+                KeyManagementException | CertificateException e) {
             throw new RuntimeException(e);
+        } catch (IOException e) {
+            throw new UncheckedIOException(e);
         }
     }
 }
diff --git a/container-disc/src/main/java/com/yahoo/container/jdisc/athenz/impl/InstanceRefreshInformation.java b/container-disc/src/main/java/com/yahoo/container/jdisc/athenz/impl/InstanceRefreshInformation.java
new file mode 100644
index 00000000000..621eafca3bb
--- /dev/null
+++ b/container-disc/src/main/java/com/yahoo/container/jdisc/athenz/impl/InstanceRefreshInformation.java
@@ -0,0 +1,24 @@
+// Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.container.jdisc.athenz.impl;
+
+import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
+import com.fasterxml.jackson.annotation.JsonInclude;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+/**
+ * @author bjorncs
+ */
+@JsonIgnoreProperties(ignoreUnknown = true)
+@JsonInclude(JsonInclude.Include.NON_NULL)
+public class InstanceRefreshInformation {
+
+    @JsonProperty("csr")
+    private final String csr;
+    @JsonProperty("token")
+    private final boolean requestServiceToken;
+
+    public InstanceRefreshInformation(String csr, boolean requestServiceToken) {
+        this.csr = csr;
+        this.requestServiceToken = requestServiceToken;
+    }
+}