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.
This commit is contained in:
George Tsiamasiotis 2024-09-17 20:57:59 +03:00
parent 53e92a7cda
commit 532bbc8f61
3 changed files with 28 additions and 27 deletions

View File

@ -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)

View File

@ -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):

View File

@ -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(
"Key redefinitions are not allowed: " + ", ".join(duplicate_keys)
)
return mapping
def clean_camera_user_pass(line: str) -> str: