From 38ff46e45cdaa686bb1ef94e4563856b9b898ebe Mon Sep 17 00:00:00 2001
From: gtsiam <gtsiam@windowslive.com>
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 <nickmowen213@gmail.com>

---------

Co-authored-by: Nicolas Mowen <nickmowen213@gmail.com>
---
 .gitignore                  |  2 +-
 frigate/config.py           |  7 +++---
 frigate/test/test_config.py |  5 +++--
 frigate/util/builtin.py     | 45 ++++++++++++++++++-------------------
 4 files changed, 30 insertions(+), 29 deletions(-)

diff --git a/.gitignore b/.gitignore
index 195708e2d4..11a147a6e7 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 ad96cebef8..af3ed07536 100644
--- a/frigate/config.py
+++ b/frigate/config.py
@@ -9,6 +9,7 @@
 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.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 @@ def parse_file(cls, config_file):
             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 @@ def parse_file(cls, config_file):
 
     @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 c703de8930..1f0ff086dd 100644
--- a/frigate/test/test_config.py
+++ b/frigate/test/test_config.py
@@ -4,13 +4,14 @@
 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 @@ def test_fails_duplicate_keys(self):
         """
 
         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 2c3051fc06..413c8d8ce2 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"""
+
+    def construct_mapping(self, node, deep=False):
+        mapping = super().construct_mapping(node, deep=deep)
+
+        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
+                )
+                if count > 1
+            ]
 
-    # https://stackoverflow.com/a/71751051
-    # important to use SafeLoader here to avoid RCE
-    class PreserveDuplicatesLoader(yaml.loader.SafeLoader):
-        pass
+            # 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
 
-    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."
-                )
-            else:
-                data[key] = val
-        return data
+            raise ValueError(
+                f"Config field duplicates are not allowed, the following fields are duplicated in the config: {', '.join(duplicate_keys)}"
+            )
 
-    PreserveDuplicatesLoader.add_constructor(
-        yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, map_constructor
-    )
-    return yaml.load(raw_config, PreserveDuplicatesLoader)
+        return mapping
 
 
 def clean_camera_user_pass(line: str) -> str: