patch: sitemap and web loader

application/parser/remote/sitemap_loader.py

@@ -36,6 +36,11 @@ class SitemapLoader(BaseRemote):
             if self.limit is not None and processed_urls >= self.limit:
                 break  # Stop processing if the limit is reached
 
+            try:
+                url = validate_url(url)
+            except SSRFError as e:
+                logging.error(f"URL validation failed for sitemap entry {url}: {e}")
+                continue
             try:
                 loader = self.loader([url])
                 documents.extend(loader.load())
@@ -90,6 +95,15 @@ class SitemapLoader(BaseRemote):
         # Check for nested sitemaps
         for sitemap in root.findall('.//sitemap/loc'):
             nested_sitemap_url = sitemap.text
+            if not nested_sitemap_url:
+                continue
+            try:
+                nested_sitemap_url = validate_url(nested_sitemap_url)
+            except SSRFError as e:
+                logging.error(
+                    f"URL validation failed for nested sitemap {nested_sitemap_url}: {e}"
+                )
+                continue
             urls.extend(self._extract_urls(nested_sitemap_url))
 
         return urls

application/parser/remote/web_loader.py

@@ -1,8 +1,8 @@
 import logging
 from application.parser.remote.base import BaseRemote
 from application.parser.schema.base import Document
+from application.core.url_validation import validate_url, SSRFError
 from langchain_community.document_loaders import WebBaseLoader
-from urllib.parse import urlparse
 
 headers = {
     "User-Agent": "Mozilla/5.0",
@@ -26,9 +26,11 @@ class WebLoader(BaseRemote):
             urls = [urls]
         documents = []
         for url in urls:
-            # Check if the URL scheme is provided, if not, assume http
-            if not urlparse(url).scheme:
-                url = "http://" + url
+            try:
+                url = validate_url(url)
+            except SSRFError as e:
+                logging.error(f"URL validation failed for {url}: {e}")
+                continue
             try:
                 loader = self.loader([url], header_template=headers)
                 loaded_docs = loader.load()