test: drive the resource API under the tenant-rooted layout #1863

README.md

@@ -232,6 +232,8 @@ without putting the raw address in the log.
 | storage.bucket              |     -      |    No    | Blob storage bucket.                                                                                                                                                                               |
 | storage.overrides.*         |     -      |    No    | Key-value pairs to override storage settings. `*` might be any specific blob storage setting to be overridden. Refer to [examples](#temporary-credentials-1) in the sections below.                |
 | storage.createBucket        |   false    |    No    | Indicates whether bucket should be created on start-up.                                                                                                                                            |
+| storageLayout.tenantRooted  |   false    |    No    | Places every bucket under a tenant root instead of storing it at the top level. Changing it on a populated deployment re-addresses existing data and requires migration; leave disabled otherwise.  |
+| storageLayout.defaultTenant |  default   |    No    | Tenant that existing buckets are placed under when `storageLayout.tenantRooted` is enabled.                                                                                                         |
 | storage.prefix              |     -      |    No    | Base prefix for all stored resources. The purpose to use the same bucket for different environments, e.g. dev, prod, pre-prod. Must not contain path separators or any invalid chars.              |
 | storage.maxUploadedFileSize | 536870912  |    No    | Maximum size in bytes of uploaded file. If a size of uploaded file exceeds the limit the server returns HTTP code 413                                                                              |
 

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

@@ -87,8 +87,12 @@
 import com.epam.aidial.core.storage.blobstore.BlobStorage;
 import com.epam.aidial.core.storage.blobstore.Storage;
 import com.epam.aidial.core.storage.cache.CacheClientFactory;
+import com.epam.aidial.core.storage.resource.LegacyStorageLayout;
 import com.epam.aidial.core.storage.resource.ResourceDescriptor;
 import com.epam.aidial.core.storage.resource.ResourceTypes;
+import com.epam.aidial.core.storage.resource.StorageLayout;
+import com.epam.aidial.core.storage.resource.StorageLayouts;
+import com.epam.aidial.core.storage.resource.TenantRootedStorageLayout;
 import com.epam.aidial.core.storage.service.LockService;
 import com.epam.aidial.core.storage.service.ResourceService;
 import com.epam.aidial.core.storage.service.TimerService;
@@ -202,6 +206,8 @@ void start() throws Exception {
                 accessTokenValidator = new AccessTokenValidator(settings("identityProviders"), vertx, taskExecutor, client, clientOptions, claimsLogLevel);
             }
 
+            StorageLayouts.useLayout(createStorageLayout(settings("storageLayout")));
+
             if (storage == null) {
                 Storage storageConfig = Json.decodeValue(settings("storage").toBuffer(), Storage.class);
                 storage = new BlobStorage(storageConfig);
@@ -489,6 +495,14 @@ private JsonObject settings(String key) {
         return settings.getJsonObject(key, new JsonObject());
     }
 
+    private static StorageLayout createStorageLayout(JsonObject settings) {
+        if (!settings.getBoolean("tenantRooted", false)) {
+            return LegacyStorageLayout.INSTANCE;
+        }
+
+        return new TenantRootedStorageLayout(settings.getString("defaultTenant", "default"));
+    }
+
     private List<String> getAllowedRedirectUris() {
         return settings("toolsets")
                 .getJsonObject("security", new JsonObject())

server/src/main/java/com/epam/aidial/core/server/config/SecretFieldProcessor.java

@@ -3,6 +3,7 @@
 import com.epam.aidial.core.config.annotation.EncryptedField;
 import com.epam.aidial.core.credentials.data.credentials.BucketInfo;
 import com.epam.aidial.core.credentials.encryption.CredentialEncryptionService;
+import com.epam.aidial.core.server.security.ResourceSecretAad;
 import com.epam.aidial.core.storage.resource.ResourceDescriptor;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.node.ArrayNode;
@@ -39,15 +40,15 @@ public void encryptFields(Object entity, ResourceDescriptor descriptor) {
         if (entity == null) {
             return;
         }
-        byte[] aad = descriptor.getAbsoluteFilePath().getBytes(StandardCharsets.UTF_8);
+        byte[] aad = ResourceSecretAad.deriveFor(descriptor);
         walk(entity, aad, true);
     }
 
     public void decryptFields(Object entity, ResourceDescriptor descriptor) {
         if (entity == null) {
             return;
         }
-        byte[] aad = descriptor.getAbsoluteFilePath().getBytes(StandardCharsets.UTF_8);
+        byte[] aad = ResourceSecretAad.deriveFor(descriptor);
         walk(entity, aad, false);
     }
 
@@ -56,7 +57,7 @@ public String resolveSecret(String value, ResourceDescriptor descriptor) {
             return null;
         }
         if (value.startsWith(ENC_PREFIX) && value.endsWith(ENC_SUFFIX)) {
-            byte[] aad = descriptor.getAbsoluteFilePath().getBytes(StandardCharsets.UTF_8);
+            byte[] aad = ResourceSecretAad.deriveFor(descriptor);
             return decryptEnvelope(value, aad, "value");
         }
         return value;

server/src/main/java/com/epam/aidial/core/server/security/ResourceSecretAad.java

@@ -0,0 +1,27 @@
+package com.epam.aidial.core.server.security;
+
+import com.epam.aidial.core.storage.resource.ResourceDescriptor;
+
+import java.nio.charset.StandardCharsets;
+
+/**
+ * Derives the additional authenticated data that binds an encrypted field to the resource holding it.
+ *
+ * <p>The AAD is the resource's physical path, so ciphertext is only readable at the path it was written
+ * to: re-addressing a resource without re-encrypting it makes its secrets unrecoverable. Anything that
+ * moves encrypted resources has to decrypt with the source path and encrypt with the destination one,
+ * which is why the path overload exists alongside the descriptor one.
+ */
+public final class ResourceSecretAad {
+
+    private ResourceSecretAad() {
+    }
+
+    public static byte[] deriveFor(ResourceDescriptor descriptor) {
+        return deriveFor(descriptor.getAbsoluteFilePath());
+    }
+
+    public static byte[] deriveFor(String absoluteFilePath) {
+        return absoluteFilePath.getBytes(StandardCharsets.UTF_8);
+    }
+}

server/src/main/java/com/epam/aidial/core/server/service/BackgroundJobService.java

@@ -15,6 +15,7 @@
 import com.epam.aidial.core.server.log.AnalyticsLogContext;
 import com.epam.aidial.core.server.log.LogStore;
 import com.epam.aidial.core.server.security.ApiKeyStore;
+import com.epam.aidial.core.server.security.ResourceSecretAad;
 import com.epam.aidial.core.server.token.TokenStatsTracker;
 import com.epam.aidial.core.server.token.TokenUsage;
 import com.epam.aidial.core.server.token.UsagePerModel;
@@ -270,14 +271,14 @@ private Future<Void> processResult(
 
     private String encryptKey(ResourceDescriptor descriptor, String key) {
         BucketInfo bucketInfo = new BucketInfo(descriptor.getBucketName(), descriptor.getBucketLocation());
-        byte[] aad = descriptor.getAbsoluteFilePath().getBytes(StandardCharsets.UTF_8);
+        byte[] aad = ResourceSecretAad.deriveFor(descriptor);
         byte[] cipher = encryptionService.encrypt(bucketInfo, key.getBytes(StandardCharsets.UTF_8), aad);
         return Base64.getEncoder().encodeToString(cipher);
     }
 
     private String decryptKey(ResourceDescriptor descriptor, String key) {
         BucketInfo bucketInfo = new BucketInfo(descriptor.getBucketName(), descriptor.getBucketLocation());
-        byte[] aad = descriptor.getAbsoluteFilePath().getBytes(StandardCharsets.UTF_8);
+        byte[] aad = ResourceSecretAad.deriveFor(descriptor);
         byte[] raw = Base64.getDecoder().decode(key);
         return new String(encryptionService.decrypt(bucketInfo, raw, aad), StandardCharsets.UTF_8);
     }

server/src/main/resources/aidial.settings.json

@@ -55,6 +55,10 @@
       "jclouds.filesystem.basedir": "data"
     }
   },
+  "storageLayout": {
+    "tenantRooted": false,
+    "defaultTenant": "default"
+  },
   "resources": {
     "maxSize" : 67108864,
     "maxSizeToCache": 1048576,

server/src/test/java/com/epam/aidial/core/server/TenantRootedLayoutApiTest.java

@@ -0,0 +1,86 @@
+package com.epam.aidial.core.server;
+
+import com.epam.aidial.core.storage.resource.LegacyStorageLayout;
+import com.epam.aidial.core.storage.resource.StorageLayouts;
+import io.vertx.core.http.HttpMethod;
+import io.vertx.core.json.JsonObject;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.List;
+import java.util.stream.Stream;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+/**
+ * Drives the resource API with {@code storageLayout.tenantRooted} enabled: the whole stack — descriptor,
+ * cache and blob store — has to agree on the tenant-rooted paths, which unit tests cannot show.
+ */
+public class TenantRootedLayoutApiTest extends ResourceBaseTest {
+
+    private static final String TENANT = "test-tenant";
+
+    @Override
+    protected JsonObject additionalSettingsOverrides() {
+        return new JsonObject().put("storageLayout", new JsonObject()
+                .put("tenantRooted", true)
+                .put("defaultTenant", TENANT));
+    }
+
+    @AfterEach
+    public void restoreLegacyLayout() {
+        StorageLayouts.useLayout(LegacyStorageLayout.INSTANCE);
+    }
+
+    @Test
+    public void testResourceRoundTrip() {
+        Response created = resourceRequest(HttpMethod.PUT, "/folder/conversation", CONVERSATION_BODY_1);
+        assertEquals(200, created.status());
+
+        Response read = resourceRequest(HttpMethod.GET, "/folder/conversation");
+        assertEquals(200, read.status());
+        assertEquals(CONVERSATION_BODY_1, read.body());
+    }
+
+    @Test
+    public void testResourceListingAndDeletion() {
+        assertEquals(200, resourceRequest(HttpMethod.PUT, "/folder/conversation", CONVERSATION_BODY_1).status());
+
+        Response listing = metadata("/folder/");
+        assertEquals(200, listing.status());
+        assertTrue(listing.body().contains("conversations/" + bucket + "/folder/conversation"),
+                () -> "Unexpected listing: " + listing.body());
+
+        assertEquals(200, resourceRequest(HttpMethod.DELETE, "/folder/conversation").status());
+        assertEquals(404, resourceRequest(HttpMethod.GET, "/folder/conversation").status());
+    }
+
+    @Test
+    public void testBlobIsStoredUnderTenantRoot() throws IOException {
+        assertEquals(200, resourceRequest(HttpMethod.PUT, "/folder/conversation", CONVERSATION_BODY_1).status());
+        // the resource is written back to the blob store asynchronously
+        Response flushed = resourceRequest(HttpMethod.GET, "/folder/conversation");
+        assertEquals(200, flushed.status());
+
+        List<Path> storedPaths = findStoredPaths("");
+        List<Path> tenantRootedPaths = storedPaths.stream()
+                .filter(path -> path.toString().contains(".org/" + TENANT))
+                .toList();
+        assertTrue(!tenantRootedPaths.isEmpty(),
+                () -> "No blob stored under the tenant root, found: " + storedPaths);
+        assertTrue(tenantRootedPaths.stream().anyMatch(path -> path.toString().contains(".conversations")),
+                () -> "Conversations are not stored in a reserved type folder: " + tenantRootedPaths);
+    }
+
+    private List<Path> findStoredPaths(String marker) throws IOException {
+        try (Stream<Path> paths = Files.walk(testDir)) {
+            return paths.filter(Files::isRegularFile)
+                    .filter(path -> path.toString().contains(marker))
+                    .toList();
+        }
+    }
+}

server/src/test/java/com/epam/aidial/core/server/security/ResourceSecretAadTest.java

@@ -0,0 +1,37 @@
+package com.epam.aidial.core.server.security;
+
+import com.epam.aidial.core.storage.resource.ResourceDescriptor;
+import com.epam.aidial.core.storage.resource.ResourceTypes;
+import org.junit.jupiter.api.Test;
+
+import java.nio.charset.StandardCharsets;
+import java.util.List;
+
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+
+public class ResourceSecretAadTest {
+
+    private static final ResourceDescriptor RESOURCE = new ResourceDescriptor(ResourceTypes.APPLICATION, "app",
+            List.of("catalog"), "bucket", "Users/u1/", false);
+
+    @Test
+    public void testAadIsResourcePath() {
+        assertArrayEquals("Users/u1/applications/catalog/app".getBytes(StandardCharsets.UTF_8),
+                ResourceSecretAad.deriveFor(RESOURCE));
+    }
+
+    @Test
+    public void testDescriptorAndPathAgree() {
+        assertArrayEquals(ResourceSecretAad.deriveFor(RESOURCE.getAbsoluteFilePath()),
+                ResourceSecretAad.deriveFor(RESOURCE));
+    }
+
+    @Test
+    public void testDifferentPathsProduceDifferentAad() {
+        byte[] legacy = ResourceSecretAad.deriveFor("Users/u1/applications/catalog/app");
+        byte[] tenantRooted = ResourceSecretAad.deriveFor(".org/default/.users/u1/.applications/catalog/app");
+
+        assertFalse(java.util.Arrays.equals(legacy, tenantRooted));
+    }
+}

storage/src/main/java/com/epam/aidial/core/storage/resource/ResourceDescriptor.java

@@ -23,10 +23,6 @@ public class ResourceDescriptor {
     public static final String PLATFORM_BUCKET = "platform";
     public static final String PLATFORM_LOCATION = PLATFORM_BUCKET + PATH_SEPARATOR;
 
-    /**
-     * Layout in force for physical paths. Swapped once the tenant-rooted layout is enabled.
-     */
-    private static final StorageLayout LAYOUT = LegacyStorageLayout.INSTANCE;
 
     ResourceType type;
     /**
@@ -134,7 +130,8 @@ public String getAbsoluteFilePath() {
      * location followed by the resource-type folder.
      */
     private String getStoragePrefix() {
-        return LAYOUT.resolveLocationPrefix(bucketLocation) + LAYOUT.resolveTypeFolder(type.group()) + PATH_SEPARATOR;
+        StorageLayout layout = StorageLayouts.resolveActive();
+        return layout.resolveLocationPrefix(bucketLocation) + layout.resolveTypeFolder(type.group()) + PATH_SEPARATOR;
     }
 
     /**

storage/src/main/java/com/epam/aidial/core/storage/resource/StorageLayouts.java

@@ -0,0 +1,22 @@
+package com.epam.aidial.core.storage.resource;
+
+/**
+ * Holds the layout every physical path is composed with. Set once during start-up, before any resource
+ * is read or written; process-wide because {@link ResourceDescriptor} is constructed everywhere and
+ * carries no configuration of its own.
+ */
+public final class StorageLayouts {
+
+    private static volatile StorageLayout active = LegacyStorageLayout.INSTANCE;
+
+    private StorageLayouts() {
+    }
+
+    public static StorageLayout resolveActive() {
+        return active;
+    }
+
+    public static void useLayout(StorageLayout layout) {
+        active = layout;
+    }
+}

storage/src/main/java/com/epam/aidial/core/storage/resource/TenantRootedStorageLayout.java

@@ -0,0 +1,28 @@
+package com.epam.aidial.core.storage.resource;
+
+/**
+ * The tenant-rooted layout: every bucket location is placed under its tenant, and resource-type folders
+ * are reserved names. Conversion rules live in {@link TenantLayoutTransform}.
+ */
+public final class TenantRootedStorageLayout implements StorageLayout {
+
+    private final String tenantId;
+
+    public TenantRootedStorageLayout(String tenantId) {
+        if (tenantId == null || tenantId.isBlank()) {
+            throw new IllegalArgumentException("Tenant id must not be blank");
+        }
+
+        this.tenantId = tenantId;
+    }
+
+    @Override
+    public String resolveLocationPrefix(String bucketLocation) {
+        return TenantLayoutTransform.toTenantLocation(bucketLocation, tenantId);
+    }
+
+    @Override
+    public String resolveTypeFolder(String group) {
+        return TenantLayoutTransform.toTenantTypeFolder(group);
+    }
+}

storage/src/test/java/com/epam/aidial/core/storage/resource/StorageLayoutsTest.java

@@ -0,0 +1,42 @@
+package com.epam.aidial.core.storage.resource;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertSame;
+
+public class StorageLayoutsTest {
+
+    @AfterEach
+    public void restoreDefaultLayout() {
+        StorageLayouts.useLayout(LegacyStorageLayout.INSTANCE);
+    }
+
+    @Test
+    public void testLegacyLayoutIsActiveByDefault() {
+        assertSame(LegacyStorageLayout.INSTANCE, StorageLayouts.resolveActive());
+    }
+
+    @Test
+    public void testActiveLayoutIsReplaceable() {
+        StorageLayout tenantRooted = new TenantRootedStorageLayout("acme");
+        StorageLayouts.useLayout(tenantRooted);
+
+        assertSame(tenantRooted, StorageLayouts.resolveActive());
+    }
+
+    @Test
+    public void testDescriptorPathFollowsActiveLayout() {
+        ResourceDescriptor file = new ResourceDescriptor(ResourceTypes.FILE, "notes.txt",
+                List.of("documents"), "bucket", "Users/u1/", false);
+
+        assertEquals("Users/u1/files/documents/notes.txt", file.getAbsoluteFilePath());
+
+        StorageLayouts.useLayout(new TenantRootedStorageLayout("acme"));
+
+        assertEquals(".org/acme/.users/u1/.files/documents/notes.txt", file.getAbsoluteFilePath());
+    }
+}