Merge branch 'development' into issue-277

epam · Mar 14, 2024 · d5768d1 · d5768d1
2 parents 098a23c + 52ba2a7
commit d5768d1
Show file tree

Hide file tree

Showing 14 changed files with 659 additions and 31 deletions.
diff --git a/src/main/java/com/epam/aidial/core/AiDial.java b/src/main/java/com/epam/aidial/core/AiDial.java
@@ -113,7 +113,7 @@ void start() throws Exception {
             LockService lockService = new LockService(redis, storage.getPrefix());
             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);
+            ShareService shareService = new ShareService(resourceService, invitationService, encryptionService, storage);
             PublicationService publicationService = new PublicationService(encryptionService, resourceService, storage, generator, clock);
             ResourceOperationService resourceOperationService = new ResourceOperationService(resourceService, storage, invitationService, shareService);
 

diff --git a/src/main/java/com/epam/aidial/core/controller/ControllerSelector.java b/src/main/java/com/epam/aidial/core/controller/ControllerSelector.java
@@ -45,10 +45,11 @@ public class ControllerSelector {
     private static final Pattern PATTERN_TOKENIZE = Pattern.compile("^/+v1/deployments/([^/]+)/tokenize$");
     private static final Pattern PATTERN_TRUNCATE_PROMPT = Pattern.compile("^/+v1/deployments/([^/]+)/truncate_prompt$");
 
-    private static final Pattern SHARE_RESOURCE_OPERATIONS = Pattern.compile("^/v1/ops/resource/share/(create|list|discard|revoke)$");
+    private static final Pattern SHARE_RESOURCE_OPERATIONS = Pattern.compile("^/v1/ops/resource/share/(create|list|discard|revoke|copy)$");
     private static final Pattern INVITATIONS = Pattern.compile("^/v1/invitations$");
     private static final Pattern INVITATION = Pattern.compile("^/v1/invitations/([a-zA-Z0-9]+)$");
     private static final Pattern PUBLICATIONS = Pattern.compile("^/v1/ops/publications/(list|get|create|delete|approve|reject)$");
+    private static final Pattern PUBLICATION_RULES = Pattern.compile("^/v1/ops/publications/rules/list$");
 
     private static final Pattern RESOURCE_OPERATIONS = Pattern.compile("^/v1/ops/resources/(move)$");
 
@@ -261,11 +262,17 @@ private static Controller selectPost(Proxy proxy, ProxyContext context, String p
                 case "create" -> controller::createPublication;
                 case "delete" -> controller::deletePublication;
                 case "approve" -> controller::approvePublication;
-                case "reject" -> controller:: rejectPublication;
+                case "reject" -> controller::rejectPublication;
                 default -> null;
             };
         }
 
