feat: attached files referenced via Data URLs are ignored in auto-sha…

…ring (#519)
epam · Oct 11, 2024 · 9be9a9b · 9be9a9b
1 parent e4ef368
commit 9be9a9b
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 3 deletions.
diff --git a/src/main/java/com/epam/aidial/core/data/ResourceType.java b/src/main/java/com/epam/aidial/core/data/ResourceType.java
@@ -23,7 +23,7 @@ public static ResourceType of(String group) {
             case "invitations" -> INVITATION;
             case "publications" -> PUBLICATION;
             case "applications" -> APPLICATION;
-            default -> throw new IllegalArgumentException("Unsupported group: " + group);
+            default -> throw new IllegalArgumentException("Unsupported resource type: " + group);
         };
     }
 }
diff --git a/src/main/java/com/epam/aidial/core/function/BaseFunction.java b/src/main/java/com/epam/aidial/core/function/BaseFunction.java
@@ -20,8 +20,8 @@ public BaseFunction(Proxy proxy, ProxyContext context) {
 
     @SneakyThrows
     public static ResourceDescription fromAnyUrl(String url, EncryptionService encryption) {
-        if (url == null || UrlUtil.isAbsoluteUrl(url)) {
-            // second check to skip public resource
+        if (url == null || UrlUtil.isAbsoluteUrl(url) || UrlUtil.isDataUrl(url)) {
+            // skipping public resources and Data URLs
             return null;
         }
 

diff --git a/src/main/java/com/epam/aidial/core/util/UrlUtil.java b/src/main/java/com/epam/aidial/core/util/UrlUtil.java
@@ -67,4 +67,8 @@ public String decodePath(String path, boolean checkUri) {
     public boolean isAbsoluteUrl(String url) {
         return ABSOLUTE_URL_PATTERN.matcher(url).find();
     }
+
+    public boolean isDataUrl(String url) {
+        return url.startsWith("data:");
+    }
 }
diff --git a/src/test/java/com/epam/aidial/core/util/UrlUtilTest.java b/src/test/java/com/epam/aidial/core/util/UrlUtilTest.java
@@ -57,4 +57,15 @@ public void testIsAbsoluteUrl() {
         assertFalse(UrlUtil.isAbsoluteUrl("/example/file.txt"));
         assertFalse(UrlUtil.isAbsoluteUrl("file"));
     }
+
+    @Test
+    public void testIsDataUrl() {
+        assertTrue(UrlUtil.isDataUrl("data:whatever"));
+        assertTrue(UrlUtil.isDataUrl("data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAMAAAADCAIAAADZSiLoAAAAF0lEQVR4nGNkYPjPwMDAwMDAxAADCBYAG10BBdmz9y8AAAAASUVORK5CYII="));
+        assertFalse(UrlUtil.isDataUrl("data"));
+        assertFalse(UrlUtil.isDataUrl("DATA:whatever"));
+        assertFalse(UrlUtil.isDataUrl(" data:whatever"));
+        assertFalse(UrlUtil.isDataUrl("data;whatever"));
+        assertFalse(UrlUtil.isDataUrl("https://example.com"));
+    }
 }