From a69b0605653bc89f963b9512d9fedde5a59cb5fb Mon Sep 17 00:00:00 2001 From: Aliaksandr Stsiapanay Date: Mon, 11 Mar 2024 19:21:28 +0300 Subject: [PATCH 1/2] feat: Support auto-sharing of folders #265 --- .../com/epam/aidial/core/config/ApiKeyData.java | 1 + .../controller/AccessControlBaseController.java | 3 +-- .../controller/DeploymentPostController.java | 6 +++++- .../aidial/core/security/AccessService.java | 17 +++++++++++++++++ 4 files changed, 24 insertions(+), 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 index 3e0e27401..0dfa5fa66 100644 --- 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 attachedFiles = new HashSet<>(); + private List 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 index 17776607c..3cd3989ed 100644 --- 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 index 3cd95586f..ca2ba5fff 100644 --- 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 index 3989a73c1..62b55419b 100644 --- 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 attachedFolders = context.getApiKeyData().getAttachedFolders(); + for (String folder : attachedFolders) { + if (resourceUrl.startsWith(folder)) { + return true; + } + } + return false; + } } From 507f71a57cf69d80cfd99207d2965ba766790a39 Mon Sep 17 00:00:00 2001 From: Aliaksandr Stsiapanay Date: Tue, 12 Mar 2024 11:15:06 +0300 Subject: [PATCH 2/2] chore: add tests --- .../com/epam/aidial/core/FileApiTest.java | 69 +++++++++++++++++++ 1 file changed, 69 insertions(+) diff --git a/src/test/java/com/epam/aidial/core/FileApiTest.java b/src/test/java/com/epam/aidial/core/FileApiTest.java index a95bd0954..f185ac1df 100644 --- 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 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 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);