From 84662bcdb486eaddb81b88f5146b3b50f4b9633b Mon Sep 17 00:00:00 2001
From: Pavel Mitrafanau <pavel_mitrafanau@epam.com>
Date: Mon, 16 Sep 2024 16:59:13 +0200
Subject: [PATCH] fix: Allow to attach unescaped links to a chat #481 (#482)

---
 .../epam/aidial/core/function/BaseFunction.java    | 11 ++++-------
 .../java/com/epam/aidial/core/util/UrlUtil.java    | 10 ++++++++++
 .../com/epam/aidial/core/util/UrlUtilTest.java     | 14 ++++++++++++++
 3 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/src/main/java/com/epam/aidial/core/function/BaseFunction.java b/src/main/java/com/epam/aidial/core/function/BaseFunction.java
index 4f1c2025e..f31eb6363 100644
--- a/src/main/java/com/epam/aidial/core/function/BaseFunction.java
+++ b/src/main/java/com/epam/aidial/core/function/BaseFunction.java
@@ -4,9 +4,9 @@
 import com.epam.aidial.core.ProxyContext;
 import com.epam.aidial.core.security.EncryptionService;
 import com.epam.aidial.core.storage.ResourceDescription;
+import com.epam.aidial.core.util.UrlUtil;
 import lombok.SneakyThrows;
 
-import java.net.URI;
 import java.util.function.Function;
 
 public abstract class BaseFunction<T, R> implements Function<T, R> {
@@ -20,14 +20,11 @@ public BaseFunction(Proxy proxy, ProxyContext context) {
 
     @SneakyThrows
     public static ResourceDescription fromAnyUrl(String url, EncryptionService encryption) {
-        if (url == null) {
-            return null;
-        }
-        URI uri = new URI(url);
-        if (uri.isAbsolute()) {
-            // skip public resource
+        if (url == null || UrlUtil.isAbsoluteUrl(url)) {
+            // second check to skip public resource
             return null;
         }
+
         return ResourceDescription.fromAnyUrl(url, encryption);
     }
 }
diff --git a/src/main/java/com/epam/aidial/core/util/UrlUtil.java b/src/main/java/com/epam/aidial/core/util/UrlUtil.java
index f35e4f6e5..a8ba00035 100644
--- a/src/main/java/com/epam/aidial/core/util/UrlUtil.java
+++ b/src/main/java/com/epam/aidial/core/util/UrlUtil.java
@@ -9,6 +9,7 @@
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.nio.charset.Charset;
+import java.util.regex.Pattern;
 
 @UtilityClass
 public class UrlUtil {
@@ -16,6 +17,11 @@ public class UrlUtil {
     private static final PercentCodec DECODER = new PercentCodec();
     private static final Escaper ENCODER = UrlEscapers.urlPathSegmentEscaper();
 
+    /**
+     * Universal, non case-sensitive, protocol-agnostic URL pattern
+     */
+    private static final Pattern ABSOLUTE_URL_PATTERN = Pattern.compile("^[a-z][a-z0-9-+.]*?://", Pattern.CASE_INSENSITIVE);
+
     @SneakyThrows
     public String encodePath(String path) {
         return ENCODER.escape(path);
@@ -39,4 +45,8 @@ public String decodePath(String path, boolean checkUri) {
         }
         return new String(DECODER.decode(path.getBytes(Charset.defaultCharset())));
     }
+
+    public boolean isAbsoluteUrl(String url) {
+        return ABSOLUTE_URL_PATTERN.matcher(url).find();
+    }
 }
diff --git a/src/test/java/com/epam/aidial/core/util/UrlUtilTest.java b/src/test/java/com/epam/aidial/core/util/UrlUtilTest.java
index 60eed56d3..0d39f9e44 100644
--- a/src/test/java/com/epam/aidial/core/util/UrlUtilTest.java
+++ b/src/test/java/com/epam/aidial/core/util/UrlUtilTest.java
@@ -3,7 +3,9 @@
 import org.junit.jupiter.api.Test;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 public class UrlUtilTest {
 
@@ -29,4 +31,16 @@ public void testPathDecoding() {
         assertThrows(IllegalArgumentException.class, () -> UrlUtil.decodePath("/folder1/file#5.txt"));
         assertThrows(IllegalArgumentException.class, () -> UrlUtil.decodePath("/folder1/q?file=5.txt"));
     }
+
+    @Test
+    public void testIsAbsoluteUrl() {
+        assertTrue(UrlUtil.isAbsoluteUrl("test://example.com"));
+        assertTrue(UrlUtil.isAbsoluteUrl("TEST://EXAMPLE.COM"));
+        assertTrue(UrlUtil.isAbsoluteUrl("test1+test2://example.com/item"));
+        assertTrue(UrlUtil.isAbsoluteUrl("test1.test2://example.com/item"));
+        assertTrue(UrlUtil.isAbsoluteUrl("test1-test2://example.com/item"));
+        assertFalse(UrlUtil.isAbsoluteUrl("//example.com/item"));
+        assertFalse(UrlUtil.isAbsoluteUrl("/example/file.txt"));
+        assertFalse(UrlUtil.isAbsoluteUrl("file"));
+    }
 }