Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_ppb / utils / check_rules.py: 78%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15 

16"""Defensive YAML loading helpers (caps nesting depth to guard against malicious input).""" 

17from typing import Union 

18 

19import yaml 

20from yaml.nodes import MappingNode, Node 

21 

22YAML_MAX_NESTING_DEPTH = 10 

23 

24 

25def _get_yaml_ast_depth(node: Node, depth: int = 0) -> int: 

26 """Recursively return the maximum nesting depth of a YAML AST.""" 

27 if isinstance(node, MappingNode): 

28 return max( 

29 (_get_yaml_ast_depth(v, depth + 1) for _, v in node.value), default=depth 

30 ) 

31 return depth 

32 

33 

34def check_yaml_depth_before_loading(yaml_str: Union[str, bytes], 

35 max_depth: int = YAML_MAX_NESTING_DEPTH) -> None: 

36 """Reject YAML documents whose nesting depth exceeds ``max_depth``. 

37 

38 Args: 

39 yaml_str: Raw YAML text. 

40 max_depth: Maximum permitted nesting depth. 

41 

42 Raises: 

43 ValueError: If the document exceeds ``max_depth`` or fails to parse. 

44 """ 

45 try: 

46 node = yaml.compose(yaml_str) 

47 if node is None: 

48 return 

49 depth = _get_yaml_ast_depth(node) 

50 if depth > max_depth: 

51 raise ValueError( 

52 f"YAML nesting depth {depth} exceeds the maximum allowed value of {max_depth}" 

53 ) 

54 except yaml.YAMLError as e: 

55 raise ValueError(f"YAML parse error: {e}") from e