feat: Support auto-sharing of folders #265 (#272)

Co-authored-by: Aliaksandr Stsiapanay <[email protected]>
epam · Mar 12, 2024 · 49b0b76 · 49b0b76
1 parent a599f62
commit 49b0b76
Show file tree

Hide file tree

Showing 5 changed files with 93 additions and 3 deletions.
diff --git a/src/main/java/com/epam/aidial/core/config/ApiKeyData.java b/src/main/java/com/epam/aidial/core/config/ApiKeyData.java
@@ -33,6 +33,7 @@ public class ApiKeyData {
     private String spanId;
     // list of attached file URLs collected from conversation history of the current request
     private Set<String> attachedFiles = new HashSet<>();
+    private List<String> attachedFolders = new ArrayList<>();
     // deployment name of the source(application/assistant/model) associated with the current request
     private String sourceDeployment;
     // Execution path of the root request

diff --git a/src/main/java/com/epam/aidial/core/controller/AccessControlBaseController.java b/src/main/java/com/epam/aidial/core/controller/AccessControlBaseController.java
@@ -49,8 +49,7 @@ public Future<?> handle(String resourceType, String bucket, String path) {
 
                     if (!checkFullAccess) {
                         // some per-request API-keys may have access to the resources implicitly
-                        boolean isAutoShared = context.getApiKeyData().getAttachedFiles().contains(resource.getUrl());
-                        if (isAutoShared) {
+                        if (proxy.getAccessService().isAutoSharedResource(resource, context)) {
                             return true;
                         }
 

diff --git a/src/main/java/com/epam/aidial/core/controller/DeploymentPostController.java b/src/main/java/com/epam/aidial/core/controller/DeploymentPostController.java
@@ -261,7 +261,11 @@ private void processAttachedFile(String url) {
         if (accessService.hasWriteAccess(resource, context)
                 || accessService.isSharedResource(resource, context)
                 || sourceApiKeyData.getAttachedFiles().contains(resourceUrl)) {
-            destApiKeyData.getAttachedFiles().add(resourceUrl);
+            if (resource.isFolder()) {
+                destApiKeyData.getAttachedFolders().add(resourceUrl);
+            } else {
+                destApiKeyData.getAttachedFiles().add(resourceUrl);
+            }
         } else {
             throw new HttpException(HttpStatus.FORBIDDEN, "Access denied to the file %s".formatted(url));
         }

diff --git a/src/main/java/com/epam/aidial/core/security/AccessService.java b/src/main/java/com/epam/aidial/core/security/AccessService.java
@@ -8,6 +8,8 @@
 import com.epam.aidial.core.util.UrlUtil;
 import lombok.AllArgsConstructor;
 
+import java.util.List;
+
 @AllArgsConstructor
 public class AccessService {
 
@@ -49,4 +51,19 @@ public boolean isReviewResource(ResourceDescription resource, ProxyContext conte
         String actualUserBucket = encryptionService.encrypt(actualUserLocation);
         return publicationService != null && publicationService.hasReviewAccess(resource, actualUserBucket, actualUserLocation);
     }
+
+    public boolean isAutoSharedResource(ResourceDescription resource, ProxyContext context) {
+        String resourceUrl = resource.getUrl();
+        boolean isAutoShared = context.getApiKeyData().getAttachedFiles().contains(resourceUrl);
+        if (isAutoShared) {
+            return true;
+        }
+        List<String> attachedFolders = context.getApiKeyData().getAttachedFolders();
+        for (String folder : attachedFolders) {
+            if (resourceUrl.startsWith(folder)) {
+                return true;
+            }
+        }
+        return false;
+    }
 }
diff --git a/src/test/java/com/epam/aidial/core/FileApiTest.java b/src/test/java/com/epam/aidial/core/FileApiTest.java
@@ -498,6 +498,75 @@ public void testDownloadSharedFile(Vertx vertx, VertxTestContext context) {
         });
     }
 
+    @Test
+    public void testDownloadFileWithinSharedFolder(Vertx vertx, VertxTestContext context) {
+        Checkpoint checkpoint = context.checkpoint(3);
+        WebClient client = WebClient.create(vertx);
+
+        // creating per-request API key with proxyKey2 as originator
+        // and proxyKey1 caller
+        ApiKeyData projectApiKeyData = apiKeyStore.getApiKeyData("proxyKey2");
+        ApiKeyData apiKeyData1 = new ApiKeyData();
+        apiKeyData1.setOriginalKey(projectApiKeyData.getOriginalKey());
+        // set deployment ID for proxyKey1
+        apiKeyData1.setSourceDeployment("EPM-RTC-GPT");
+        apiKeyData1.setAttachedFolders(List.of("files/7G9WZNcoY26Vy9D7bEgbv6zqbJGfyDp9KZyEbJR4XMZt/folder1/"));
+        apiKeyStore.assignApiKey(apiKeyData1);
+
+        String apiKey1 = apiKeyData1.getPerRequestKey();
+
+        FileMetadata expectedFileMetadata = new FileMetadata("7G9WZNcoY26Vy9D7bEgbv6zqbJGfyDp9KZyEbJR4XMZt",
+                "file.txt", "folder1", "files/7G9WZNcoY26Vy9D7bEgbv6zqbJGfyDp9KZyEbJR4XMZt/folder1/file.txt", 17, "text/plain");
+
+        Future.succeededFuture().compose((mapper) -> {
+            Promise<Void> promise = Promise.promise();
+            // proxyKey2 uploads file
+            client.put(serverPort, "localhost", "/v1/files/7G9WZNcoY26Vy9D7bEgbv6zqbJGfyDp9KZyEbJR4XMZt/folder1/file.txt")
+                    .putHeader("Api-key", "proxyKey2")
+                    .as(BodyCodec.json(FileMetadata.class))
+                    .sendMultipartForm(generateMultipartForm("file.txt", TEST_FILE_CONTENT),
+                            context.succeeding(response -> {
+                                context.verify(() -> {
+                                    assertEquals(200, response.statusCode());
+                                    assertEquals(expectedFileMetadata, response.body());
+                                    checkpoint.flag();
+                                    promise.complete();
+                                });
+                            })
+                    );
+
+            return promise.future();
+        }).compose((mapper) -> {
+            Promise<Void> promise = Promise.promise();
+            // verify caller can't download shared file with own api-key
+            client.get(serverPort, "localhost", "/v1/files/7G9WZNcoY26Vy9D7bEgbv6zqbJGfyDp9KZyEbJR4XMZt/folder1/file.txt")
+                    .putHeader("Api-key", "proxyKey1")
+                    .as(BodyCodec.string())
+                    .send(context.succeeding(response -> {
+                        context.verify(() -> {
+                            assertEquals(403, response.statusCode());
+                            assertEquals("You don't have an access to the FILE 7G9WZNcoY26Vy9D7bEgbv6zqbJGfyDp9KZyEbJR4XMZt/folder1/file.txt", response.body());
+                            checkpoint.flag();
+                            promise.complete();
+                        });
+                    }));
+
+            return promise.future();
+        }).andThen((result) -> {
+            // verify pre-request api key can download shared file
+            client.get(serverPort, "localhost", "/v1/files/7G9WZNcoY26Vy9D7bEgbv6zqbJGfyDp9KZyEbJR4XMZt/folder1/file.txt")
+                    .putHeader("Api-key", apiKey1)
+                    .as(BodyCodec.string())
+                    .send(context.succeeding(response -> {
+                        context.verify(() -> {
+                            assertEquals(200, response.statusCode());
+                            assertEquals(TEST_FILE_CONTENT, response.body());
+                            checkpoint.flag();
+                        });
+                    }));
+        });
+    }
+
     @Test
     public void testFileDownload(Vertx vertx, VertxTestContext context) {
         Checkpoint checkpoint = context.checkpoint(2);