fix: unintentional sharing; clean up invitations and revoke shared access

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

@@ -69,8 +69,6 @@ public class AiDial {
 
     private BlobStorage storage;
     private ResourceService resourceService;
-    private InvitationService invitationService;
-    private ShareService shareService;
 
     @VisibleForTesting
     void start() throws Exception {
@@ -98,14 +96,10 @@ void start() throws Exception {
 
             redis = openRedis();
 
-            if (redis != null) {
-                LockService lockService = new LockService(redis);
-                resourceService = new ResourceService(vertx, redis, storage, lockService, settings("resources"), storage.getPrefix());
-                invitationService = new InvitationService(resourceService, encryptionService, settings("invitations"));
-                shareService = new ShareService(resourceService, invitationService, encryptionService);
-            } else {
-                log.warn("Redis config is not found, some features may be unavailable");
-            }
+            LockService lockService = new LockService(redis);
+            resourceService = new ResourceService(vertx, redis, storage, lockService, settings("resources"), storage.getPrefix());
+            InvitationService invitationService = new InvitationService(resourceService, encryptionService, settings("invitations"));
+            ShareService shareService = new ShareService(resourceService, invitationService, encryptionService);
 
             AccessService accessService = new AccessService(encryptionService, shareService);
 
@@ -114,7 +108,7 @@ void start() throws Exception {
             proxy = new Proxy(vertx, client, configStore, logStore,
                     rateLimiter, upstreamBalancer, accessTokenValidator,
                     storage, encryptionService, apiKeyStore, tokenStatsTracker, resourceService, invitationService,
-                    shareService, accessService);
+                    shareService, accessService, lockService);
 
             server = vertx.createHttpServer(new HttpServerOptions(settings("server"))).requestHandler(proxy);
             open(server, HttpServer::listen);
@@ -130,7 +124,7 @@ void start() throws Exception {
     private RedissonClient openRedis() throws IOException {
         JsonObject conf = settings("redis");
         if (conf.isEmpty()) {
-            return null;
+            throw new IllegalArgumentException("Redis configuration not found");
         }
 
         ConfigSupport support = new ConfigSupport();

src/main/java/com/epam/aidial/core/Proxy.java

@@ -13,6 +13,7 @@
 import com.epam.aidial.core.security.EncryptionService;
 import com.epam.aidial.core.security.ExtractedClaims;
 import com.epam.aidial.core.service.InvitationService;
+import com.epam.aidial.core.service.LockService;
 import com.epam.aidial.core.service.ResourceService;
 import com.epam.aidial.core.service.ShareService;
 import com.epam.aidial.core.storage.BlobStorage;
@@ -76,6 +77,7 @@ public class Proxy implements Handler<HttpServerRequest> {
     private final InvitationService invitationService;
     private final ShareService shareService;
     private final AccessService accessService;
+    private final LockService lockService;
 
     @Override
     public void handle(HttpServerRequest request) {

src/main/java/com/epam/aidial/core/controller/DeleteFileController.java

@@ -4,6 +4,8 @@
 import com.epam.aidial.core.ProxyContext;
 import com.epam.aidial.core.data.ResourceLink;
 import com.epam.aidial.core.data.ResourceLinkCollection;
+import com.epam.aidial.core.service.InvitationService;
+import com.epam.aidial.core.service.LockService;
 import com.epam.aidial.core.service.ShareService;
 import com.epam.aidial.core.storage.BlobStorage;
 import com.epam.aidial.core.storage.ResourceDescription;
@@ -18,10 +20,14 @@
 public class DeleteFileController extends AccessControlBaseController {
 
     private final ShareService shareService;
+    private final InvitationService invitationService;
+    private final LockService lockService;
 
     public DeleteFileController(Proxy proxy, ProxyContext context) {
         super(proxy, context, true);
         this.shareService = proxy.getShareService();
+        this.invitationService = proxy.getInvitationService();
+        this.lockService = proxy.getLockService();
     }
 
     @Override
@@ -34,18 +40,19 @@ protected Future<?> handle(ResourceDescription resource) {
 
         BlobStorage storage = proxy.getStorage();
         Future<Void> result = proxy.getVertx().executeBlocking(() -> {
+            String bucketName = resource.getBucketName();
+            String bucketLocation = resource.getBucketLocation();
             try {
-                // clean shared access
-                // TODO remove check when redis become mandatory
-                if (shareService != null) {
-                    Set<ResourceLink> resourceLinks = new HashSet<>();
-                    resourceLinks.add(new ResourceLink(resource.getUrl()));
-                    shareService.revokeSharedAccess(resource.getBucketName(), resource.getBucketLocation(), new ResourceLinkCollection(resourceLinks));
-                }
-                storage.delete(absoluteFilePath);
-                return null;
+                Set<ResourceLink> resourceLinks = new HashSet<>();
+                resourceLinks.add(new ResourceLink(resource.getUrl()));
+                return lockService.underBucketLock(proxy, bucketLocation, () -> {
+                    invitationService.cleanUpResourceLinks(bucketName, bucketLocation, resourceLinks);
+                    shareService.revokeSharedAccess(bucketName, bucketLocation, new ResourceLinkCollection(resourceLinks));
+                    storage.delete(absoluteFilePath);
+                    return null;
+                });
             } catch (Exception ex) {
-                log.error("Failed to delete file  %s/%s".formatted(resource.getBucketName(), resource.getOriginalPath()), ex);
+                log.error("Failed to delete file  %s/%s".formatted(bucketName, resource.getOriginalPath()), ex);
                 throw new RuntimeException(ex);
             }
         });

src/main/java/com/epam/aidial/core/controller/InvitationController.java

@@ -4,10 +4,12 @@
 import com.epam.aidial.core.ProxyContext;
 import com.epam.aidial.core.security.EncryptionService;
 import com.epam.aidial.core.service.InvitationService;
+import com.epam.aidial.core.service.LockService;
 import com.epam.aidial.core.service.PermissionDeniedException;
 import com.epam.aidial.core.service.ResourceNotFoundException;
 import com.epam.aidial.core.service.ShareService;
 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;
 
@@ -18,13 +20,15 @@ public class InvitationController {
     private final InvitationService invitationService;
     private final ShareService shareService;
     private final EncryptionService encryptionService;
+    private final LockService lockService;
 
     public InvitationController(Proxy proxy, ProxyContext context) {
         this.proxy = proxy;
         this.context = context;
         this.invitationService = proxy.getInvitationService();
         this.shareService = proxy.getShareService();
         this.encryptionService = proxy.getEncryptionService();
+        this.lockService = proxy.getLockService();
     }
 
     public Future<?> getInvitations() {
@@ -40,14 +44,20 @@ public Future<?> getInvitations() {
     }
 
     public Future<?> getOrAcceptInvitation(String invitationId) {
-        String accept = context.getRequest().getParam("accept");
-        if (accept != null) {
+        boolean accept = Boolean.parseBoolean(context.getRequest().getParam("accept"));
+        if (accept) {
             proxy.getVertx()
                     .executeBlocking(() -> {
                         String bucketLocation = BlobStorageUtil.buildInitiatorBucket(context);
                         String bucket = encryptionService.encrypt(bucketLocation);
-                        shareService.acceptSharedResources(bucket, bucketLocation, invitationId);
-                        return null;
+                        ResourceDescription invitationResource = invitationService.getInvitationResource(invitationId);
+                        if (invitationResource == null) {
+                            throw new ResourceNotFoundException();
+                        }
+                        return lockService.underBucketLock(proxy, invitationResource.getBucketLocation(), () -> {
+                            shareService.acceptSharedResources(bucket, bucketLocation, invitationId);
+                            return null;
+                        });
                     })
                     .onSuccess(ignore -> context.respond(HttpStatus.OK))
                     .onFailure(error -> {
@@ -78,8 +88,10 @@ public Future<?> deleteInvitation(String invitationId) {
                 .executeBlocking(() -> {
                     String bucketLocation = BlobStorageUtil.buildInitiatorBucket(context);
                     String bucket = encryptionService.encrypt(bucketLocation);
-                    invitationService.deleteInvitation(bucket, bucketLocation, invitationId);
-                    return null;
+                    return lockService.underBucketLock(proxy, bucketLocation, () -> {
+                        invitationService.deleteInvitation(bucket, bucketLocation, invitationId);
+                        return null;
+                    });
                 })
                 .onSuccess(ignore -> context.respond(HttpStatus.OK))
                 .onFailure(error -> {

src/main/java/com/epam/aidial/core/controller/ResourceController.java

@@ -4,6 +4,8 @@
 import com.epam.aidial.core.ProxyContext;
 import com.epam.aidial.core.data.ResourceLink;
 import com.epam.aidial.core.data.ResourceLinkCollection;
+import com.epam.aidial.core.service.InvitationService;
+import com.epam.aidial.core.service.LockService;
 import com.epam.aidial.core.service.ResourceService;
 import com.epam.aidial.core.service.ShareService;
 import com.epam.aidial.core.storage.ResourceDescription;
@@ -27,6 +29,8 @@ public class ResourceController extends AccessControlBaseController {
     private final Vertx vertx;
     private final ResourceService service;
     private final ShareService shareService;
+    private final LockService lockService;
+    private final InvitationService invitationService;
     private final boolean metadata;
 
     public ResourceController(Proxy proxy, ProxyContext context, boolean metadata) {
@@ -35,6 +39,8 @@ public ResourceController(Proxy proxy, ProxyContext context, boolean metadata) {
         this.vertx = proxy.getVertx();
         this.service = proxy.getResourceService();
         this.shareService = proxy.getShareService();
+        this.lockService = proxy.getLockService();
+        this.invitationService = proxy.getInvitationService();
         this.metadata = metadata;
     }
 
@@ -158,9 +164,14 @@ private Future<?> deleteResource(ResourceDescription descriptor) {
         return vertx.executeBlocking(() -> {
                     Set<ResourceLink> resourceLinks = new HashSet<>();
                     resourceLinks.add(new ResourceLink(descriptor.getUrl()));
-                    shareService.revokeSharedAccess(descriptor.getBucketName(), descriptor.getBucketLocation(),
-                            new ResourceLinkCollection(resourceLinks));
-                    return service.deleteResource(descriptor);
+                    String bucketName = descriptor.getBucketName();
+                    String bucketLocation = descriptor.getBucketLocation();
+                    return lockService.underBucketLock(proxy, bucketLocation, () -> {
+                        invitationService.cleanUpResourceLinks(bucketName, bucketLocation, resourceLinks);
+                        shareService.revokeSharedAccess(bucketName, bucketLocation,
+                                new ResourceLinkCollection(resourceLinks));
+                        return service.deleteResource(descriptor);
+                    });
                 })
                 .onSuccess(deleted -> {
                     if (deleted) {

src/main/java/com/epam/aidial/core/controller/ShareController.java

@@ -6,6 +6,8 @@
 import com.epam.aidial.core.data.ResourceLinkCollection;
 import com.epam.aidial.core.data.ShareResourcesRequest;
 import com.epam.aidial.core.security.EncryptionService;
+import com.epam.aidial.core.service.InvitationService;
+import com.epam.aidial.core.service.LockService;
 import com.epam.aidial.core.service.ShareService;
 import com.epam.aidial.core.storage.BlobStorageUtil;
 import com.epam.aidial.core.util.HttpException;
@@ -26,15 +28,18 @@ public class ShareController {
     private final ProxyContext context;
     private final ShareService shareService;
     private final EncryptionService encryptionService;
+    private final LockService lockService;
+    private final InvitationService invitationService;
 
     public ShareController(Proxy proxy, ProxyContext context) {
         this.proxy = proxy;
         this.context = context;
         this.shareService = proxy.getShareService();
         this.encryptionService = proxy.getEncryptionService();
+        this.lockService = proxy.getLockService();
+        this.invitationService = proxy.getInvitationService();
     }
 
-
     public Future<?> handle(Operation operation) {
         switch (operation) {
             case LIST -> listSharedResources();
@@ -122,10 +127,11 @@ public Future<?> revokeSharedResources() {
                     String bucketLocation = BlobStorageUtil.buildInitiatorBucket(context);
                     String bucket = encryptionService.encrypt(bucketLocation);
                     return proxy.getVertx()
-                            .executeBlocking(() -> {
+                            .executeBlocking(() -> lockService.underBucketLock(proxy, bucketLocation, () -> {
+                                invitationService.cleanUpResourceLinks(bucket, bucketLocation, request.getResources());
                                 shareService.revokeSharedAccess(bucket, bucketLocation, request);
                                 return null;
-                            });
+                            }));
                 })
                 .onSuccess(response -> context.respond(HttpStatus.OK))
                 .onFailure(this::handleServiceError);

src/main/java/com/epam/aidial/core/service/InvitationService.java

@@ -15,6 +15,7 @@
 
 import java.time.Instant;
 import java.time.temporal.ChronoUnit;
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -125,11 +126,35 @@ public InvitationCollection getMyInvitations(String bucket, String location) {
         return new InvitationCollection(new HashSet<>(invitationMap.getInvitations().values()));
     }
 
+    public void cleanUpResourceLinks(String bucket, String location, Set<ResourceLink> resourcesToCleanUp) {
+        ResourceDescription resource = ResourceDescription.fromDecoded(ResourceType.INVITATION, bucket, location, INVITATION_RESOURCE_FILENAME);
+        resourceService.computeResource(resource, state -> {
+            InvitationsMap invitations = ProxyUtil.convertToObject(state, InvitationsMap.class);
+            if (invitations == null) {
+                return null;
+            }
+            Map<String, Invitation> invitationMap = invitations.getInvitations();
+            List<Invitation> invitationsToRemove = new ArrayList<>();
+            for (Invitation invitation : invitationMap.values()) {
+                Set<ResourceLink> invitationResourceLinks = invitation.getResources();
+                invitationResourceLinks.removeAll(resourcesToCleanUp);
+
+                if (invitationResourceLinks.isEmpty()) {
+                    invitationsToRemove.add(invitation);
+                }
+            }
+
+            invitationsToRemove.forEach(invitationToRemove -> invitationMap.remove(invitationToRemove.getId()));
+
+            return ProxyUtil.convertToString(invitations);
+        });
+    }
+
     private void cleanUpExpiredInvitations(ResourceDescription resource, Collection<String> idsToEvict) {
         resourceService.computeResource(resource, state -> {
             InvitationsMap invitations = ProxyUtil.convertToObject(state, InvitationsMap.class);
             if (invitations == null) {
-                invitations = new InvitationsMap(new HashMap<>());
+                return null;
             }
             Map<String, Invitation> invitationMap = invitations.getInvitations();
             idsToEvict.forEach(invitationMap::remove);
@@ -139,7 +164,7 @@ private void cleanUpExpiredInvitations(ResourceDescription resource, Collection<
     }
 
     @Nullable
-    private ResourceDescription getInvitationResource(String invitationId) {
+    public ResourceDescription getInvitationResource(String invitationId) {
         // decrypt invitation ID to obtain its location
         String decryptedInvitationPath = encryptionService.decrypt(invitationId);
         if (decryptedInvitationPath == null) {

src/main/java/com/epam/aidial/core/service/LockService.java

@@ -1,5 +1,7 @@
 package com.epam.aidial.core.service;
 
+import com.epam.aidial.core.Proxy;
+import com.epam.aidial.core.storage.BlobStorageUtil;
 import lombok.extern.slf4j.Slf4j;
 import org.redisson.api.RScript;
 import org.redisson.api.RedissonClient;
@@ -9,6 +11,7 @@
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.LockSupport;
+import java.util.function.Supplier;
 import javax.annotation.Nullable;
 
 /**
@@ -42,6 +45,17 @@ public Lock lock(String key) {
         return () -> unlock(id, owner);
     }
 
+    public <T> T underBucketLock(Proxy proxy, String bucketLocation, Supplier<T> function) {
+        return underBucketLock(bucketLocation, proxy.getStorage().getPrefix(), function);
+    }
+
+    private <T> T underBucketLock(String bucketLocation, @Nullable String prefix, Supplier<T> function) {
+        String key = BlobStorageUtil.toStoragePath(prefix, bucketLocation);
+        try (var ignored = lock(key)) {
+            return function.get();
+        }
+    }
+
     @Nullable
     public Lock tryLock(String key) {
         String id = id(key);

src/main/java/com/epam/aidial/core/service/ResourceService.java

@@ -437,10 +437,7 @@ private void redisSync(String key) {
     }
 
     private String redisKey(ResourceDescription descriptor) {
-        String resourcePath = descriptor.getAbsoluteFilePath();
-        if (prefix != null) {
-            resourcePath = prefix + BlobStorageUtil.PATH_SEPARATOR + resourcePath;
-        }
+        String resourcePath = BlobStorageUtil.toStoragePath(prefix, descriptor.getAbsoluteFilePath());
         return descriptor.getType().name().toLowerCase() + ":" + resourcePath;
     }
 

src/main/java/com/epam/aidial/core/storage/BlobStorage.java

@@ -317,6 +317,6 @@ private void createBucketIfNeeded(Storage config) {
      * @return a full storage path
      */
     private String getStorageLocation(String absoluteFilePath) {
-        return prefix == null ? absoluteFilePath : prefix + BlobStorageUtil.PATH_SEPARATOR + absoluteFilePath;
+        return BlobStorageUtil.toStoragePath(prefix, absoluteFilePath);
     }
 }

src/main/java/com/epam/aidial/core/storage/BlobStorageUtil.java

@@ -57,4 +57,12 @@ public static String buildInitiatorBucket(ProxyContext context) {
     public boolean isFolder(String path) {
         return path.endsWith(PATH_SEPARATOR);
     }
+
+    public String toStoragePath(@Nullable String prefix, String absoluteResourcePath) {
+        if (prefix == null) {
+            return absoluteResourcePath;
+        }
+
+        return prefix + PATH_SEPARATOR + absoluteResourcePath;
+    }
 }