feat: Autosharing of the output files #131

epam · Jan 12, 2024 · 6a705a5 · 6a705a5
1 parent 7075561
commit 6a705a5
Show file tree

Hide file tree

Showing 11 changed files with 62 additions and 10 deletions.
diff --git a/src/main/java/com/epam/aidial/core/ProxyContext.java b/src/main/java/com/epam/aidial/core/ProxyContext.java
@@ -36,6 +36,7 @@ public class ProxyContext {
     private final ApiKeyData apiKeyData;
     private final String spanId;
     private final String parentSpanId;
+    private final String sourceDeployment;
 
     private Deployment deployment;
     private String userSub;
@@ -67,10 +68,12 @@ public ProxyContext(Config config, HttpServerRequest request, ApiKeyData apiKeyD
             initExtractedClaims(apiKeyData.getExtractedClaims());
             this.traceId = apiKeyData.getTraceId();
             this.parentSpanId = apiKeyData.getSpanId();
+            this.sourceDeployment = apiKeyData.getSourceDeployment();
         } else {
             initExtractedClaims(extractedClaims);
             this.traceId = traceId;
             this.parentSpanId = null;
+            this.sourceDeployment = null;
         }
         this.spanId = spanId;
     }

diff --git a/src/main/java/com/epam/aidial/core/config/ApiKeyData.java b/src/main/java/com/epam/aidial/core/config/ApiKeyData.java
@@ -14,6 +14,7 @@ public class ApiKeyData {
     private String traceId;
     private String spanId;
     private Set<String> attachedFiles = new HashSet<>();
+    private String sourceDeployment;
 
     public ApiKeyData() {
     }

diff --git a/src/main/java/com/epam/aidial/core/controller/AccessControlBaseController.java b/src/main/java/com/epam/aidial/core/controller/AccessControlBaseController.java
@@ -17,6 +17,7 @@ public abstract class AccessControlBaseController {
 
     final Proxy proxy;
     final ProxyContext context;
+    final boolean checkFullAccess;
 
 
     /**
@@ -25,10 +26,12 @@ public abstract class AccessControlBaseController {
      */
     public Future<?> handle(String bucket, String filePath) {
         String urlDecodedBucket = UrlUtil.decodePath(bucket);
-        String expectedUserBucket = BlobStorageUtil.buildUserBucket(context);
         String decryptedBucket = proxy.getEncryptionService().decrypt(urlDecodedBucket);
+        boolean hasReadAccess = hasReadAccess(bucket, filePath);
+        boolean hasWriteAccess = hasWriteAccess(filePath, decryptedBucket);
+        boolean hasAccess = checkFullAccess ? hasWriteAccess : hasReadAccess || hasWriteAccess;
 
-        if (!expectedUserBucket.equals(decryptedBucket)) {
+        if (!hasAccess) {
             return context.respond(HttpStatus.FORBIDDEN, "You don't have an access to the bucket " + bucket);
         }
 
@@ -45,4 +48,21 @@ public Future<?> handle(String bucket, String filePath) {
 
     protected abstract Future<?> handle(ResourceDescription resource);
 
+    protected boolean hasReadAccess(String bucket, String filePath) {
+        String url = bucket + BlobStorageUtil.PATH_SEPARATOR + filePath;
+        return context.getApiKeyData().getAttachedFiles().contains(url);
+    }
+
+    protected boolean hasWriteAccess(String filePath, String decryptedBucket) {
+        String expectedUserBucket = BlobStorageUtil.buildUserBucket(context);
+        if (expectedUserBucket.equals(decryptedBucket)) {
+            return true;
+        }
+        String expectedAppDataBucket = BlobStorageUtil.buildAppDataBucket(context);
+        if (expectedAppDataBucket != null && expectedAppDataBucket.equals(decryptedBucket)) {
+            return filePath.startsWith(BlobStorageUtil.APPDATA_PATTERN.formatted(context.getSourceDeployment()));
+        }
+        return false;
+    }
+
 }
diff --git a/src/main/java/com/epam/aidial/core/controller/BucketController.java b/src/main/java/com/epam/aidial/core/controller/BucketController.java
@@ -19,7 +19,14 @@ 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));
+        String appDataBucket = BlobStorageUtil.buildAppDataBucket(context);
+        String appDataLocation;
+        if (appDataBucket == null) {
+            appDataLocation = null;
+        } else {
+            String encryptedAppDataBucket = encryptionService.encrypt(bucketLocation);
+            appDataLocation = encryptedAppDataBucket + String.format(BlobStorageUtil.APPDATA_PATTERN.formatted(context.getSourceDeployment()));
+        }
+        return context.respond(HttpStatus.OK, new Bucket(encryptedBucket, appDataLocation));
     }
 }