+        match = match(PUBLICATION_RULES, path);
+        if (match != null) {
+            PublicationController controller = new PublicationController(proxy, context);
+            return controller::listRules;
+        }
+
         match = match(RESOURCE_OPERATIONS, path);
         if (match != null) {
             ResourceOperationController controller = new ResourceOperationController(proxy, context);

diff --git a/src/main/java/com/epam/aidial/core/controller/PublicationController.java b/src/main/java/com/epam/aidial/core/controller/PublicationController.java
@@ -6,6 +6,7 @@
 import com.epam.aidial.core.data.Publications;
 import com.epam.aidial.core.data.ResourceLink;
 import com.epam.aidial.core.data.ResourceType;
+import com.epam.aidial.core.data.Rules;
 import com.epam.aidial.core.security.AccessService;
 import com.epam.aidial.core.security.EncryptionService;
 import com.epam.aidial.core.service.LockService;
@@ -133,6 +134,21 @@ public Future<?> rejectPublication() {
         return Future.succeededFuture();
     }
 
+    public Future<?> listRules() {
+        context.getRequest()
+                .body()
+                .compose(body -> {
+                    String url = ProxyUtil.convertToObject(body, ResourceLink.class).url();
+                    ResourceDescription rule = decodeRule(url);
+                    checkRuleAccess(rule);
+                    return vertx.executeBlocking(() -> publicationService.listRules(rule));
+                })
+                .onSuccess(rules -> context.respond(HttpStatus.OK, new Rules(rules)))
+                .onFailure(error -> respondError("Can't list rules", error));
+
+        return Future.succeededFuture();
+    }
+
     private void respondError(String message, Throwable error) {
         HttpStatus status = HttpStatus.INTERNAL_SERVER_ERROR;
         String body = null;
@@ -172,6 +188,25 @@ private ResourceDescription decodePublication(String path, boolean allowPublic)
         return resource;
     }
 
+    private ResourceDescription decodeRule(String path) {
+        try {
+            if (!path.startsWith(BlobStorageUtil.PUBLIC_LOCATION)) {
+                throw new IllegalArgumentException();
+            }
+
+            String folder = path.substring(BlobStorageUtil.PUBLIC_LOCATION.length());
+            ResourceDescription resource = ResourceDescription.fromEncoded(ResourceType.RULES, BlobStorageUtil.PUBLIC_BUCKET, BlobStorageUtil.PUBLIC_LOCATION, folder);
+
+            if (!resource.isFolder()) {
+                throw new IllegalArgumentException();
+            }
+
+            return resource;
+        } catch (IllegalArgumentException e) {
+            throw new IllegalArgumentException("Invalid resource: " + path, e);
+        }
+    }
+
     private void checkAccess(ResourceDescription resource, boolean allowUser) {
         boolean hasAccess = accessService.hasAdminAccess(context);
 
@@ -184,4 +219,10 @@ private void checkAccess(ResourceDescription resource, boolean allowUser) {
             throw new HttpException(HttpStatus.FORBIDDEN, "Forbidden resource: " + resource.getUrl());
         }
     }
+
+    private void checkRuleAccess(ResourceDescription rule) {
+        if (!accessService.hasReadAccess(rule, context)) {
+            throw new HttpException(HttpStatus.FORBIDDEN, "Forbidden resource: " + rule.getUrl());
+        }
+    }
 }
diff --git a/src/main/java/com/epam/aidial/core/controller/ShareController.java b/src/main/java/com/epam/aidial/core/controller/ShareController.java
@@ -2,17 +2,22 @@
 
 import com.epam.aidial.core.Proxy;
 import com.epam.aidial.core.ProxyContext;
+import com.epam.aidial.core.data.CopySharedAccessRequest;
 import com.epam.aidial.core.data.ListSharedResourcesRequest;
 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.ResourceService;
 import com.epam.aidial.core.service.ShareService;
+import com.epam.aidial.core.storage.BlobStorage;
 import com.epam.aidial.core.storage.BlobStorageUtil;
+import com.epam.aidial.core.storage.ResourceDescription;
 import com.epam.aidial.core.util.HttpException;
 import com.epam.aidial.core.util.HttpStatus;
 import com.epam.aidial.core.util.ProxyUtil;
+import com.epam.aidial.core.util.ResourceUtil;
 import io.vertx.core.Future;
 import io.vertx.core.buffer.Buffer;
 import lombok.extern.slf4j.Slf4j;
