Coverage for lode / helpers / spool.py: 87%

1# lode/helpers/spool.py 

2"""Disk cache (spool) for uploaded files and fetched URLs. 

3 

4Lifecycle: 

5- Entries expire after TTL seconds (default 4 h). 

6- Total size is capped at MAX_BYTES (default 1 GB). 

7- Location is controlled by LODE_SPOOL_DIR; defaults to the system 

8 temp directory so it works out of the box after pip install without 

9 needing write access to the package directory. 

10 

11No dependency on FastAPI or any HTTP layer. 

12""" 

13import os 

14import tempfile 

15import time 

16from uuid import uuid4 

17 

18from lode.exceptions import ArtefactValidationError 

19 

20SPOOL_DIR: str = os.path.realpath( 

21 os.getenv("LODE_SPOOL_DIR", os.path.join(tempfile.gettempdir(), "lode-spool")) 

22) 

23os.makedirs(SPOOL_DIR, exist_ok=True) 

24 

25TTL = 4 * 60 * 60 # seconds before an entry is considered stale 

26MAX_BYTES = 1024 ** 3 # 1 GB total budget 

27 

28 

29def get_path(token: str) -> str: 

30 """Return the absolute spool path for a token. 

31 

32 Validates that the resolved path stays inside SPOOL_DIR to prevent 

33 path-injection via a crafted token. 

34 """ 

35 p = os.path.realpath(os.path.join(SPOOL_DIR, f"{token}.rdf")) 

36 if os.path.commonpath((SPOOL_DIR, p)) != SPOOL_DIR: 

37 raise ArtefactValidationError( 

38 "Invalid upload token", context={"token": token} 

39 ) 

40 return p 

41 

42 

43def prune() -> None: 

44 """Evict expired entries, then enforce the size budget. 

45 

46 Deletes the oldest entries (by mtime) until total size is back 

47 under MAX_BYTES. Best-effort: OSError races are silently ignored. 

48 """ 

49 cutoff = time.time() - TTL 

50 survivors: list[tuple[float, int, str]] = [] 

51 

52 for name in os.listdir(SPOOL_DIR): 

53 p = os.path.join(SPOOL_DIR, name) 

54 try: 

55 st = os.stat(p) 

56 except OSError: 

57 continue 

58 if st.st_mtime < cutoff: 

59 try: 

60 os.unlink(p) 

61 except OSError: 

62 pass 

63 continue 

64 survivors.append((st.st_mtime, st.st_size, p)) 

65 

66 total = sum(size for _, size, _ in survivors) 

67 if total <= MAX_BYTES: 

68 return 

69 

70 for _, size, p in sorted(survivors): # oldest first 

71 if total <= MAX_BYTES: 

72 break 

73 try: 

74 os.unlink(p) 

75 total -= size 

76 except OSError: 

77 pass 

78 

79 

80def save(content: bytes) -> str: 

81 """Persist raw bytes in the spool and return the token.""" 

82 prune() 

83 token = uuid4().hex 

84 with open(get_path(token), "wb") as f: 

85 f.write(content) 

86 return token