diff --git a/src/main/java/com/epam/aidial/core/controller/DeleteFileController.java b/src/main/java/com/epam/aidial/core/controller/DeleteFileController.java
@@ -12,7 +12,7 @@
 public class DeleteFileController extends AccessControlBaseController {
 
     public DeleteFileController(Proxy proxy, ProxyContext context) {
-        super(proxy, context);
+        super(proxy, context, true);
     }
 
     @Override

diff --git a/src/main/java/com/epam/aidial/core/controller/DownloadFileController.java b/src/main/java/com/epam/aidial/core/controller/DownloadFileController.java
@@ -20,7 +20,7 @@
 public class DownloadFileController extends AccessControlBaseController {
 
     public DownloadFileController(Proxy proxy, ProxyContext context) {
-        super(proxy, context);
+        super(proxy, context, false);
     }
 
     @Override

diff --git a/src/main/java/com/epam/aidial/core/controller/FileMetadataController.java b/src/main/java/com/epam/aidial/core/controller/FileMetadataController.java
@@ -13,7 +13,7 @@
 public class FileMetadataController extends AccessControlBaseController {
 
     public FileMetadataController(Proxy proxy, ProxyContext context) {
-        super(proxy, context);
+        super(proxy, context, false);
     }
 
     @Override

diff --git a/src/main/java/com/epam/aidial/core/controller/UploadFileController.java b/src/main/java/com/epam/aidial/core/controller/UploadFileController.java
@@ -16,7 +16,7 @@
 public class UploadFileController extends AccessControlBaseController {
 
     public UploadFileController(Proxy proxy, ProxyContext context) {
-        super(proxy, context);
+        super(proxy, context, true);
     }
 
     @Override

diff --git a/src/main/java/com/epam/aidial/core/data/Bucket.java b/src/main/java/com/epam/aidial/core/data/Bucket.java
@@ -1,4 +1,4 @@
 package com.epam.aidial.core.data;
 
-public record Bucket(String bucket) {
+public record Bucket(String bucket, String appdata) {
 }
diff --git a/src/main/java/com/epam/aidial/core/storage/BlobStorageUtil.java b/src/main/java/com/epam/aidial/core/storage/BlobStorageUtil.java
@@ -4,9 +4,13 @@
 import io.vertx.core.http.impl.MimeMapping;
 import lombok.experimental.UtilityClass;
 
+import javax.annotation.Nullable;
+
 @UtilityClass
 public class BlobStorageUtil {
 
+    public static final String APPDATA_PATTERN = "appdata/%s";
+
     private static final String USER_BUCKET_PATTERN = "Users/%s/";
 
     private static final String API_KEY_BUCKET_PATTERN = "Keys/%s/";
@@ -19,6 +23,23 @@ public String getContentType(String fileName) {
     }
 
     public String buildUserBucket(ProxyContext context) {
+        if (context.getApiKeyData().getPerRequestKey() == null) {
+            return buildInitiatorBucket(context);
+        } else {
+            return API_KEY_BUCKET_PATTERN.formatted(context.getSourceDeployment());
+        }
+    }
+
+    @Nullable
+    public String buildAppDataBucket(ProxyContext context) {
+        if (context.getApiKeyData().getPerRequestKey() == null) {
+            return null;
+        } else {
+            return buildInitiatorBucket(context);
+        }
+    }
+
+    public static String buildInitiatorBucket(ProxyContext context) {
         String userSub = context.getUserSub();
         String apiKeyId = context.getProject();
 

diff --git a/src/test/java/com/epam/aidial/core/FileApiTest.java b/src/test/java/com/epam/aidial/core/FileApiTest.java
@@ -77,7 +77,7 @@ public void testBucket(Vertx vertx, VertxTestContext context) {
                 .send(context.succeeding(response -> {
                     context.verify(() -> {
                         assertEquals(200, response.statusCode());
-                        assertEquals(new Bucket("7G9WZNcoY26Vy9D7bEgbv6zqbJGfyDp9KZyEbJR4XMZt"), response.body());
+                        assertEquals(new Bucket("7G9WZNcoY26Vy9D7bEgbv6zqbJGfyDp9KZyEbJR4XMZt", null), response.body());
                         context.completeNow();
                     });
                 }));