@@ -30,6 +35,8 @@ public class ShareController {
     private final EncryptionService encryptionService;
     private final LockService lockService;
     private final InvitationService invitationService;
+    private final ResourceService resourceService;
+    private final BlobStorage storage;
 
     public ShareController(Proxy proxy, ProxyContext context) {
         this.proxy = proxy;
@@ -38,6 +45,8 @@ public ShareController(Proxy proxy, ProxyContext context) {
         this.encryptionService = proxy.getEncryptionService();
         this.lockService = proxy.getLockService();
         this.invitationService = proxy.getInvitationService();
+        this.resourceService = proxy.getResourceService();
+        this.storage = proxy.getStorage();
     }
 
     public Future<?> handle(Operation operation) {
@@ -46,6 +55,7 @@ public Future<?> handle(Operation operation) {
             case CREATE -> createSharedResources();
             case REVOKE -> revokeSharedResources();
             case DISCARD -> discardSharedResources();
+            case COPY -> copySharedAccess();
             default ->
                     context.respond(HttpStatus.INTERNAL_SERVER_ERROR, "Operation %s is not supported".formatted(operation));
         }
@@ -137,6 +147,49 @@ public Future<?> revokeSharedResources() {
                 .onFailure(this::handleServiceError);
     }
 
+    public Future<?> copySharedAccess() {
+        return context.getRequest()
+                .body()
+                .compose(buffer -> {
+                    CopySharedAccessRequest request;
+                    try {
+                        request = ProxyUtil.convertToObject(buffer, CopySharedAccessRequest.class);
+                    } catch (Exception e) {
+                        log.error("Invalid request body provided", e);
+                        throw new IllegalArgumentException("Can't initiate copy shared access request. Incorrect body provided");
+                    }
+
+                    String sourceUrl = request.sourceUrl();
+                    if (sourceUrl == null) {
+                        throw new IllegalArgumentException("sourceUrl must be provided");
+                    }
+                    String destinationUrl = request.destinationUrl();
+                    if (destinationUrl == null) {
+                        throw new IllegalArgumentException("destinationUrl must be provided");
+                    }
+
+                    String bucketLocation = BlobStorageUtil.buildInitiatorBucket(context);
+                    String bucket = encryptionService.encrypt(bucketLocation);
+
+                    ResourceDescription source = ResourceDescription.fromPrivateUrl(sourceUrl, encryptionService);
+                    if (!bucket.equals(source.getBucketName())) {
+                        throw new IllegalArgumentException("sourceUrl does not belong to the user");
+                    }
+                    ResourceDescription destination = ResourceDescription.fromPrivateUrl(destinationUrl, encryptionService);
+                    if (!bucket.equals(destination.getBucketName())) {
+                        throw new IllegalArgumentException("destinationUrl does not belong to the user");
+                    }
+
+                    return proxy.getVertx().executeBlocking(() ->
+                            lockService.underBucketLock(bucketLocation, () -> {
+                                shareService.copySharedAccess(bucket, bucketLocation, source, destination);
+                                return null;
+                            }));
+                })
+                .onSuccess(ignore -> context.respond(HttpStatus.OK))
+                .onFailure(this::handleServiceError);
+    }
+
     private void handleServiceError(Throwable error) {
         if (error instanceof IllegalArgumentException) {
             context.respond(HttpStatus.BAD_REQUEST, error.getMessage());
@@ -157,7 +210,11 @@ private ResourceLinkCollection getResourceLinkCollection(Buffer buffer, Operatio
         }
     }
 
+    private boolean hasResource(ResourceDescription resource) {
+        return ResourceUtil.hasResource(resource, resourceService, storage);
+    }
+
     public enum Operation {
-        CREATE, LIST, DISCARD, REVOKE
+        CREATE, LIST, DISCARD, REVOKE, COPY
     }
 }
diff --git a/src/main/java/com/epam/aidial/core/data/CopySharedAccessRequest.java b/src/main/java/com/epam/aidial/core/data/CopySharedAccessRequest.java
@@ -0,0 +1,4 @@
+package com.epam.aidial.core.data;
+
+public record CopySharedAccessRequest(String sourceUrl, String destinationUrl) {
+}
diff --git a/src/main/java/com/epam/aidial/core/data/Rules.java b/src/main/java/com/epam/aidial/core/data/Rules.java
@@ -0,0 +1,7 @@
+package com.epam.aidial.core.data;
+
+import java.util.List;
+import java.util.Map;
+
+public record Rules(Map<String, List<Rule>> rules) {
+}
diff --git a/src/main/java/com/epam/aidial/core/service/PublicationService.java b/src/main/java/com/epam/aidial/core/service/PublicationService.java
@@ -4,6 +4,7 @@
 import com.epam.aidial.core.data.MetadataBase;
 import com.epam.aidial.core.data.Publication;
 import com.epam.aidial.core.data.ResourceFolderMetadata;
+import com.epam.aidial.core.data.ResourceItemMetadata;
 import com.epam.aidial.core.data.ResourceType;
 import com.epam.aidial.core.data.ResourceUrl;
 import com.epam.aidial.core.data.Rule;
@@ -17,6 +18,7 @@
 import com.fasterxml.jackson.core.type.TypeReference;
 import lombok.RequiredArgsConstructor;
 import org.apache.commons.lang3.mutable.MutableObject;
+import org.apache.commons.lang3.tuple.Pair;
 
 import java.util.Collection;
 import java.util.HashMap;
@@ -25,6 +27,8 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.TreeMap;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.LongSupplier;
 import java.util.function.Supplier;
 import javax.annotation.Nullable;
@@ -53,6 +57,10 @@ public class PublicationService {
 
     private static final Set<ResourceType> ALLOWED_RESOURCES = Set.of(ResourceType.FILE, ResourceType.CONVERSATION, ResourceType.PROMPT);
 
+    /**
+     * Key is updated time from the metadata. Value is decoded map (folder path, list of rules).
+     */
+    private final AtomicReference<Pair<Long, Map<String, List<Rule>>>> cachedRules = new AtomicReference<>();
     private final EncryptionService encryption;
     private final ResourceService resources;
     private final BlobStorage files;
@@ -78,7 +86,7 @@ public boolean hasPublicAccess(ProxyContext context, ResourceDescription resourc
             return false;
         }
 
-        Map<String, List<Rule>> rules = decodeRules(resources.getResource(PUBLIC_RULES));
+        Map<String, List<Rule>> rules = getCachedRules();
         Map<String, Boolean> cache = new HashMap<>();
 
         return evaluate(context, resource, rules, cache);
@@ -89,7 +97,7 @@ public void filterForbidden(ProxyContext context, ResourceDescription folder, Re
             return;
         }
 
-        Map<String, List<Rule>> rules = decodeRules(resources.getResource(PUBLIC_RULES));
+        Map<String, List<Rule>> rules = getCachedRules();
         Map<String, Boolean> cache = new HashMap<>();
         cache.put(ruleUrl(folder), true);
 
@@ -101,6 +109,27 @@ public void filterForbidden(ProxyContext context, ResourceDescription folder, Re
         metadata.setItems(filtered);
     }
 
+    public Map<String, List<Rule>> listRules(ResourceDescription resource) {
+        if (!resource.isFolder() || !resource.isPublic()) {
+            throw new IllegalArgumentException("Bad rule url: " + resource.getUrl());
+        }
+
+        Map<String, List<Rule>> rules = getCachedRules();
+        Map<String, List<Rule>> result = new TreeMap<>();
+
+        while (resource != null) {
+            String url = ruleUrl(resource);;
+            List<Rule> list = rules.get(url);
+            resource = resource.getParent();
+
+            if (list != null) {
+                result.put(url, list);
+            }
+        }
+
+        return result;
+    }
+
     public Collection<Publication> listPublications(ResourceDescription resource) {
         if (resource.getType() != ResourceType.PUBLICATION || !resource.isRootFolder()) {
             throw new IllegalArgumentException("Bad publication url: " + resource.getUrl());
@@ -473,6 +502,24 @@ private String encodeReviewBucket(ResourceDescription bucket, String id) {
         return encryption.encrypt(path);
     }
 
+    private Map<String, List<Rule>> getCachedRules() {
+        ResourceItemMetadata meta = resources.getResourceMetadata(PUBLIC_RULES);
+        long key = (meta == null) ? Long.MIN_VALUE : meta.getUpdatedAt();
+        Pair<Long, Map<String, List<Rule>>> current = cachedRules.get();
+
+        if (current == null || current.getKey() != key) {
+            Pair<ResourceItemMetadata, String> resource = resources.getResourceWithMetadata(PUBLIC_RULES);
+            Pair<Long, Map<String, List<Rule>>> next = (resource == null)
+                    ? Pair.of(Long.MIN_VALUE, decodeRules(null))
+                    : Pair.of(resource.getKey().getUpdatedAt(), decodeRules(resource.getValue()));
+
+            cachedRules.compareAndSet(current, next);
+            current = next;
+        }
+
+        return current.getValue();
+    }
+
     private static boolean evaluate(ProxyContext context,
                                     ResourceDescription resource,
                                     Map<String, List<Rule>> rules,
@@ -542,17 +589,7 @@ private static String encodePublications(Map<String, Publication> publications)
 
     private static Map<String, List<Rule>> decodeRules(String json) {
         Map<String, List<Rule>> rules = ProxyUtil.convertToObject(json, RULES_TYPE);
-
-        if (rules == null) {
-            Rule rule = new Rule();
-            rule.setSource("roles");
-            rule.setFunction(Rule.Function.TRUE);
-            rule.setTargets(List.of());
-            rules = new LinkedHashMap<>();
-            rules.put(PUBLIC_LOCATION, List.of(rule));
-        }
-
-        return rules;
+        return (rules == null) ? new LinkedHashMap<>() : rules;
     }
 
     private static String encodeRules(Map<String, List<Rule>> rules) {