Implement S3Downloader

5 files changed, 111 insertions, 21 deletions

diff --git a/config-proxy/pom.xml b/config-proxy/pom.xml
index 2984b4b1b65..466515388d4 100644
--- a/config-proxy/pom.xml
+++ b/config-proxy/pom.xml
@@ -14,6 +14,11 @@
   <version>8-SNAPSHOT</version>
   <dependencies>
     <dependency>
+      <groupId>com.amazonaws</groupId>
+      <artifactId>aws-java-sdk-s3</artifactId>
+      <version>${aws-sdk.vespa.version}</version>
+    </dependency>
+    <dependency>
       <groupId>com.yahoo.vespa</groupId>
       <artifactId>config-lib</artifactId>
       <version>${project.version}</version>
diff --git a/config-proxy/src/main/java/com/yahoo/vespa/config/proxy/filedistribution/Downloader.java b/config-proxy/src/main/java/com/yahoo/vespa/config/proxy/filedistribution/Downloader.java
new file mode 100644
index 00000000000..0692d3ee499
--- /dev/null
+++ b/config-proxy/src/main/java/com/yahoo/vespa/config/proxy/filedistribution/Downloader.java
@@ -0,0 +1,31 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.config.proxy.filedistribution;
+
+import java.io.File;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.net.HttpURLConnection;
+import java.net.URL;
+import java.nio.channels.Channels;
+import java.nio.channels.ReadableByteChannel;
+import java.nio.file.Files;
+import java.util.Optional;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+/**
+ *
+ * @author hmusum
+ */
+interface Downloader {
+
+    Optional<File> downloadFile(String url, File downloadDir) throws IOException;
+
+    default String fileName() { return "contents"; }
+
+    default boolean alreadyDownloaded(Downloader downloader, File downloadDir) {
+        File contents = new File(downloadDir, downloader.fileName());
+        return contents.exists() && contents.length() > 0;
+    }
+
+}
diff --git a/config-proxy/src/main/java/com/yahoo/vespa/config/proxy/filedistribution/S3Downloader.java b/config-proxy/src/main/java/com/yahoo/vespa/config/proxy/filedistribution/S3Downloader.java
index 76a6c29a56a..a999601ada5 100644
--- a/config-proxy/src/main/java/com/yahoo/vespa/config/proxy/filedistribution/S3Downloader.java
+++ b/config-proxy/src/main/java/com/yahoo/vespa/config/proxy/filedistribution/S3Downloader.java
@@ -1,16 +1,77 @@
 package com.yahoo.vespa.config.proxy.filedistribution;
 
+import com.amazonaws.auth.AWSCredentials;
+import com.amazonaws.auth.AWSCredentialsProvider;
+import com.amazonaws.auth.AWSSessionCredentials;
+import com.amazonaws.auth.BasicSessionCredentials;
+import com.amazonaws.services.s3.AmazonS3;
+import com.amazonaws.services.s3.AmazonS3ClientBuilder;
+import com.amazonaws.services.s3.AmazonS3URI;
+import com.amazonaws.services.s3.model.S3Object;
+import com.yahoo.slime.Cursor;
+import com.yahoo.slime.Slime;
+import com.yahoo.slime.SlimeUtils;
+import com.yahoo.vespa.defaults.Defaults;
+
 import java.io.File;
 import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
 import java.util.Optional;
+import java.util.concurrent.atomic.AtomicReference;
+
+public class S3Downloader implements Downloader {
+
+    private final AmazonS3 s3Client;
+
+    S3Downloader() {
+        this.s3Client = AmazonS3ClientBuilder.standard()
+                .withRegion(System.getenv("VESPA_CLOUD_NATIVE_REGION"))
+                .withCredentials(new CredentialsProvider())
+                .build();
+    }
+
+    @Override
+    public Optional<File> downloadFile(String url, File targetDir) throws IOException {
+        AmazonS3URI s3URI = new AmazonS3URI(url);
+        S3Object s3Object = s3Client.getObject(s3URI.getBucket(), s3URI.getKey());
+        File file = new File(targetDir, fileName());
+        Files.copy(s3Object.getObjectContent(), file.toPath());
+        return Optional.of(file);
+    }
+
+    private static class CredentialsProvider implements AWSCredentialsProvider {
+
+        private static final String DEFAULT_CREDENTIALS_PATH = Defaults.getDefaults()
+                .underVespaHome("var/vespa/aws/credentials.json");
+
+        private final AtomicReference<AWSCredentials> credentials = new AtomicReference<>();
+        private final Path credentialsPath;
+
+        public CredentialsProvider() {
+            this.credentialsPath = Path.of(DEFAULT_CREDENTIALS_PATH);
+        }
+
+        @Override
+        public AWSCredentials getCredentials() { return credentials.getAndSet(readCredentials()); }
 
-public class S3Downloader {
+        @Override
+        public void refresh() { readCredentials(); }
 
-    // TODO: Avoid hardcoding
-    private static final String ZTS_URL = "https://zts.athenz.ouroath.com:4443/zts/v1";
+        private AWSSessionCredentials readCredentials() {
+            try {
+                Slime slime = SlimeUtils.jsonToSlime(Files.readAllBytes(credentialsPath));
+                Cursor cursor = slime.get();
+                String accessKey = cursor.field("awsAccessKey").asString();
+                String secretKey = cursor.field("awsSecretKey").asString();
+                String sessionToken = cursor.field("sessionToken").asString();
+                return new BasicSessionCredentials(accessKey, secretKey, sessionToken);
+            } catch (IOException e) {
+                throw new UncheckedIOException(e);
+            }
+        }
 
-    Optional<File> downloadFile(String fileName, File targetDir) throws IOException {
-        throw new UnsupportedOperationException("Download of S3 urls not implemented");
     }
 
 }
diff --git a/config-proxy/src/main/java/com/yahoo/vespa/config/proxy/filedistribution/UrlDownloadRpcServer.java b/config-proxy/src/main/java/com/yahoo/vespa/config/proxy/filedistribution/UrlDownloadRpcServer.java
index e48ad3c8759..59481df1aa8 100644
--- a/config-proxy/src/main/java/com/yahoo/vespa/config/proxy/filedistribution/UrlDownloadRpcServer.java
+++ b/config-proxy/src/main/java/com/yahoo/vespa/config/proxy/filedistribution/UrlDownloadRpcServer.java
@@ -13,7 +13,6 @@ import com.yahoo.yolean.Exceptions;
 import net.jpountz.xxhash.XXHashFactory;
 
 import java.io.File;
-import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.Optional;
 import java.util.concurrent.ExecutorService;
@@ -36,7 +35,6 @@ import static java.util.logging.Level.WARNING;
 class UrlDownloadRpcServer {
 
     private static final Logger log = Logger.getLogger(UrlDownloadRpcServer.class.getName());
-    private static final String CONTENTS_FILE_NAME = "contents";
     static final File defaultDownloadDirectory = new File(Defaults.getDefaults().underVespaHome("var/db/vespa/download"));
 
     private final File rootDownloadDir;
@@ -70,15 +68,16 @@ class UrlDownloadRpcServer {
     private void downloadFile(Request req) {
         String url = req.parameters().get(0).asString();
         File downloadDir = new File(rootDownloadDir, urlToDirName(url));
-        if (alreadyDownloaded(downloadDir)) {
+        Downloader downloader = downloader(url);
+        if (downloader.alreadyDownloaded(downloader, downloadDir)) {
             log.log(Level.INFO, "URL '" + url + "' already downloaded");
-            req.returnValues().add(new StringValue(new File(downloadDir, CONTENTS_FILE_NAME).getAbsolutePath()));
+            req.returnValues().add(new StringValue(new File(downloadDir, downloader.fileName()).getAbsolutePath()));
             req.returnRequest();
             return;
         }
 
         try {
-            Optional<File> file = downloadFile(url, downloadDir);
+            Optional<File> file = downloader.downloadFile(url, downloadDir);
             if (file.isPresent())
                 req.returnValues().add(new StringValue(file.get().getAbsolutePath()));
             else
@@ -91,10 +90,8 @@ class UrlDownloadRpcServer {
         req.returnRequest();
     }
 
-    private static Optional<File> downloadFile(String url, File downloadDir) throws IOException {
-        return (url.startsWith("s3://"))
-                ? new S3Downloader().downloadFile(url, downloadDir)
-                : new UrlDownloader().downloadFile(url, downloadDir);
+    private static Downloader downloader(String url) {
+        return url.startsWith("s3://") ? new S3Downloader() : new UrlDownloader();
     }
 
     private static void logAndSetRpcError(Request req, String url, Throwable e, int rpcErrorCode) {
@@ -107,9 +104,4 @@ class UrlDownloadRpcServer {
         return String.valueOf(XXHashFactory.fastestJavaInstance().hash64().hash(ByteBuffer.wrap(Utf8.toBytes(uri)), 0));
     }
 
-    private static boolean alreadyDownloaded(File downloadDir) {
-        File contents = new File(downloadDir, CONTENTS_FILE_NAME);
-        return contents.exists() && contents.length() > 0;
-    }
-
 }
diff --git a/config-proxy/src/main/java/com/yahoo/vespa/config/proxy/filedistribution/UrlDownloader.java b/config-proxy/src/main/java/com/yahoo/vespa/config/proxy/filedistribution/UrlDownloader.java
index 5fa2a12f608..3a50545f0ae 100644
--- a/config-proxy/src/main/java/com/yahoo/vespa/config/proxy/filedistribution/UrlDownloader.java
+++ b/config-proxy/src/main/java/com/yahoo/vespa/config/proxy/filedistribution/UrlDownloader.java
@@ -18,12 +18,13 @@ import java.util.logging.Logger;
  *
  * @author hmusum
  */
-class UrlDownloader {
+class UrlDownloader implements Downloader {
 
     private static final Logger log = Logger.getLogger(UrlDownloader.class.getName());
     private static final String CONTENTS_FILE_NAME = "contents";
 
-    Optional<File> downloadFile(String url, File downloadDir) throws IOException {
+    @Override
+    public Optional<File> downloadFile(String url, File downloadDir) throws IOException {
         long start = System.currentTimeMillis();
         HttpURLConnection connection = (HttpURLConnection) new URL(url).openConnection();
         if (connection.getResponseCode() != 200)