Enable PDF scraping through the Docker API

crawl4ai/processors/pdf/__init__.py

@@ -7,6 +7,8 @@
 from crawl4ai.content_scraping_strategy import ContentScrapingStrategy
 from .processor import NaivePDFProcessorStrategy  # Assuming your current PDF code is in pdf_processor.py
 
+MAX_PDF_DOWNLOAD_REDIRECTS = 10  # Max redirect hops to follow when downloading a PDF
+
 class PDFCrawlerStrategy(AsyncCrawlerStrategy):
     """Crawler strategy for PDF documents.
 
@@ -68,8 +70,10 @@ def __init__(self,
                  extract_images : bool = False,
                  image_save_dir : str = None,
                  batch_size: int = 4,
-                 logger: AsyncLogger = None):
+                 logger: AsyncLogger = None,
+                 url_validator=None):
         self.logger = logger
+        self.url_validator = url_validator # vets the download URL before fetch
         self.pdf_processor = NaivePDFProcessorStrategy(
             save_images_locally=save_images_locally,
             extract_images=extract_images,
@@ -164,7 +168,25 @@ def _get_pdf_path(self, url: str) -> str:
 
                 # Download PDF with streaming and timeout
                 # Connection timeout: 10s, Read timeout: 300s (5 minutes for large PDFs)
-                response = requests.get(url, stream=True, timeout=(20, 60 * 10))
+                # Redirects are followed manually so url_validator (when set) can vet every hop BEFORE it is fetched
+                from urllib.parse import urljoin
+                current_url = url
+                for _ in range(MAX_PDF_DOWNLOAD_REDIRECTS):
+                    if self.url_validator:
+                        self.url_validator(current_url)
+                    response = requests.get(current_url, stream=True, timeout=(20, 60 * 10),
+                                            allow_redirects=False)
+                    if response.is_redirect:
+                        location = response.headers.get("location")
+                        response.close()  # hop response holds its connection open (stream=True)
+                        if not location:
+                            raise RuntimeError(f"Redirect without Location header from {current_url}")
+                        current_url = urljoin(current_url, location)
+                        continue
+                    break
+                else:
+                    raise RuntimeError(
+                        f"Too many redirects (>{MAX_PDF_DOWNLOAD_REDIRECTS}) downloading PDF from {url}")
                 response.raise_for_status()
 
                 # Get file size if available

deploy/docker/api.py

@@ -85,6 +85,7 @@ def _attach_declarative_hooks(crawler, hooks_config: dict) -> dict:
     get_llm_base_url,
     get_redis_task_ttl,
     validate_url_destination,
+    datetime_handler,
 )
 from webhook import WebhookDeliveryService
 
@@ -602,11 +603,19 @@ def create_task_response(task: dict, task_id: str, base_url: str) -> dict:
 
     return response
 
+async def _dispose_crawler(crawler):
+    """Close a dedicated PDF crawler (not pooled) or release a pooled one."""
+    from crawl4ai.processors.pdf import PDFCrawlerStrategy
+    from crawler_pool import release_crawler
+    if isinstance(crawler.crawler_strategy, PDFCrawlerStrategy):
+        await crawler.close()
+    else:
+        await release_crawler(crawler)
+
 async def stream_results(crawler: AsyncWebCrawler, results_gen: AsyncGenerator) -> AsyncGenerator[bytes, None]:
     """Stream results with heartbeats and completion markers."""
     import json
     from utils import datetime_handler
-    from crawler_pool import release_crawler
 
     try:
         async for result in results_gen:
@@ -634,7 +643,7 @@ async def stream_results(crawler: AsyncWebCrawler, results_gen: AsyncGenerator)
         logger.warning("Client disconnected during streaming")
     finally:
         if crawler:
-            await release_crawler(crawler)
+            await _dispose_crawler(crawler)
 
 
 def _normalize_and_validate_seeds(urls: List[str]) -> List[str]:
@@ -659,6 +668,7 @@ async def handle_crawl_request(
     # Track request start
     request_id = f"req_{uuid4().hex[:8]}"
     crawler = None
+    is_pdf_crawl = False
     try:
         from monitor import get_monitor
         await get_monitor().track_request_start(
@@ -689,7 +699,19 @@ async def handle_crawl_request(
         )
 
         from crawler_pool import get_crawler, release_crawler
-        crawler = await get_crawler(browser_config)
+        from crawl4ai.processors.pdf import PDFContentScrapingStrategy, PDFCrawlerStrategy
+        is_pdf_crawl = isinstance(crawler_config.scraping_strategy, PDFContentScrapingStrategy)
+        if is_pdf_crawl:
+            if hooks_config:
+                # PDFCrawlerStrategy has no browser page, so hooks can't attach to it
+                raise HTTPException(status_code=400, detail="Hooks are not supported with PDFContentScrapingStrategy")
+            # SSRF protection: vet the PDF download URL and every redirect hop
+            crawler_config.scraping_strategy.url_validator = validate_url_destination
+            # Use PDFCrawlerStrategy when scraping PDFs, as headless Chromium can't render PDFs inline
+            crawler = AsyncWebCrawler(crawler_strategy=PDFCrawlerStrategy())
+            await crawler.start()
+        else:
+            crawler = await get_crawler(browser_config)
 
         # Attach declarative hooks if provided
         hooks_status = {}
@@ -709,6 +731,9 @@ async def handle_crawl_request(
                         current_value = getattr(cfg, key)
                         if current_value is None or current_value == "":
                             setattr(cfg, key, value)
+                # SSRF: per-URL PDF strategies need the validator wired too
+                if isinstance(cfg.scraping_strategy, PDFContentScrapingStrategy):
+                    cfg.scraping_strategy.url_validator = validate_url_destination
             effective_config = config_list
         else:
             # Single config (original behavior)
@@ -771,6 +796,10 @@ async def handle_crawl_request(
                 if result_dict.get('pdf') is not None and isinstance(result_dict.get('pdf'), bytes):
                     result_dict['pdf'] = b64encode(result_dict['pdf']).decode('utf-8')
 
+                if is_pdf_crawl:
+                    # PDF metadata contains datetimes that JSONResponse can't serialize
+                    result_dict = json.loads(json.dumps(result_dict, default=datetime_handler))
+
                 processed_results.append(result_dict)
             except Exception as e:
                 logger.error(f"Error processing result: {e}")
@@ -851,7 +880,10 @@ async def handle_crawl_request(
         )
     finally:
         if crawler:
-            await release_crawler(crawler)
+            if is_pdf_crawl:
+                await crawler.close()  # not pooled; release_crawler would be a no-op
+            else:
+                await release_crawler(crawler)
 
 async def handle_stream_crawl_request(
     urls: List[str],
@@ -894,8 +926,19 @@ async def handle_stream_crawl_request(
                 ),
             )
 
-        from crawler_pool import get_crawler, release_crawler
-        crawler = await get_crawler(browser_config)
+        from crawler_pool import get_crawler
+        from crawl4ai.processors.pdf import PDFContentScrapingStrategy, PDFCrawlerStrategy
+        if isinstance(crawler_config.scraping_strategy, PDFContentScrapingStrategy):
+            if hooks_config:
+                # PDFCrawlerStrategy has no browser page, so hooks can't attach to it
+                raise HTTPException(status_code=400, detail="Hooks are not supported with PDFContentScrapingStrategy")
+            # SSRF protection: vet the PDF download URL and every redirect hop
+            crawler_config.scraping_strategy.url_validator = validate_url_destination
+            # Use PDFCrawlerStrategy when scraping PDFs, as headless Chromium can't render PDFs inline
+            crawler = AsyncWebCrawler(crawler_strategy=PDFCrawlerStrategy())
+            await crawler.start()
+        else:
+            crawler = await get_crawler(browser_config)
 
         # Attach declarative hooks if provided
         if hooks_config:
@@ -928,21 +971,21 @@ async def handle_stream_crawl_request(
 
     except (UntrustedConfigError, HookValidationError) as e:
         if crawler:
-            await release_crawler(crawler)
+            await _dispose_crawler(crawler)
         raise HTTPException(status_code=400, detail=f"Rejected request: {e}")
 
     except HTTPException:
         # Deliberate status (e.g. 400 SSRF "URL blocked") must pass through
         # rather than be genericized to 500 by the handler below.
         if crawler:
-            await release_crawler(crawler)
+            await _dispose_crawler(crawler)
         raise
 
     except Exception as e:
-        # Release crawler on setup error (for successful streams,
-        # release happens in stream_results finally block)
+        # Dispose crawler on setup error (for successful streams,
+        # disposal happens in stream_results finally block)
         if crawler:
-            await release_crawler(crawler)
+            await _dispose_crawler(crawler)
         logger.error(f"Stream crawl error: {str(e)}", exc_info=True)
         # Raising HTTPException here will prevent streaming response
         raise HTTPException(