From 38ff46e45cdaa686bb1ef94e4563856b9b898ebe Mon Sep 17 00:00:00 2001 From: gtsiam Date: Tue, 17 Sep 2024 23:52:55 +0300 Subject: [PATCH] Rewrite yaml loader (#13803) * Ignore entire __pycache__ folder instead of individual *.pyc files * Rewrite the yaml loader to match PyYAML The old implementation would fail in weird ways with configs that were incorrect in just the right way. The new implementation just does what PyYAML would do, only diverging in case of duplicate keys. * Clarify duplicate yaml key ValueError message Co-authored-by: Nicolas Mowen --------- Co-authored-by: Nicolas Mowen --- .gitignore | 2 +- frigate/config.py | 7 +++--- frigate/test/test_config.py | 5 +++-- frigate/util/builtin.py | 43 ++++++++++++++++++------------------- 4 files changed, 29 insertions(+), 28 deletions(-) diff --git a/.gitignore b/.gitignore index 195708e2d..11a147a6e 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,5 @@ .DS_Store -*.pyc +__pycache__ *.swp debug .vscode/* diff --git a/frigate/config.py b/frigate/config.py index ad96cebef..af3ed0753 100644 --- a/frigate/config.py +++ b/frigate/config.py @@ -9,6 +9,7 @@ from pathlib import Path from typing import Any, Dict, List, Optional, Tuple, Union import numpy as np +import yaml from pydantic import ( BaseModel, ConfigDict, @@ -41,11 +42,11 @@ from frigate.ffmpeg_presets import ( ) from frigate.plus import PlusApi from frigate.util.builtin import ( + NoDuplicateKeysLoader, deep_merge, escape_special_characters, generate_color_palette, get_ffmpeg_arg_list, - load_config_with_no_duplicates, ) from frigate.util.config import StreamInfoRetriever, get_relative_coordinates from frigate.util.image import create_mask @@ -1764,7 +1765,7 @@ class FrigateConfig(FrigateBaseModel): raw_config = f.read() if config_file.endswith(YAML_EXT): - config = load_config_with_no_duplicates(raw_config) + config = yaml.load(raw_config, NoDuplicateKeysLoader) elif config_file.endswith(".json"): config = json.loads(raw_config) @@ -1772,5 +1773,5 @@ class FrigateConfig(FrigateBaseModel): @classmethod def parse_raw(cls, raw_config): - config = load_config_with_no_duplicates(raw_config) + config = yaml.load(raw_config, NoDuplicateKeysLoader) return cls.model_validate(config) diff --git a/frigate/test/test_config.py b/frigate/test/test_config.py index c703de893..1f0ff086d 100644 --- a/frigate/test/test_config.py +++ b/frigate/test/test_config.py @@ -4,13 +4,14 @@ import unittest from unittest.mock import patch import numpy as np +import yaml from pydantic import ValidationError from frigate.config import BirdseyeModeEnum, FrigateConfig from frigate.const import MODEL_CACHE_DIR from frigate.detectors import DetectorTypeEnum from frigate.plus import PlusApi -from frigate.util.builtin import deep_merge, load_config_with_no_duplicates +from frigate.util.builtin import NoDuplicateKeysLoader, deep_merge class TestConfig(unittest.TestCase): @@ -1537,7 +1538,7 @@ class TestConfig(unittest.TestCase): """ self.assertRaises( - ValueError, lambda: load_config_with_no_duplicates(raw_config) + ValueError, lambda: yaml.load(raw_config, NoDuplicateKeysLoader) ) def test_object_filter_ratios_work(self): diff --git a/frigate/util/builtin.py b/frigate/util/builtin.py index 2c3051fc0..413c8d8ce 100644 --- a/frigate/util/builtin.py +++ b/frigate/util/builtin.py @@ -89,32 +89,31 @@ def deep_merge(dct1: dict, dct2: dict, override=False, merge_lists=False) -> dic return merged -def load_config_with_no_duplicates(raw_config) -> dict: - """Get config ensuring duplicate keys are not allowed.""" +class NoDuplicateKeysLoader(yaml.loader.SafeLoader): + """A yaml SafeLoader that disallows duplicate keys""" - # https://stackoverflow.com/a/71751051 - # important to use SafeLoader here to avoid RCE - class PreserveDuplicatesLoader(yaml.loader.SafeLoader): - pass + def construct_mapping(self, node, deep=False): + mapping = super().construct_mapping(node, deep=deep) - def map_constructor(loader, node, deep=False): - keys = [loader.construct_object(node, deep=deep) for node, _ in node.value] - vals = [loader.construct_object(node, deep=deep) for _, node in node.value] - key_count = Counter(keys) - data = {} - for key, val in zip(keys, vals): - if key_count[key] > 1: - raise ValueError( - f"Config input {key} is defined multiple times for the same field, this is not allowed." + if len(node.value) != len(mapping): + # There's a duplicate key somewhere. Find it. + duplicate_keys = [ + key + for key, count in Counter( + self.construct_object(key, deep=deep) for key, _ in node.value ) - else: - data[key] = val - return data + if count > 1 + ] - PreserveDuplicatesLoader.add_constructor( - yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, map_constructor - ) - return yaml.load(raw_config, PreserveDuplicatesLoader) + # This might be possible if PyYAML's construct_mapping() changes the node + # afterwards for some reason? I don't see why, but better safe than sorry. + assert len(duplicate_keys) > 0 + + raise ValueError( + f"Config field duplicates are not allowed, the following fields are duplicated in the config: {', '.join(duplicate_keys)}" + ) + + return mapping def clean_camera_user_pass(line: str) -> str: