diff --git a/api/apps/restful_apis/document_api.py b/api/apps/restful_apis/document_api.py index 4642ce2a57..192064101f 100644 --- a/api/apps/restful_apis/document_api.py +++ b/api/apps/restful_apis/document_api.py @@ -519,7 +519,11 @@ async def _upload_web_document(dataset_id, kb, tenant_id): if not is_valid_url(url): return get_error_data_result(message="The URL format is invalid", code=RetCode.ARGUMENT_ERROR) - blob = html2pdf(url) + try: + blob = await thread_pool_exec(html2pdf, url) + except Exception as e: + logging.warning("html2pdf failed for %s, %s", dataset_id, str(e)) + return get_error_data_result(message=str(e), code=RetCode.SERVER_ERROR) if not blob: return server_error_response(ValueError("Download failure.")) diff --git a/api/db/services/file_service.py b/api/db/services/file_service.py index 8c80466d7d..72728e7183 100644 --- a/api/db/services/file_service.py +++ b/api/db/services/file_service.py @@ -762,73 +762,75 @@ class FileService(CommonService): if url: import requests as _requests from urllib.parse import urljoin as _urljoin + from api.utils.web_utils import BROWSER_FETCH_TIMEOUT, browser_fetch_slot _MAX_CRAWL_REDIRECTS = 10 - # Pre-resolve the full redirect chain so that AsyncWebCrawler never - # follows a server-sent redirect to an unvalidated (potentially - # internal) host. Each hop is SSRF-checked before being followed; - # the validated (hostname, ip) pairs are pinned via Chromium's - # --host-resolver-rules so the browser cannot re-resolve any of them - # through a fresh DNS query. - current_url = url - current_hostname, current_ip = FileService._validate_url_for_crawl(current_url) - # Accumulate MAP rules for every hostname we encounter in the chain. - host_pins: dict[str, str] = {current_hostname: current_ip} + with browser_fetch_slot(): + # Pre-resolve the full redirect chain so that AsyncWebCrawler never + # follows a server-sent redirect to an unvalidated (potentially + # internal) host. Each hop is SSRF-checked before being followed; + # the validated (hostname, ip) pairs are pinned via Chromium's + # --host-resolver-rules so the browser cannot re-resolve any of them + # through a fresh DNS query. + current_url = url + current_hostname, current_ip = FileService._validate_url_for_crawl(current_url) + # Accumulate MAP rules for every hostname we encounter in the chain. + host_pins: dict[str, str] = {current_hostname: current_ip} - for _ in range(_MAX_CRAWL_REDIRECTS): - try: - _resp = _requests.get( - current_url, - timeout=10, - allow_redirects=False, + for _ in range(_MAX_CRAWL_REDIRECTS): + try: + _resp = _requests.get( + current_url, + timeout=10, + allow_redirects=False, + ) + except _requests.RequestException as _exc: + raise ValueError(f"Failed to fetch {current_url!r}: {_exc}") from _exc + + if _resp.status_code not in (301, 302, 303, 307, 308): + break + + _location = _resp.headers.get("Location") + if not _location: + break + + _next_url = _urljoin(current_url, _location) + _next_hostname, _next_ip = FileService._validate_url_for_crawl(_next_url) + host_pins[_next_hostname] = _next_ip + current_url = _next_url + else: + raise ValueError(f"Exceeded {_MAX_CRAWL_REDIRECTS} redirects fetching {url!r}") + + # Build a single MAP rule string covering every validated hostname + # in the redirect chain. Chromium uses the pinned IP for each, + # skipping DNS entirely and eliminating the rebinding window. + _map_rules = ",".join(f"MAP {h} {ip}" for h, ip in host_pins.items()) + + from crawl4ai import AsyncWebCrawler, BrowserConfig, CrawlerRunConfig, DefaultMarkdownGenerator, PruningContentFilter, CrawlResult + + filename = re.sub(r"\?.*", "", url.split("/")[-1]) + + async def adownload(): + browser_config = BrowserConfig( + headless=True, + verbose=False, + extra_args=[f"--host-resolver-rules={_map_rules}"], ) - except _requests.RequestException as _exc: - raise ValueError(f"Failed to fetch {current_url!r}: {_exc}") from _exc + async with AsyncWebCrawler(config=browser_config) as crawler: + crawler_config = CrawlerRunConfig(markdown_generator=DefaultMarkdownGenerator(content_filter=PruningContentFilter()), pdf=True, screenshot=False) + # Use the final resolved URL so the browser starts at the + # redirect destination rather than re-following the chain. + result: CrawlResult = await asyncio.wait_for(crawler.arun(url=current_url, config=crawler_config), timeout=BROWSER_FETCH_TIMEOUT) + return result - if _resp.status_code not in (301, 302, 303, 307, 308): - break + page = asyncio.run(adownload()) + if page.pdf: + if filename.split(".")[-1].lower() != "pdf": + filename += ".pdf" + return structured(filename, "pdf", page.pdf, page.response_headers["content-type"]) - _location = _resp.headers.get("Location") - if not _location: - break - - _next_url = _urljoin(current_url, _location) - _next_hostname, _next_ip = FileService._validate_url_for_crawl(_next_url) - host_pins[_next_hostname] = _next_ip - current_url = _next_url - else: - raise ValueError(f"Exceeded {_MAX_CRAWL_REDIRECTS} redirects fetching {url!r}") - - # Build a single MAP rule string covering every validated hostname - # in the redirect chain. Chromium uses the pinned IP for each, - # skipping DNS entirely and eliminating the rebinding window. - _map_rules = ",".join(f"MAP {h} {ip}" for h, ip in host_pins.items()) - - from crawl4ai import AsyncWebCrawler, BrowserConfig, CrawlerRunConfig, DefaultMarkdownGenerator, PruningContentFilter, CrawlResult - - filename = re.sub(r"\?.*", "", url.split("/")[-1]) - - async def adownload(): - browser_config = BrowserConfig( - headless=True, - verbose=False, - extra_args=[f"--host-resolver-rules={_map_rules}"], - ) - async with AsyncWebCrawler(config=browser_config) as crawler: - crawler_config = CrawlerRunConfig(markdown_generator=DefaultMarkdownGenerator(content_filter=PruningContentFilter()), pdf=True, screenshot=False) - # Use the final resolved URL so the browser starts at the - # redirect destination rather than re-following the chain. - result: CrawlResult = await crawler.arun(url=current_url, config=crawler_config) - return result - - page = asyncio.run(adownload()) - if page.pdf: - if filename.split(".")[-1].lower() != "pdf": - filename += ".pdf" - return structured(filename, "pdf", page.pdf, page.response_headers["content-type"]) - - return structured(filename, "html", str(page.markdown).encode("utf-8"), page.response_headers["content-type"]) + return structured(filename, "html", str(page.markdown).encode("utf-8"), page.response_headers["content-type"]) DocumentService.check_doc_health(user_id, file.filename) return structured(file.filename, filename_type(file.filename), file.read(), file.content_type) diff --git a/api/utils/web_utils.py b/api/utils/web_utils.py index c0ab4b00d7..5b551015cf 100644 --- a/api/utils/web_utils.py +++ b/api/utils/web_utils.py @@ -15,8 +15,12 @@ # import base64 +from contextlib import contextmanager import json +import logging +import os import re +import threading import aiosmtplib from email.mime.text import MIMEText from email.header import Header @@ -38,6 +42,24 @@ OTP_TTL_SECONDS = 5 * 60 # valid for 5 minutes ATTEMPT_LIMIT = 5 # maximum attempts ATTEMPT_LOCK_SECONDS = 30 * 60 # lock for 30 minutes RESEND_COOLDOWN_SECONDS = 60 # cooldown for 1 minute +BROWSER_FETCH_CONCURRENCY = max(1, int(os.getenv("RAGFLOW_BROWSER_FETCH_CONCURRENCY", "2"))) +BROWSER_FETCH_ACQUIRE_TIMEOUT = float(os.getenv("RAGFLOW_BROWSER_FETCH_ACQUIRE_TIMEOUT", "5")) +BROWSER_FETCH_TIMEOUT = float(os.getenv("RAGFLOW_BROWSER_FETCH_TIMEOUT", "60")) +_BROWSER_FETCH_SEMAPHORE = threading.BoundedSemaphore(BROWSER_FETCH_CONCURRENCY) + + +class BrowserFetchBusy(RuntimeError): + pass + + +@contextmanager +def browser_fetch_slot(timeout: float = BROWSER_FETCH_ACQUIRE_TIMEOUT): + if not _BROWSER_FETCH_SEMAPHORE.acquire(timeout=timeout): + raise BrowserFetchBusy("Too many concurrent browser fetch requests") + try: + yield + finally: + _BROWSER_FETCH_SEMAPHORE.release() from api.utils.file_response import ( # noqa: F401 @@ -69,7 +91,8 @@ def html2pdf( install_driver: bool = True, print_options: dict = {}, ): - result = __get_pdf_from_html(source, timeout, install_driver, print_options) + with browser_fetch_slot(): + result = __get_pdf_from_html(source, timeout, install_driver, print_options) return result @@ -96,20 +119,23 @@ def __get_pdf_from_html(path: str, timeout: int, install_driver: bool, print_opt webdriver_prefs["profile.default_content_settings"] = {"images": 2} - if install_driver: - service = Service(ChromeDriverManager().install()) - driver = webdriver.Chrome(service=service, options=webdriver_options) - else: - driver = webdriver.Chrome(options=webdriver_options) - - driver.get(path) - + driver = None try: - WebDriverWait(driver, timeout).until(staleness_of(driver.find_element(by=By.TAG_NAME, value="html"))) - except TimeoutException: - pass + if install_driver: + service = Service(ChromeDriverManager().install()) + driver = webdriver.Chrome(service=service, options=webdriver_options) + else: + driver = webdriver.Chrome(options=webdriver_options) + + driver.set_page_load_timeout(BROWSER_FETCH_TIMEOUT) + driver.set_script_timeout(BROWSER_FETCH_TIMEOUT) + + try: + driver.get(path) + WebDriverWait(driver, timeout).until(staleness_of(driver.find_element(by=By.TAG_NAME, value="html"))) + except TimeoutException: + logging.warning("Timed out loading %s; printing current page state", path) - try: calculated_print_options = { "landscape": False, "displayHeaderFooter": False, @@ -120,7 +146,8 @@ def __get_pdf_from_html(path: str, timeout: int, install_driver: bool, print_opt result = __send_devtools(driver, "Page.printToPDF", calculated_print_options) return base64.b64decode(result["data"]) finally: - driver.quit() + if driver is not None: + driver.quit() def is_valid_url(url: str) -> bool: