feat: implementation of basic access control for Files API

README.md

@@ -49,7 +49,8 @@ Static settings are used on startup and cannot be changed while application is r
 | storage.credential                      | -                  |Blob storage secret key
 | storage.bucket                          | -                  |Blob storage bucket
 | storage.createBucket                    | false              |Indicates whether bucket should be created on start-up
-
+| encryption.password                     | -                  |Password used for AES encryption
+| encryption.salt                         | -                  |Salt used for AES encryption
 ### Dynamic settings
 Dynamic settings are stored in JSON files, specified via "config.files" static setting, and reloaded at interval, specified via "config.reload" static setting.
 Dynamic settings include:

src/main/java/com/epam/aidial/core/AiDial.java

@@ -1,13 +1,14 @@
 package com.epam.aidial.core;
 
-import com.auth0.jwk.UrlJwkProvider;
 import com.epam.aidial.core.config.ConfigStore;
+import com.epam.aidial.core.config.Encryption;
 import com.epam.aidial.core.config.FileConfigStore;
 import com.epam.aidial.core.config.Storage;
 import com.epam.aidial.core.limiter.RateLimiter;
 import com.epam.aidial.core.log.GfLogStore;
 import com.epam.aidial.core.log.LogStore;
 import com.epam.aidial.core.security.AccessTokenValidator;
+import com.epam.aidial.core.security.EncryptionService;
 import com.epam.aidial.core.storage.BlobStorage;
 import com.epam.aidial.core.upstream.UpstreamBalancer;
 import com.epam.deltix.gflog.core.LogConfigurator;
@@ -68,7 +69,8 @@ void start() throws Exception {
                 Storage storageConfig = Json.decodeValue(settings("storage").toBuffer(), Storage.class);
                 storage = new BlobStorage(storageConfig);
             }
-            Proxy proxy = new Proxy(vertx, client, configStore, logStore, rateLimiter, upstreamBalancer, accessTokenValidator, storage);
+            EncryptionService encryptionService = new EncryptionService(Json.decodeValue(settings("encryption").toBuffer(), Encryption.class));
+            Proxy proxy = new Proxy(vertx, client, configStore, logStore, rateLimiter, upstreamBalancer, accessTokenValidator, storage, encryptionService);
 
             server = vertx.createHttpServer(new HttpServerOptions(settings("server"))).requestHandler(proxy);
             open(server, HttpServer::listen);

src/main/java/com/epam/aidial/core/Proxy.java

@@ -8,6 +8,7 @@
 import com.epam.aidial.core.limiter.RateLimiter;
 import com.epam.aidial.core.log.LogStore;
 import com.epam.aidial.core.security.AccessTokenValidator;
+import com.epam.aidial.core.security.EncryptionService;
 import com.epam.aidial.core.security.ExtractedClaims;
 import com.epam.aidial.core.storage.BlobStorage;
 import com.epam.aidial.core.upstream.UpstreamBalancer;
@@ -27,6 +28,7 @@
 
 import java.net.URLDecoder;
 import java.nio.charset.StandardCharsets;
+import java.util.Set;
 
 @Slf4j
 @Getter
@@ -46,6 +48,8 @@ public class Proxy implements Handler<HttpServerRequest> {
     public static final int REQUEST_BODY_MAX_SIZE_BYTES = 16 * 1024 * 1024;
     public static final int FILES_REQUEST_BODY_MAX_SIZE_BYTES = 512 * 1024 * 1024;
 
+    private static final Set<HttpMethod> ALLOWED_HTTP_METHODS = Set.of(HttpMethod.GET, HttpMethod.POST, HttpMethod.PUT, HttpMethod.DELETE);
+
     private final Vertx vertx;
     private final HttpClient client;
     private final ConfigStore configStore;
@@ -54,6 +58,7 @@ public class Proxy implements Handler<HttpServerRequest> {
     private final UpstreamBalancer upstreamBalancer;
     private final AccessTokenValidator tokenValidator;
     private final BlobStorage storage;
+    private final EncryptionService encryptionService;
 
     @Override
     public void handle(HttpServerRequest request) {
@@ -79,7 +84,7 @@ private void handleRequest(HttpServerRequest request) {
         }
 
         HttpMethod requestMethod = request.method();
-        if (requestMethod != HttpMethod.GET && requestMethod != HttpMethod.POST && requestMethod != HttpMethod.DELETE) {
+        if (!ALLOWED_HTTP_METHODS.contains(requestMethod)) {
             respond(request, HttpStatus.METHOD_NOT_ALLOWED);
             return;
         }

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

@@ -0,0 +1,9 @@
+package com.epam.aidial.core.config;
+
+import lombok.Data;
+
+@Data
+public class Encryption {
+    String password;
+    String salt;
+}

src/main/java/com/epam/aidial/core/controller/BucketController.java

@@ -0,0 +1,25 @@
+package com.epam.aidial.core.controller;
+
+import com.epam.aidial.core.Proxy;
+import com.epam.aidial.core.ProxyContext;
+import com.epam.aidial.core.data.Bucket;
+import com.epam.aidial.core.security.EncryptionService;
+import com.epam.aidial.core.storage.BlobStorageUtil;
+import com.epam.aidial.core.util.HttpStatus;
+import io.vertx.core.Future;
+import lombok.AllArgsConstructor;
+
+@AllArgsConstructor
+public class BucketController {
+
+    private final Proxy proxy;
+    private final ProxyContext context;
+
+    public Future<?> getBucket() {
+        EncryptionService encryptionService = proxy.getEncryptionService();
+        String bucketLocation = BlobStorageUtil.buildUserBucket(context);
+        String encryptedBucket = encryptionService.encrypt(bucketLocation);
+
+        return context.respond(HttpStatus.OK, new Bucket(encryptedBucket));
+    }
+}

src/main/java/com/epam/aidial/core/controller/ControllerSelector.java

@@ -33,7 +33,11 @@ public class ControllerSelector {
     private static final Pattern PATTERN_APPLICATIONS = Pattern.compile("/+openai/applications");
 
 
-    private static final Pattern PATTERN_FILES = Pattern.compile("/v1/files(.*)");
+    private static final Pattern PATTERN_BUCKET = Pattern.compile("/v1/bucket");
+
+    private static final Pattern PATTERN_FILES = Pattern.compile("/v1/files/([^/]*)/(.*)");
+
+    private static final Pattern PATTERN_FILES_METADATA = Pattern.compile("/v1/files/metadata/([^/]*)/(.*)");
 
     private static final Pattern PATTERN_RATE_RESPONSE = Pattern.compile("/+v1/([-.@a-zA-Z0-9]+)/rate");
     private static final Pattern PATTERN_TOKENIZE = Pattern.compile("/+v1/deployments/([-.@a-zA-Z0-9]+)/tokenize");
@@ -50,6 +54,8 @@ public Controller select(Proxy proxy, ProxyContext context) {
             controller = selectPost(proxy, context, path);
         } else if (method == HttpMethod.DELETE) {
             controller = selectDelete(proxy, context, path);
+        } else if (method == HttpMethod.PUT) {
+            controller = selectPut(proxy, context, path);
         }
 
         return (controller == null) ? new RouteController(proxy, context) : controller;
@@ -123,17 +129,26 @@ private static Controller selectGet(Proxy proxy, ProxyContext context, String pa
             return controller::getApplications;
         }
 
+        match = match(PATTERN_FILES_METADATA, path);
+        if (match != null) {
+            String bucket = match.group(1);
+            String filePath = match.group(2);
+            FileMetadataController controller = new FileMetadataController(proxy, context);
+            return () -> controller.list(bucket, filePath);
+        }
+
         match = match(PATTERN_FILES, path);
         if (match != null) {
-            String filePath = match.group(1);
-            String purpose = context.getRequest().params().get(DownloadFileController.PURPOSE_FILE_QUERY_PARAMETER);
-            if (DownloadFileController.QUERY_METADATA_QUERY_PARAMETER_VALUE.equals(purpose)) {
-                FileMetadataController controller = new FileMetadataController(proxy, context);
-                return () -> controller.list(filePath);
-            } else {
-                DownloadFileController controller = new DownloadFileController(proxy, context);
-                return () -> controller.download(filePath);
-            }
+            String bucket = match.group(1);
+            String filePath = match.group(2);
+            DownloadFileController controller = new DownloadFileController(proxy, context);
+            return () -> controller.download(bucket, filePath);
+        }
+
+        match = match(PATTERN_BUCKET, path);
+        if (match != null) {
+            BucketController controller = new BucketController(proxy, context);
+            return controller::getBucket;
         }
 
         return null;
@@ -148,22 +163,15 @@ private static Controller selectPost(Proxy proxy, ProxyContext context, String p
             return () -> controller.handle(deploymentId, deploymentApi);
         }
 
-        match = match(PATTERN_FILES, path);
-        if (match != null) {
-            String relativeFilePath = match.group(1);
-            UploadFileController controller = new UploadFileController(proxy, context);
-            return () -> controller.upload(relativeFilePath);
-        }
-
         match = match(PATTERN_RATE_RESPONSE, path);
         if (match != null) {
             String deploymentId = match.group(1);
 
             Function<Deployment, String> getter = (model) -> {
                 return Optional.ofNullable(model)
-                    .map(d -> d.getFeatures())
-                    .map(t -> t.getRateEndpoint())
-                    .orElse(null);
+                        .map(d -> d.getFeatures())
+                        .map(t -> t.getRateEndpoint())
+                        .orElse(null);
             };
 
             DeploymentFeatureController controller = new DeploymentFeatureController(proxy, context);
@@ -176,9 +184,9 @@ private static Controller selectPost(Proxy proxy, ProxyContext context, String p
 
             Function<Deployment, String> getter = (model) -> {
                 return Optional.ofNullable(model)
-                    .map(d -> d.getFeatures())
-                    .map(t -> t.getTokenizeEndpoint())
-                    .orElse(null);
+                        .map(d -> d.getFeatures())
+                        .map(t -> t.getTokenizeEndpoint())
+                        .orElse(null);
             };
 
             DeploymentFeatureController controller = new DeploymentFeatureController(proxy, context);
@@ -191,9 +199,9 @@ private static Controller selectPost(Proxy proxy, ProxyContext context, String p
 
             Function<Deployment, String> getter = (model) -> {
                 return Optional.ofNullable(model)
-                    .map(d -> d.getFeatures())
-                    .map(t -> t.getTruncatePromptEndpoint())
-                    .orElse(null);
+                        .map(d -> d.getFeatures())
+                        .map(t -> t.getTruncatePromptEndpoint())
+                        .orElse(null);
             };
 
             DeploymentFeatureController controller = new DeploymentFeatureController(proxy, context);
@@ -206,9 +214,22 @@ private static Controller selectPost(Proxy proxy, ProxyContext context, String p
     private static Controller selectDelete(Proxy proxy, ProxyContext context, String path) {
         Matcher match = match(PATTERN_FILES, path);
         if (match != null) {
-            String relativeFilePath = match.group(1);
+            String bucket = match.group(1);
+            String filePath = match.group(2);
             DeleteFileController controller = new DeleteFileController(proxy, context);
-            return () -> controller.delete(relativeFilePath);
+            return () -> controller.delete(bucket, filePath);
+        }
+
+        return null;
+    }
+
+    private static Controller selectPut(Proxy proxy, ProxyContext context, String path) {
+        Matcher match = match(PATTERN_FILES, path);
+        if (match != null) {
+            String bucket = match.group(1);
+            String filePath = match.group(2);
+            UploadFileController controller = new UploadFileController(proxy, context);
+            return () -> controller.upload(bucket, filePath);
         }
 
         return null;

src/main/java/com/epam/aidial/core/controller/DeleteFileController.java

@@ -4,6 +4,8 @@
 import com.epam.aidial.core.ProxyContext;
 import com.epam.aidial.core.storage.BlobStorage;
 import com.epam.aidial.core.storage.BlobStorageUtil;
+import com.epam.aidial.core.storage.ResourceDescription;
+import com.epam.aidial.core.storage.ResourceType;
 import com.epam.aidial.core.util.HttpStatus;
 import io.vertx.core.Future;
 import lombok.AllArgsConstructor;
@@ -17,12 +19,20 @@ public class DeleteFileController {
 
     /**
      * Deletes file from storage.
-     * Current API implementation requires a relative path, absolute path will be calculated based on authentication context
      *
-     * @param path relative path, for example: /inputs/data.csv
+     * @param filePath relative path, for example: inputs/data.csv
      */
-    public Future<?> delete(String path) {
-        String absoluteFilePath = BlobStorageUtil.buildAbsoluteFilePath(context, path);
+    public Future<?> delete(String bucket, String filePath) {
+        String expectedUserBucket = BlobStorageUtil.buildUserBucket(context);
+        String decryptedBucket = proxy.getEncryptionService().decrypt(bucket);
+
+        if (!expectedUserBucket.equals(decryptedBucket)) {
+            return context.respond(HttpStatus.FORBIDDEN, "You don't have an access to the bucket " + bucket);
+        }
+
+        ResourceDescription resource = ResourceDescription.from(ResourceType.FILES, bucket, decryptedBucket, filePath);
+        String absoluteFilePath = resource.getAbsoluteFilePath();
+
         BlobStorage storage = proxy.getStorage();
         Future<Void> result = proxy.getVertx().executeBlocking(() -> {
             try {

src/main/java/com/epam/aidial/core/controller/DownloadFileController.java

@@ -4,6 +4,8 @@
 import com.epam.aidial.core.ProxyContext;
 import com.epam.aidial.core.storage.BlobStorageUtil;
 import com.epam.aidial.core.storage.InputStreamReader;
+import com.epam.aidial.core.storage.ResourceDescription;
+import com.epam.aidial.core.storage.ResourceType;
 import com.epam.aidial.core.util.HttpStatus;
 import io.vertx.core.Future;
 import io.vertx.core.Promise;
@@ -21,13 +23,6 @@
 @AllArgsConstructor
 public class DownloadFileController {
 
-    private static final String PATH_TYPE_QUERY_PARAMETER = "path";
-    private static final String RELATIVE_PATH_TYPE = "relative";
-
-    static final String PURPOSE_FILE_QUERY_PARAMETER = "purpose";
-
-    static final String QUERY_METADATA_QUERY_PARAMETER_VALUE = "metadata";
-
     private final Proxy proxy;
     private final ProxyContext context;
 
@@ -36,18 +31,20 @@ public class DownloadFileController {
      * Path can be either absolute or relative.
      * Path type determined by "path" query parameter which can be "absolute" or "relative"(default value)
      *
-     * @param path file path; absolute or relative
+     * @param filePath file path; absolute or relative
      */
-    public Future<?> download(String path) {
-        String pathType = context.getRequest().params().get(PATH_TYPE_QUERY_PARAMETER);
-        String absoluteFilePath;
-        if (RELATIVE_PATH_TYPE.equals(pathType)) {
-            absoluteFilePath = BlobStorageUtil.buildAbsoluteFilePath(context, path);
-        } else {
-            absoluteFilePath = path;
+    public Future<?> download(String bucket, String filePath) {
+        String expectedUserBucket = BlobStorageUtil.buildUserBucket(context);
+        String decryptedBucket = proxy.getEncryptionService().decrypt(bucket);
+
+        if (!expectedUserBucket.equals(decryptedBucket)) {
+            return context.respond(HttpStatus.FORBIDDEN, "You don't have an access to the bucket " + bucket);
         }
+
+        ResourceDescription resource = ResourceDescription.from(ResourceType.FILES, bucket, decryptedBucket, filePath);
+
         Future<Blob> blobFuture = proxy.getVertx().executeBlocking(() ->
-                proxy.getStorage().load(BlobStorageUtil.removeLeadingPathSeparator(absoluteFilePath)));
+                proxy.getStorage().load(resource.getAbsoluteFilePath()));
 
         Promise<Void> result = Promise.promise();
         blobFuture.onSuccess(blob -> {
@@ -78,7 +75,7 @@ public Future<?> download(String path) {
                 result.fail(e);
             }
         }).onFailure(error -> context.respond(HttpStatus.INTERNAL_SERVER_ERROR,
-                "Failed to fetch file with path " + path));
+                "Failed to fetch file with path %s/%s".formatted(bucket, filePath)));
 
         return result.future();
     }

src/main/java/com/epam/aidial/core/controller/FileMetadataController.java

@@ -5,13 +5,13 @@
 import com.epam.aidial.core.data.FileMetadataBase;
 import com.epam.aidial.core.storage.BlobStorage;
 import com.epam.aidial.core.storage.BlobStorageUtil;
+import com.epam.aidial.core.storage.ResourceDescription;
+import com.epam.aidial.core.storage.ResourceType;
 import com.epam.aidial.core.util.HttpStatus;
 import io.vertx.core.Future;
 import lombok.AllArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 
-import java.util.List;
-
 @AllArgsConstructor
 @Slf4j
 public class FileMetadataController {
@@ -20,17 +20,28 @@ public class FileMetadataController {
 
     /**
      * Lists all files and folders that belong to the provided path.
-     * Current API implementation requires a relative path, absolute path will be calculated based on authentication context
      *
-     * @param path relative path, for example: /inputs
+     * @param path relative path, for example: inputs/
      */
-    public Future<?> list(String path) {
+    public Future<?> list(String bucket, String path) {
+        String expectedUserBucket = BlobStorageUtil.buildUserBucket(context);
+        String decryptedBucket = proxy.getEncryptionService().decrypt(bucket);
+
+        if (!expectedUserBucket.equals(decryptedBucket)) {
+            return context.respond(HttpStatus.FORBIDDEN, "You don't have an access to the bucket " + bucket);
+        }
+
         BlobStorage storage = proxy.getStorage();
         return proxy.getVertx().executeBlocking(() -> {
             try {
-                String absolutePath = BlobStorageUtil.buildAbsoluteFilePath(context, path);
-                List<FileMetadataBase> metadata = storage.listMetadata(absolutePath);
-                context.respond(HttpStatus.OK, metadata);
+                String filePath = path.isEmpty() ? BlobStorageUtil.PATH_SEPARATOR : path;
+                ResourceDescription resource = ResourceDescription.from(ResourceType.FILES, bucket, decryptedBucket, filePath);
+                FileMetadataBase metadata = storage.listMetadata(resource);
+                if (metadata != null) {
+                    context.respond(HttpStatus.OK, metadata);
+                } else {
+                    context.respond(HttpStatus.NOT_FOUND);
+                }
             } catch (Exception ex) {
                 log.error("Failed to list files", ex);
                 context.respond(HttpStatus.INTERNAL_SERVER_ERROR, "Failed to list files by path %s".formatted(path));