From ff9e1da1deffb0b5718e49893e9cb9392080f509 Mon Sep 17 00:00:00 2001
From: Josh Hawkins <32435876+hawkeye217@users.noreply.github.com>
Date: Tue, 17 Sep 2024 16:03:22 -0500
Subject: [PATCH] Revert "Rewrite yaml loader (#13803)" (#13805)

This reverts commit 38ff46e45cdaa686bb1ef94e4563856b9b898ebe.
---
 .gitignore                  |  2 +-
 frigate/config.py           |  7 +++---
 frigate/test/test_config.py |  5 ++---
 frigate/util/builtin.py     | 43 +++++++++++++++++++------------------
 4 files changed, 28 insertions(+), 29 deletions(-)

diff --git a/.gitignore b/.gitignore
index 11a147a6e..195708e2d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,5 +1,5 @@
 .DS_Store
-__pycache__
+*.pyc
 *.swp
 debug
 .vscode/*
diff --git a/frigate/config.py b/frigate/config.py
index af3ed0753..ad96cebef 100644
--- a/frigate/config.py
+++ b/frigate/config.py
@@ -9,7 +9,6 @@ from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import numpy as np
-import yaml
 from pydantic import (
     BaseModel,
     ConfigDict,
@@ -42,11 +41,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
@@ -1765,7 +1764,7 @@ class FrigateConfig(FrigateBaseModel):
             raw_config = f.read()
 
         if config_file.endswith(YAML_EXT):
-            config = yaml.load(raw_config, NoDuplicateKeysLoader)
+            config = load_config_with_no_duplicates(raw_config)
         elif config_file.endswith(".json"):
             config = json.loads(raw_config)
 
@@ -1773,5 +1772,5 @@ class FrigateConfig(FrigateBaseModel):
 
     @classmethod
     def parse_raw(cls, raw_config):
-        config = yaml.load(raw_config, NoDuplicateKeysLoader)
+        config = load_config_with_no_duplicates(raw_config)
         return cls.model_validate(config)
diff --git a/frigate/test/test_config.py b/frigate/test/test_config.py
index 1f0ff086d..c703de893 100644
--- a/frigate/test/test_config.py
+++ b/frigate/test/test_config.py
@@ -4,14 +4,13 @@ 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 NoDuplicateKeysLoader, deep_merge
+from frigate.util.builtin import deep_merge, load_config_with_no_duplicates
 
 
 class TestConfig(unittest.TestCase):
@@ -1538,7 +1537,7 @@ class TestConfig(unittest.TestCase):
         """
 
         self.assertRaises(
-            ValueError, lambda: yaml.load(raw_config, NoDuplicateKeysLoader)
+            ValueError, lambda: load_config_with_no_duplicates(raw_config)
         )
 
     def test_object_filter_ratios_work(self):
diff --git a/frigate/util/builtin.py b/frigate/util/builtin.py
index 413c8d8ce..2c3051fc0 100644
--- a/frigate/util/builtin.py
+++ b/frigate/util/builtin.py
@@ -89,31 +89,32 @@ def deep_merge(dct1: dict, dct2: dict, override=False, merge_lists=False) -> dic
     return merged
 
 
-class NoDuplicateKeysLoader(yaml.loader.SafeLoader):
-    """A yaml SafeLoader that disallows duplicate keys"""
+def load_config_with_no_duplicates(raw_config) -> dict:
+    """Get config ensuring duplicate keys are not allowed."""
 
-    def construct_mapping(self, node, deep=False):
-        mapping = super().construct_mapping(node, deep=deep)
+    # https://stackoverflow.com/a/71751051
+    # important to use SafeLoader here to avoid RCE
+    class PreserveDuplicatesLoader(yaml.loader.SafeLoader):
+        pass
 
-        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
+    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 count > 1
-            ]
+            else:
+                data[key] = val
+        return data
 
-            # 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
+    PreserveDuplicatesLoader.add_constructor(
+        yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, map_constructor
+    )
+    return yaml.load(raw_config, PreserveDuplicatesLoader)
 
 
 def clean_camera_user_pass(line: str) -> str: