From 73bc81a369d1c976398c18174b2eaf5ae461607a Mon Sep 17 00:00:00 2001 From: Harald Musum Date: Mon, 18 Sep 2023 11:27:54 +0200 Subject: Implement S3Downloader --- config-proxy/pom.xml | 5 ++ .../config/proxy/filedistribution/Downloader.java | 31 ++++++++++ .../proxy/filedistribution/S3Downloader.java | 71 ++++++++++++++++++++-- .../filedistribution/UrlDownloadRpcServer.java | 20 ++---- .../proxy/filedistribution/UrlDownloader.java | 5 +- 5 files changed, 111 insertions(+), 21 deletions(-) create mode 100644 config-proxy/src/main/java/com/yahoo/vespa/config/proxy/filedistribution/Downloader.java 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 @@ -13,6 +13,11 @@ jar 8-SNAPSHOT + + com.amazonaws + aws-java-sdk-s3 + ${aws-sdk.vespa.version} + com.yahoo.vespa config-lib 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 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 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 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 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 = downloadFile(url, downloadDir); + Optional 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 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 downloadFile(String url, File downloadDir) throws IOException { + @Override + public Optional downloadFile(String url, File downloadDir) throws IOException { long start = System.currentTimeMillis(); HttpURLConnection connection = (HttpURLConnection) new URL(url).openConnection(); if (connection.getResponseCode() != 200) -- cgit v1.2.3