feat: support aws-s3 EC2 instance metadata authentication (#215)

.gitignore

@@ -3,5 +3,5 @@
 .vscode/
 build/
 bin/
-.tmp/
+data/
 *.log

README.md

@@ -50,10 +50,10 @@ Static settings are used on startup and cannot be changed while application is r
 | vertx.*                                    | -                  | Vertx settings.                                                                                                   
 | server.*                                   | -                  | Vertx HTTP server settings for incoming requests.                                                                 
 | client.*                                   | -                  | Vertx HTTP client settings for outbound requests.                                                                 
-| storage.provider                           | -                  | Specifies blob storage provider. Supported providers: s3, aws-s3, azureblob, google-cloud-storage                 
+| storage.provider                           | -                  | Specifies blob storage provider. Supported providers: s3, aws-s3, azureblob, google-cloud-storage, filesystem                 
 | storage.endpoint                           | -                  | Optional. Specifies endpoint url for s3 compatible storages                                                       
-| storage.identity                           | -                  | Blob storage access key                                                                                           
-| storage.credential                         | -                  | Blob storage secret key                                                                                           
+| storage.identity                           | -                  | Blob storage access key. Can be optional for filesystem and aws-s3 providers                                                                                           
+| storage.credential                         | -                  | Blob storage secret key. Can be optional for filesystem and aws-s3 providers                                                                                           
 | storage.bucket                             | -                  | Blob storage bucket  
 | storage.overrides.*                        | -                  | Key-value pairs to override storage settings
 | storage.createBucket                       | false              | Indicates whether bucket should be created on start-up                                                            

src/main/java/com/epam/aidial/core/config/Storage.java

@@ -17,12 +17,14 @@ public class Storage {
     @Nullable
     String endpoint;
     /**
-     * api key
+     * Api key. Optional for filesystem and aws-s3 (will try to get token from EC2 instance metadata)
      */
+    @Nullable
     String identity;
     /**
-     * secret key
+     * Secret key. Optional for filesystem and aws-s3 (will try to get token from EC2 instance metadata)
      */
+    @Nullable
     String credential;
     /**
      * container name/root bucket
@@ -35,7 +37,7 @@ public class Storage {
     boolean createBucket;
 
     /**
-     * Optional. Collection of key-value pairs for overrides, for example: "jclouds.filesystem.basedir": "/tmp/data"
+     * Optional. Collection of key-value pairs for overrides, for example: "jclouds.filesystem.basedir": "data"
      */
     @Nullable
     Properties overrides;

src/main/java/com/epam/aidial/core/storage/BlobStorage.java

@@ -43,15 +43,17 @@ public class BlobStorage implements Closeable {
     private final String bucketName;
 
     public BlobStorage(Storage config) {
-        ContextBuilder builder = ContextBuilder.newBuilder(config.getProvider());
+        String provider = config.getProvider();
+        ContextBuilder builder = ContextBuilder.newBuilder(provider);
         if (config.getEndpoint() != null) {
             builder.endpoint(config.getEndpoint());
         }
         Properties overrides = config.getOverrides();
         if (overrides != null) {
             builder.overrides(overrides);
         }
-        builder.credentials(config.getIdentity(), config.getCredential());
+        CredentialProvider credentialProvider = getCredentialProvider(StorageProvider.from(provider), config.getIdentity(), config.getCredential());
+        builder.credentialsSupplier(credentialProvider::getCredentials);
         this.storeContext = builder.buildView(BlobStoreContext.class);
         this.blobStore = storeContext.getBlobStore();
         this.bucketName = config.getBucket();
@@ -254,6 +256,20 @@ private static ContentMetadata buildContentMetadata(String contentType) {
         return BaseMutableContentMetadata.fromContentMetadata(contentMetadata);
     }
 
+    private static CredentialProvider getCredentialProvider(StorageProvider provider, String identity, String credential) {
+        return switch (provider) {
+            case S3, AZURE_BLOB, GOOGLE_CLOUD_STORAGE -> new DefaultCredentialProvider(identity, credential);
+            case FILESYSTEM -> new DefaultCredentialProvider("identity", "credential");
+            case AWS_S3 -> {
+                if (identity != null && credential != null) {
+                    yield new DefaultCredentialProvider(identity, credential);
+                } else {
+                    yield new Ec2InstanceMetadataCredentialProvider();
+                }
+            }
+        };
+    }
+
     private void createBucketIfNeeded(Storage config) {
         if (config.isCreateBucket() && !storeContext.getBlobStore().containerExists(bucketName)) {
             storeContext.getBlobStore().createContainerInLocation(null, bucketName);

src/main/java/com/epam/aidial/core/storage/CredentialProvider.java

@@ -0,0 +1,8 @@
+package com.epam.aidial.core.storage;
+
+import org.jclouds.domain.Credentials;
+
+public interface CredentialProvider {
+
+    Credentials getCredentials();
+}

src/main/java/com/epam/aidial/core/storage/DefaultCredentialProvider.java

@@ -0,0 +1,19 @@
+package com.epam.aidial.core.storage;
+
+import org.jclouds.domain.Credentials;
+
+import java.util.Objects;
+
+public class DefaultCredentialProvider implements CredentialProvider {
+
+    private final Credentials credentials;
+
+    public DefaultCredentialProvider(String identity, String credential) {
+        this.credentials = new Credentials(Objects.requireNonNull(identity), Objects.requireNonNull(credential));
+    }
+
+    @Override
+    public Credentials getCredentials() {
+        return credentials;
+    }
+}

src/main/java/com/epam/aidial/core/storage/Ec2InstanceMetadataCredentialProvider.java

@@ -0,0 +1,117 @@
+package com.epam.aidial.core.storage;
+
+import com.epam.aidial.core.util.ProxyUtil;
+import com.fasterxml.jackson.databind.PropertyNamingStrategies;
+import com.fasterxml.jackson.databind.annotation.JsonNaming;
+import lombok.AllArgsConstructor;
+import lombok.Data;
+import lombok.NoArgsConstructor;
+import lombok.extern.slf4j.Slf4j;
+import org.jclouds.aws.domain.SessionCredentials;
+
+import java.io.IOException;
+import java.net.URI;
+import java.net.http.HttpClient;
+import java.net.http.HttpRequest;
+import java.net.http.HttpResponse;
+import java.time.Duration;
+import java.time.Instant;
+import java.time.temporal.ChronoUnit;
+import java.util.Date;
+
+/**
+ * Implementation of EC2 Instance Metadata credentials provider by following
+ * see <a href="https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instance-metadata-v2-how-it-works.html">AWS spec</a>
+ */
+@Slf4j
+public class Ec2InstanceMetadataCredentialProvider implements CredentialProvider {
+
+    private static final String EC2_INSTANCE_METADATA_BASE_URL = "http://169.254.169.254/latest/";
+    private static final String EC2_INSTANCE_METADATA_CREDENTIALS_URL = EC2_INSTANCE_METADATA_BASE_URL + "meta-data/iam/security-credentials/";
+    private static final String EC2_TOKEN_TTL_HEADER_NAME = "X-aws-ec2-metadata-token-ttl-seconds";
+    private static final String EC2_METADATA_TOKEN_HEADER_NAME = "X-aws-ec2-metadata-token";
+    private static final Duration DEFAULT_REQUEST_TIMEOUT = Duration.of(10, ChronoUnit.SECONDS);
+
+    private final HttpClient httpClient;
+
+    private SessionCredentials credentials;
+
+    public Ec2InstanceMetadataCredentialProvider(HttpClient httpClient) {
+        this.httpClient = httpClient;
+    }
+
+    public Ec2InstanceMetadataCredentialProvider() {
+        this(HttpClient.newHttpClient());
+    }
+
+    @Override
+    public synchronized SessionCredentials getCredentials() {
+        try {
+            // if token present and not expired
+            if (credentials != null && credentials.getExpiration().isPresent() && Date.from(Instant.now()).after(credentials.getExpiration().get())) {
+                return credentials;
+            }
+            String token = getToken();
+            String roleName = getRoleName(token);
+            AwsCredentials awsCredentials = getAwsCredentials(token, roleName);
+
+            credentials = SessionCredentials.builder()
+                    .accessKeyId(awsCredentials.getAccessKeyId())
+                    .expiration(Date.from(Instant.parse(awsCredentials.getExpiration())))
+                    .secretAccessKey(awsCredentials.getSecretAccessKey())
+                    .sessionToken(awsCredentials.getToken()).build();
+
+            return credentials;
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private String getToken() throws IOException, InterruptedException {
+        HttpRequest request = HttpRequest.newBuilder()
+                .uri(URI.create(EC2_INSTANCE_METADATA_BASE_URL + "api/token"))
+                .setHeader(EC2_TOKEN_TTL_HEADER_NAME, "21600")
+                .timeout(DEFAULT_REQUEST_TIMEOUT)
+                .PUT(HttpRequest.BodyPublishers.noBody())
+                .build();
+
+        return httpClient.send(request, HttpResponse.BodyHandlers.ofString()).body();
+    }
+
+    private String getRoleName(String token) throws IOException, InterruptedException {
+        HttpRequest request = HttpRequest.newBuilder()
+                .uri(URI.create(EC2_INSTANCE_METADATA_CREDENTIALS_URL))
+                .setHeader(EC2_METADATA_TOKEN_HEADER_NAME, token)
+                .timeout(DEFAULT_REQUEST_TIMEOUT)
+                .GET()
+                .build();
+
+        return httpClient.send(request, HttpResponse.BodyHandlers.ofString()).body();
+    }
+
+    private AwsCredentials getAwsCredentials(String token, String roleName) throws IOException, InterruptedException {
+        HttpRequest request = HttpRequest.newBuilder()
+                .uri(URI.create(EC2_INSTANCE_METADATA_CREDENTIALS_URL + roleName))
+                .setHeader(EC2_METADATA_TOKEN_HEADER_NAME, token)
+                .timeout(DEFAULT_REQUEST_TIMEOUT)
+                .GET()
+                .build();
+
+        HttpResponse<String> response = httpClient.send(request, HttpResponse.BodyHandlers.ofString());
+        return ProxyUtil.convertToObject(response.body(), AwsCredentials.class);
+    }
+
+    @Data
+    @NoArgsConstructor
+    @AllArgsConstructor
+    @JsonNaming(PropertyNamingStrategies.UpperCamelCaseStrategy.class)
+    static class AwsCredentials {
+        String code;
+        String lastUpdated;
+        String type;
+        String accessKeyId;
+        String secretAccessKey;
+        String token;
+        String expiration;
+    }
+}

src/main/java/com/epam/aidial/core/storage/StorageProvider.java

@@ -0,0 +1,16 @@
+package com.epam.aidial.core.storage;
+
+public enum StorageProvider {
+    S3, AWS_S3, FILESYSTEM, GOOGLE_CLOUD_STORAGE, AZURE_BLOB;
+
+    public static StorageProvider from(String storageProviderName) {
+        return switch (storageProviderName) {
+            case "s3" -> S3;
+            case "aws-s3" -> AWS_S3;
+            case "azureblob" -> AZURE_BLOB;
+            case "google-cloud-storage" -> GOOGLE_CLOUD_STORAGE;
+            case "filesystem" -> FILESYSTEM;
+            default -> throw new IllegalArgumentException("Unknown storage provider");
+        };
+    }
+}

src/main/resources/aidial.settings.json

@@ -43,12 +43,10 @@
   },
   "storage": {
     "provider" : "filesystem",
-    "identity": "access-key",
-    "credential": "secret-key",
     "bucket": "dial",
     "createBucket": true,
     "overrides": {
-      "jclouds.filesystem.basedir": ".tmp/data"
+      "jclouds.filesystem.basedir": "data"
     }
   },
   "resources": {