Coverage for src/local_deep_research/security/toml_limits.py: 100%

1"""Limits checked before tomllib parses TOML that a user supplied. 

2 

3tomllib is pure Python. It rebuilds a key's tuple for every part and, on a 

4dotted key/value line, also builds and keeps every prefix of the key, so 

5one long key costs time (and, on a key/value line, memory) that grows with 

6the square of its length. Uploaded TOML documents and the raw configuration 

7endpoint both parse through ``loads_toml_within_limits``. 

8""" 

9 

10import re 

11import tomllib 

12from itertools import islice 

13from typing import Any 

14 

15# Largest TOML document parsed, in UTF-8 bytes. TOML files are mostly 

16# configuration: a full pyproject.toml is about 30 KB, and gitleaks' rule 

17# set, one of the larger common ones, about 100 KB, so this leaves fivefold 

18# headroom. tomllib is pure Python, and hostile but valid documents were 

19# measured holding several hundred bytes of parser memory per input byte; 

20# neither the general upload limit (3 GB by default) nor the default 16 MB 

21# JSON body limit bounds that. 

22MAX_TOML_BYTES = 512 * 1024 

23 

24# Most dotted parts allowed in one key or table header ("a.b.c" has three). 

25# The deepest key in 252 sampled TOML files had five parts. Under this cap 

26# the total work stays proportional to the document size. A number such as 

27# 1.5 also scans as a two-part run, so keep this at least 2. 

28MAX_TOML_KEY_PARTS = 32 

29 

30# One key part: bare, or a one-line basic or literal string that ends 

31# where tomllib ends it. 

32_KEY_PART = ( 

33 r"(?:[A-Za-z0-9_-]++" 

34 r"|\x22(?:[^\x22\\\n]++|\\[^\n])*+\x22" 

35 r"|\x27[^\x27\n]*+\x27)" 

36) 

37_KEY_PART_RE = re.compile(_KEY_PART) 

38 

39# Some alternative matches at every position, so the tokens cover the text 

40# in order and, up to where tomllib would reject the document, string and 

41# comment contents are skipped exactly where tomllib skips them. A "key" 

42# token is a run of dot-separated key parts. 

43_TOKEN_RE = re.compile( 

44 rf""" 

45 (?P<string> 

46 \x22{{3}} (?: [^\x22\\]++ | \\[\s\S] | \x22(?!\x22{{2}}) )*+ \x22{{3,5}} 

47 | \x27{{3}} (?: [^\x27]++ | \x27(?!\x27{{2}}) )*+ \x27{{3,5}} 

48 ) 

49 | (?P<key> 

50 (?!\x22{{3}}|\x27{{3}}) 

51 {_KEY_PART} (?: [ \t]*+ \. [ \t]*+ {_KEY_PART} )*+ 

52 ) 

53 | (?P<comment> \# [^\n]*+ ) 

54 | (?P<other> [^A-Za-z0-9_\-\x22\x27\#]++ ) 

55 | (?P<unterminated> [\x22\x27] ) 

56 """, 

57 re.VERBOSE, 

58) 

59 

60 

61class TOMLLimitError(ValueError): 

62 """TOML input exceeds a limit that is checked before tomllib runs.""" 

63 

64 

65def _has_overlong_key(text: str, max_parts: int) -> bool: 

66 """Return whether a key or table header has more than *max_parts* parts. 

67 

68 The scan stops at an unterminated string, where tomllib rejects the 

69 document, so it stays linear in the length of *text*. 

70 """ 

71 for token in _TOKEN_RE.finditer(text): 

72 kind = token.lastgroup 

73 if kind == "unterminated": 

74 return False 

75 start, end = token.span() 

76 if kind == "key" and text.find(".", start, end) != -1: 

77 parts = _KEY_PART_RE.finditer(text, start, end) 

78 if sum(1 for _ in islice(parts, max_parts + 1)) > max_parts: 

79 return True 

80 return False 

81 

82 

83def check_toml_size(size: int) -> None: 

84 """Raise TOMLLimitError when a *size*-byte document is over the limit.""" 

85 if size > MAX_TOML_BYTES: 

86 raise TOMLLimitError( 

87 f"TOML document is larger than {MAX_TOML_BYTES} bytes" 

88 ) 

89 

90 

91def loads_toml_within_limits(text: str) -> dict[str, Any]: 

92 """Parse *text* with tomllib unless it exceeds the limits. 

93 

94 Raises TOMLLimitError, before tomllib runs, when the UTF-8 form of 

95 *text* is larger than MAX_TOML_BYTES or a key or table header has more 

96 than MAX_TOML_KEY_PARTS dotted parts, and TypeError when *text* is not 

97 a string. Otherwise it returns or raises what 

98 tomllib.loads does, such as TOMLDecodeError for invalid TOML. 

99 """ 

100 if not isinstance(text, str): 

101 raise TypeError(f"Expected str object, not '{type(text).__qualname__}'") 

102 # A character takes at least one UTF-8 byte, so the length check refuses 

103 # most oversized text before its encoded size is measured. 

104 check_toml_size(len(text)) 

105 check_toml_size(len(text.encode("utf-8", "surrogatepass"))) 

106 # tomllib.loads folds CRLF before it parses; scan that same text. 

107 text = text.replace("\r\n", "\n") 

108 if _has_overlong_key(text, MAX_TOML_KEY_PARTS): 

109 raise TOMLLimitError( 

110 "TOML key or table header has more than " 

111 f"{MAX_TOML_KEY_PARTS} dotted parts" 

112 ) 

113 return tomllib.loads(text)