Coverage for src/local_deep_research/research_library/services/pdf_storage_manager.py: 94%
208 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""
2PDF Storage Manager for Research Library
4Handles PDF storage across three modes:
5- none: Don't store PDFs (text-only)
6- filesystem: Store PDFs unencrypted on disk (fast, external tool compatible)
7- database: Store PDFs encrypted in SQLCipher database (secure, portable)
8"""
10import hashlib
11import re
12from datetime import datetime, UTC
13from pathlib import Path
14from typing import Optional, Tuple
15from urllib.parse import urlparse
17from loguru import logger
18from sqlalchemy.orm import Session
20from ...constants import FILE_PATH_SENTINELS
21from ...utilities.arxiv import extract_arxiv_id
22from ...database.models.library import Document, DocumentBlob
23from ...security.path_validator import PathValidator, resolve_allow_missing
26# Default storage cap for individual PDFs (megabytes). Mirrors the
27# upload-validator cap (`FileUploadValidator.MAX_FILE_SIZE`, configurable
28# via `LDR_SECURITY_UPLOAD_MAX_FILE_SIZE_MB`) so a file that passes the
29# upload step won't be silently dropped at storage time. The runtime
30# value comes from the `research_library.max_pdf_size_mb` setting; this
31# constant is the shared fallback used by every code-level default so the
32# limit doesn't drift across files.
33DEFAULT_MAX_PDF_SIZE_MB = 3072 # 3 GB
36def filesystem_pdf_storage_allowed() -> bool:
37 """Return whether the operator enabled unencrypted filesystem PDF storage.
39 Environment-only operator gate mirroring
40 ``policy.allow_unprotected_egress``: the ``filesystem`` mode writes
41 library-downloaded third-party PDFs as PLAINTEXT to disk (cleartext
42 storage of sensitive information, CWE-312), so it is disabled by default
43 and cannot be turned on through the user-writable settings API — only
44 via the environment variable
45 ``LDR_RESEARCH_LIBRARY_ALLOW_FILESYSTEM_PDF_STORAGE=true``.
46 """
47 from ...settings.env_registry import get_env_setting
49 return bool(
50 get_env_setting("research_library.allow_filesystem_pdf_storage", False)
51 )
54_COERCION_WARNED = False
57def resolve_pdf_storage_mode(mode):
58 """Resolve the effective PDF storage mode, enforcing the operator gate.
60 This is the runtime enforcement point (PEP) for the unencrypted-storage
61 gate: hiding the ``filesystem`` option in the settings UI is not enough,
62 because a value stored before the gate existed — or an
63 ``LDR_RESEARCH_LIBRARY_PDF_STORAGE_MODE=filesystem`` env override —
64 could still resolve to ``filesystem``. Callers that are about to WRITE
65 a PDF must resolve the mode through here first.
67 When the mode is ``filesystem`` and the gate is off, it is coerced to
68 the encrypted ``database`` default. Reads are unaffected:
69 ``PDFStorageManager.load_pdf`` checks the database first and then falls
70 back to the filesystem regardless of ``storage_mode``, so any
71 previously-written plaintext files stay readable after coercion.
72 """
73 global _COERCION_WARNED
74 if isinstance(mode, str) and mode.strip().lower() == "filesystem":
75 if not filesystem_pdf_storage_allowed():
76 if not _COERCION_WARNED:
77 _COERCION_WARNED = True
78 logger.warning(
79 "Unencrypted filesystem PDF storage is operator-gated and "
80 "disabled; coercing pdf_storage_mode 'filesystem' -> 'database' "
81 "(set LDR_RESEARCH_LIBRARY_ALLOW_FILESYSTEM_PDF_STORAGE=true to "
82 "opt in)."
83 )
84 return "database"
85 _COERCION_WARNED = False
86 return mode
89def infer_storage_mode(document: Document) -> str:
90 """Infer where a document's PDF lives when ``storage_mode`` is unset.
92 Documents written before ``storage_mode`` existed have it NULL: a blob
93 means database storage, a real ``file_path`` means filesystem storage,
94 and anything else means no stored PDF.
95 """
96 # If there's a blob, it's database storage
97 if hasattr(document, "blob") and document.blob:
98 return "database"
99 # If there's a file_path (and not a sentinel), it's filesystem
100 if document.file_path and document.file_path not in FILE_PATH_SENTINELS:
101 return "filesystem"
102 # Otherwise no storage
103 return "none"
106def effective_storage_mode(document: Document) -> str:
107 """Return ``storage_mode``, inferring it when the column is NULL."""
108 return document.storage_mode or infer_storage_mode(document)
111# Keep the historical paper-derived basename; filesystem writes add the
112# content hash so different versions cannot overwrite one another.
113_ARXIV_VERSION_SUFFIX = re.compile(r"v\d+$")
116class PDFStorageManager:
117 """Unified interface for PDF storage across all modes."""
119 def __init__(
120 self,
121 library_root: Path,
122 storage_mode: str,
123 max_pdf_size_mb: int = DEFAULT_MAX_PDF_SIZE_MB,
124 *,
125 username: Optional[str] = None,
126 ):
127 """
128 Initialize PDF storage manager.
130 Args:
131 library_root: Base directory for filesystem storage. With per-user
132 library isolation (issue #5521) this is the per-user
133 directory; new PDFs are always WRITTEN here.
134 storage_mode: One of 'none', 'filesystem', 'database'
135 max_pdf_size_mb: Maximum PDF file size in MB. Should not
136 exceed `FileUploadValidator.MAX_FILE_SIZE` (the upload
137 validator's per-file cap, default 3 GB) — uploads above
138 that cap are rejected before they reach this layer.
139 username: Owning user for destructive operations. When falsy,
140 ``library_root`` cannot be guaranteed to be a per-user root
141 (``apply_user_subdir`` returns the bare shared root for an
142 empty username), so ``delete_pdf`` fails closed rather than
143 unlink within a possibly-shared directory.
144 """
145 self.library_root = Path(library_root).resolve()
146 self.storage_mode = storage_mode
147 self.username = username
148 self.max_pdf_size_bytes = max_pdf_size_mb * 1024 * 1024
149 if storage_mode not in ("none", "filesystem", "database"):
150 logger.warning(
151 f"Unknown storage mode '{storage_mode}', defaulting to 'none'"
152 )
153 self.storage_mode = "none"
155 def _safe_path_in_root(
156 self, relative_path: str, root: Path
157 ) -> Optional[Path]:
158 """Validate ``relative_path`` inside ``root`` (traversal + symlink
159 safe). Returns the validated Path (which may not exist), or None if
160 the path is invalid/unsafe.
162 Thin delegate to the shared research-library implementation
163 (``research_library.utils._resolve_within_root``): the sanctioned
164 trusted-"pdfs" leaf layout (a symlinked pdfs directory onto a
165 larger disk) must resolve identically on the storage-manager side
166 and through every settings-based resolver, so the containment
167 rule — including its carve-out — lives in exactly one place
168 (GHSA-gv87-qcw7-gjr8).
169 """
170 from ..utils import _resolve_within_root
172 return _resolve_within_root(relative_path, root)
174 def _get_safe_file_path(self, relative_path: str) -> Optional[Path]:
175 """Resolve only inside this manager's library root.
177 A missing per-user file must not fall back to a shared legacy root:
178 colliding relative filenames do not establish ownership. Legacy files
179 require explicit recovery after the operator verifies their owner.
180 """
181 if not relative_path or relative_path in FILE_PATH_SENTINELS:
182 return None
183 return self._safe_path_in_root(relative_path, self.library_root)
185 def save_pdf(
186 self,
187 pdf_content: bytes,
188 document: Document,
189 session: Session,
190 filename: str,
191 url: Optional[str] = None,
192 resource_id: Optional[int] = None,
193 ) -> Tuple[Optional[str], int]:
194 """
195 Save PDF based on configured storage mode.
197 Args:
198 pdf_content: Raw PDF bytes
199 document: Document model instance
200 session: Database session
201 filename: Filename to use for saving
202 url: Source URL (for generating better filenames)
203 resource_id: Resource ID (for generating better filenames)
205 Returns:
206 Tuple of (file_path or storage indicator, file_size)
207 - For filesystem: relative path string
208 - For database: "database"
209 - For none: None
210 """
211 file_size = len(pdf_content)
213 # Check file size limit
214 if file_size > self.max_pdf_size_bytes:
215 max_mb = self.max_pdf_size_bytes / (1024 * 1024)
216 logger.warning(
217 f"PDF size ({file_size / (1024 * 1024):.1f}MB) exceeds limit "
218 f"({max_mb:.0f}MB), skipping storage"
219 )
220 return None, file_size
222 if self.storage_mode == "none":
223 logger.debug("PDF storage mode is 'none' - skipping PDF save")
224 return None, file_size
226 if self.storage_mode == "filesystem":
227 file_path = self._save_to_filesystem(
228 pdf_content, filename, url, resource_id
229 )
230 relative_path = str(file_path.relative_to(self.library_root))
231 document.storage_mode = "filesystem"
232 document.file_path = relative_path
233 logger.info(f"PDF saved to filesystem: {relative_path}")
234 return relative_path, file_size
236 if self.storage_mode == "database": 236 ↛ 243line 236 didn't jump to line 243 because the condition on line 236 was always true
237 self._save_to_database(pdf_content, document, session)
238 document.storage_mode = "database"
239 document.file_path = None # No filesystem path
240 logger.info(f"PDF saved to database for document {document.id}")
241 return "database", file_size
243 return None, file_size
245 def load_pdf(self, document: Document, session: Session) -> Optional[bytes]:
246 """
247 Load PDF - check database first, then filesystem.
249 Smart retrieval: doesn't rely on storage_mode column, actually checks
250 where the PDF exists.
252 Args:
253 document: Document model instance
254 session: Database session
256 Returns:
257 PDF bytes or None if not available
258 """
259 # 1. Check database first
260 pdf_bytes = self._load_from_database(document, session)
261 if pdf_bytes:
262 logger.debug(f"Loaded PDF from database for document {document.id}")
263 return pdf_bytes
265 # 2. Fallback to filesystem
266 pdf_bytes = self._load_from_filesystem(document)
267 if pdf_bytes:
268 logger.debug(
269 f"Loaded PDF from filesystem for document {document.id}"
270 )
271 return pdf_bytes
273 logger.debug(f"No PDF available for document {document.id}")
274 return None
276 def has_pdf(self, document: Document, session: Session) -> bool:
277 """
278 Check if PDF is available without loading the actual bytes.
280 Args:
281 document: Document model instance
282 session: Database session
284 Returns:
285 True if PDF is available (in database or filesystem)
286 """
287 # Must be a PDF file type
288 if document.file_type != "pdf":
289 return False
291 # Check database first (has blob?)
292 from ...database.models.library import DocumentBlob
294 has_blob = (
295 session.query(DocumentBlob.document_id)
296 .filter_by(document_id=document.id)
297 .first()
298 is not None
299 )
300 if has_blob:
301 return True
303 # Check filesystem (with path traversal protection)
304 file_path = self._get_safe_file_path(document.file_path)
305 if file_path and file_path.is_file():
306 return True
308 return False
310 @classmethod
311 def pdf_exists(cls, library_root, document, session):
312 """Check if a PDF exists in any storage backend.
314 Use this when you need to check PDF availability without a specific
315 storage mode — e.g. generating document URLs in search results.
316 Filesystem reads stay inside the supplied per-user library root.
317 """
318 manager = cls(library_root, "none")
319 return manager.has_pdf(document, session)
321 def _infer_storage_mode(self, document: Document) -> str:
322 """
323 Infer storage mode for documents without explicit mode set.
324 Used for backward compatibility with existing documents.
325 """
326 return infer_storage_mode(document)
328 def _save_to_filesystem(
329 self,
330 pdf_content: bytes,
331 filename: str,
332 url: Optional[str] = None,
333 resource_id: Optional[int] = None,
334 ) -> Path:
335 """
336 Save PDF to filesystem with organized structure.
338 Returns:
339 Absolute path to saved file
340 """
341 # Generate better filename if URL is provided
342 if url:
343 filename = self._generate_filename(url, resource_id, filename)
345 # Different snapshots can share a resource or paper-derived name.
346 # New writes include their content identity so saving a later version
347 # cannot replace bytes still referenced by an older document. Existing
348 # stored paths remain readable and are not moved.
349 # nosemgrep: semgrep.rules.path-traversal-risk, reason: this only constructs a filename; validate_safe_path checks the final path and PDF extension inside library_root before write_file_verified below
350 path = Path(filename)
351 content_hash = hashlib.sha256(pdf_content).hexdigest()
352 filename = str(
353 path.with_name(f"{path.stem}_{content_hash}{path.suffix}")
354 )
356 # Create simple flat directory structure - all PDFs in one folder
357 from ...security.directory_creation import create_directory
358 from ..utils import _TRUSTED_STORAGE_LEAF
360 pdf_path = self.library_root / _TRUSTED_STORAGE_LEAF
361 create_directory(
362 pdf_path,
363 context="library PDF storage directory",
364 )
366 # Validate the filename against the pdfs directory itself: the
367 # hard-coded "pdfs" leaf is the sanctioned operator symlink layout
368 # (see _safe_path_in_root), so containment runs against the
369 # resolved directory rather than library_root. The validated name
370 # is then re-joined on the *lexical* pdfs path so the stored
371 # relative path stays expressible against library_root. Stripping
372 # leading slashes preserves the old "pdfs/" + filename join's
373 # normalization of absolute-looking inputs INTO the store.
374 validated_name = PathValidator.validate_safe_path(
375 filename.lstrip("/"),
376 base_dir=str(pdf_path),
377 required_extensions=(".pdf",),
378 )
379 # Re-express the validated remainder (which comes back under the
380 # RESOLVED pdfs directory) on the *lexical* pdfs path, so the
381 # stored relative path stays expressible against library_root and
382 # sub-directory-carrying names keep their shape. The directory is
383 # resolved exactly as validate_safe_path resolved its base.
384 validated_path = pdf_path / validated_name.relative_to(
385 resolve_allow_missing(pdf_path)
386 )
388 # Write the PDF file with security verification
389 # Pass current storage_mode as snapshot since we already validated it
390 from ...security.file_write_verifier import write_file_verified
392 write_file_verified(
393 validated_path,
394 pdf_content,
395 "research_library.pdf_storage_mode",
396 "filesystem",
397 "library PDF storage",
398 mode="wb",
399 settings_snapshot={
400 "research_library.pdf_storage_mode": self.storage_mode
401 },
402 )
404 return Path(validated_path)
406 def _save_to_database(
407 self, pdf_content: bytes, document: Document, session: Session
408 ) -> None:
409 """Store PDF in document_blobs table."""
410 # Check if blob already exists
411 existing_blob = (
412 session.query(DocumentBlob)
413 .filter_by(document_id=document.id)
414 .first()
415 )
417 if existing_blob:
418 # Update existing blob
419 existing_blob.pdf_binary = pdf_content
420 existing_blob.blob_hash = hashlib.sha256(pdf_content).hexdigest()
421 existing_blob.stored_at = datetime.now(UTC)
422 logger.debug(f"Updated existing blob for document {document.id}")
423 else:
424 # Create new blob
425 blob = DocumentBlob(
426 document_id=document.id,
427 pdf_binary=pdf_content,
428 blob_hash=hashlib.sha256(pdf_content).hexdigest(),
429 stored_at=datetime.now(UTC),
430 )
431 session.add(blob)
432 logger.debug(f"Created new blob for document {document.id}")
434 def _load_from_filesystem(self, document: Document) -> Optional[bytes]:
435 """Load PDF from filesystem with path traversal protection."""
436 # Use safe path resolution to prevent path traversal attacks
437 file_path = self._get_safe_file_path(document.file_path)
438 if not file_path:
439 return None
441 if not file_path.is_file():
442 logger.warning(f"PDF file not found: {file_path}")
443 return None
445 try:
446 return file_path.read_bytes()
447 except Exception:
448 logger.exception(f"Failed to read PDF from {file_path}")
449 return None
451 def _load_from_database(
452 self, document: Document, session: Session
453 ) -> Optional[bytes]:
454 """Load PDF from document_blobs table."""
455 blob = (
456 session.query(DocumentBlob)
457 .filter_by(document_id=document.id)
458 .first()
459 )
461 if not blob:
462 logger.debug(f"No blob found for document {document.id}")
463 return None
465 # Update last accessed timestamp
466 blob.last_accessed = datetime.now(UTC)
468 return blob.pdf_binary
470 def _generate_filename(
471 self, url: str, resource_id: Optional[int], fallback_filename: str
472 ) -> str:
473 """Generate a meaningful filename from URL."""
474 parsed_url = urlparse(url)
475 hostname = parsed_url.hostname or ""
476 timestamp = datetime.now(UTC).strftime("%Y%m%d")
478 # Paper identity rather than host membership (#6413): an arXiv URL that
479 # names no paper is handled by the generic pipeline, so naming its file
480 # `arxiv_...` claimed a provenance nothing else agreed with.
481 arxiv_id = extract_arxiv_id(url)
482 if arxiv_id:
483 # Two shapes the previous `\d{4}\.\d{4,5}` search did not produce:
484 # old-style identifiers carry a slash (`cs.AI/0701001`), which is a
485 # path separator, and every identifier may carry a version. The
486 # version is dropped from the historical basename. New filesystem
487 # writes append a content hash to distinguish saved snapshots.
488 bare = _ARXIV_VERSION_SUFFIX.sub("", arxiv_id)
489 return f"arxiv_{bare.replace('/', '_')}.pdf"
491 if hostname == "ncbi.nlm.nih.gov" and "/pmc" in parsed_url.path:
492 # Extract PMC ID
493 match = re.search(r"(PMC\d+)", url)
494 if match:
495 return f"pmc_{match.group(1)}.pdf"
496 return f"pubmed_{timestamp}_{resource_id or 'unknown'}.pdf"
498 # Use fallback filename
499 return fallback_filename
501 def delete_pdf(self, document: Document, session: Session) -> bool:
502 """
503 Delete PDF for a document.
505 Args:
506 document: Document model instance
507 session: Database session
509 Returns:
510 True if deletion succeeded
511 """
512 storage_mode = effective_storage_mode(document)
514 try:
515 if storage_mode == "filesystem":
516 # Fail closed without a user context: an empty username means
517 # library_root may be the bare shared root (apply_user_subdir
518 # returns the shared base for an empty username), where a
519 # colliding relative path can belong to another tenant. Refuse
520 # the unlink rather than risk deleting another user's file.
521 if not self.username:
522 logger.warning(
523 "Refusing filesystem PDF delete for document "
524 f"{document.id}: no user context, cannot confirm the "
525 "library root is per-user (would risk unlinking a "
526 "shared-root file)."
527 )
528 return False
529 # Use safe path resolution to prevent path traversal attacks.
530 # Delete within our own root only — never the shared legacy
531 # root, where a colliding path can be another tenant's file.
532 file_path = self._get_safe_file_path(document.file_path)
533 if file_path and file_path.is_file():
534 file_path.unlink()
535 logger.info(f"Deleted PDF file: {file_path}")
536 document.file_path = None
537 document.storage_mode = "none"
538 return True
540 if storage_mode == "database":
541 blob = (
542 session.query(DocumentBlob)
543 .filter_by(document_id=document.id)
544 .first()
545 )
546 if blob: 546 ↛ 549line 546 didn't jump to line 549 because the condition on line 546 was always true
547 session.delete(blob)
548 logger.info(f"Deleted PDF blob for document {document.id}")
549 document.storage_mode = "none"
550 return True
552 return True # Nothing to delete for 'none' mode
554 except Exception:
555 logger.exception(f"Failed to delete PDF for document {document.id}")
556 return False
558 def upgrade_to_pdf(
559 self, document: Document, pdf_content: bytes, session: Session
560 ) -> bool:
561 """
562 Upgrade a text-only document to include PDF storage.
564 If document already has a PDF stored, returns False (no action needed).
565 If document is text-only, adds the PDF blob and updates storage_mode.
567 Args:
568 document: Document model instance
569 pdf_content: Raw PDF bytes
570 session: Database session
572 Returns:
573 True if PDF was added, False if already had PDF or failed
574 """
575 # Only upgrade if document is currently text-only. A NULL
576 # storage_mode predates the column, so infer it: a legacy row whose
577 # file_path names a stored PDF is not text-only, and replacing its
578 # file_path would orphan that file.
579 storage_mode = effective_storage_mode(document)
580 if storage_mode != "none":
581 logger.debug(
582 f"Document {document.id} already has storage_mode={storage_mode}"
583 )
584 return False
586 # Check if blob already exists (shouldn't happen, but be safe)
587 existing_blob = (
588 session.query(DocumentBlob)
589 .filter_by(document_id=document.id)
590 .first()
591 )
592 if existing_blob: 592 ↛ 593line 592 didn't jump to line 593 because the condition on line 592 was never true
593 logger.debug(f"Document {document.id} already has a blob")
594 return False
596 # Check file size
597 file_size = len(pdf_content)
598 if file_size > self.max_pdf_size_bytes:
599 max_mb = self.max_pdf_size_bytes / (1024 * 1024)
600 logger.warning(
601 f"PDF size ({file_size / (1024 * 1024):.1f}MB) exceeds limit "
602 f"({max_mb:.0f}MB), skipping upgrade"
603 )
604 return False
606 try:
607 # Add the PDF blob
608 self._save_to_database(pdf_content, document, session)
609 document.storage_mode = "database"
610 document.file_path = None
611 logger.info(f"Upgraded document {document.id} with PDF blob")
612 return True
613 except Exception:
614 logger.exception(
615 f"Failed to upgrade document {document.id} with PDF"
616 )
617 return False