feat: implementation of basic access control for Files API (#91) (#90)

Co-authored-by: Maksim_Hadalau <[email protected]>
epam · Dec 18, 2023 · 11f209a · 11f209a
1 parent b061ecb
commit 11f209a
Show file tree

Hide file tree

Showing 27 changed files with 948 additions and 364 deletions.
diff --git a/README.md b/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:

diff --git a/src/main/java/com/epam/aidial/core/AiDial.java b/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);

diff --git a/src/main/java/com/epam/aidial/core/Proxy.java b/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;
 
 import static com.epam.aidial.core.security.AccessTokenValidator.extractTokenFromHeader;
 
@@ -48,6 +50,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;
@@ -56,6 +60,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) {
@@ -81,7 +86,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;
         }

diff --git a/src/main/java/com/epam/aidial/core/config/Encryption.java b/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;
+}
diff --git a/src/main/java/com/epam/aidial/core/controller/AccessControlBaseController.java b/src/main/java/com/epam/aidial/core/controller/AccessControlBaseController.java
@@ -0,0 +1,39 @@
+package com.epam.aidial.core.controller;
+
+import com.epam.aidial.core.Proxy;
+import com.epam.aidial.core.ProxyContext;
+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;
+
+@AllArgsConstructor
+public abstract class AccessControlBaseController {
+
+    final Proxy proxy;
+    final ProxyContext context;
+
+
+    public Future<?> handle(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;
+        try {
+            resource = ResourceDescription.from(ResourceType.FILE, bucket, decryptedBucket, filePath);
+        } catch (Exception ex) {
+            return context.respond(HttpStatus.BAD_REQUEST, "Invalid file url provided");
+        }
+
+        return handle(resource);
+    }
+
+    protected abstract Future<?> handle(ResourceDescription resource);
+
+}
diff --git a/src/main/java/com/epam/aidial/core/controller/BucketController.java b/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));
+    }
+}
diff --git a/src/main/java/com/epam/aidial/core/controller/ControllerSelector.java b/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/([a-zA-Z0-9]+)/(.*)");
+
+    private static final Pattern PATTERN_FILES_METADATA = Pattern.compile("/v1/files/metadata/([a-zA-Z0-9]+)/(.*)");
 
     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.handle(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.handle(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.handle(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.handle(bucket, filePath);
         }
 
         return null;

diff --git a/src/main/java/com/epam/aidial/core/controller/DeleteFileController.java b/src/main/java/com/epam/aidial/core/controller/DeleteFileController.java
@@ -3,26 +3,26 @@
 import com.epam.aidial.core.Proxy;
 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.util.HttpStatus;
 import io.vertx.core.Future;
-import lombok.AllArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 
 @Slf4j
-@AllArgsConstructor
-public class DeleteFileController {
-    private final Proxy proxy;
-    private final ProxyContext context;
+public class DeleteFileController extends AccessControlBaseController {
+
+    public DeleteFileController(Proxy proxy, ProxyContext context) {
+        super(proxy, context);
+    }
+
+    @Override
+    protected Future<?> handle(ResourceDescription resource) {
+        if (resource.isFolder()) {
+            return context.respond(HttpStatus.BAD_REQUEST, "Can't delete a folder");
+        }
+
+        String absoluteFilePath = resource.getAbsoluteFilePath();
 
-    /**
-     * 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
-     */
-    public Future<?> delete(String path) {
-        String absoluteFilePath = BlobStorageUtil.buildAbsoluteFilePath(context, path);
         BlobStorage storage = proxy.getStorage();
         Future<Void> result = proxy.getVertx().executeBlocking(() -> {
             try {