From a40115249bffb1e6886121186c45a6ea41a6e5ec Mon Sep 17 00:00:00 2001
From: Erik Johnson <erjohnso@akamai.com>
Date: Fri, 26 Apr 2024 12:19:15 -0500
Subject: [PATCH] Don't use Salt's custom YAML loader to load static grains
 file

---
 salt/grains/extra.py                      |  5 ++--
 salt/modules/grains.py                    |  7 ++++--
 tests/pytests/unit/grains/test_extra.py   | 30 +++++++++++++++++++++++
 tests/pytests/unit/modules/test_grains.py | 29 ++++++++++++++++++++++
 4 files changed, 67 insertions(+), 4 deletions(-)
 create mode 100644 tests/pytests/unit/grains/test_extra.py

diff --git a/salt/grains/extra.py b/salt/grains/extra.py
index b030d6fc869d..3947590e2b61 100644
--- a/salt/grains/extra.py
+++ b/salt/grains/extra.py
@@ -2,12 +2,13 @@
 import logging
 import os
 
+import yaml
+
 import salt.utils
 import salt.utils.data
 import salt.utils.files
 import salt.utils.path
 import salt.utils.platform
-import salt.utils.yaml
 
 __proxyenabled__ = ["*"]
 log = logging.getLogger(__name__)
@@ -56,7 +57,7 @@ def config():
         log.debug("Loading static grains from %s", gfn)
         with salt.utils.files.fopen(gfn, "rb") as fp_:
             try:
-                return salt.utils.data.decode(salt.utils.yaml.safe_load(fp_))
+                return salt.utils.data.decode(yaml.safe_load(fp_))
             except Exception:  # pylint: disable=broad-except
                 log.warning("Bad syntax in grains file! Skipping.")
                 return {}
diff --git a/salt/modules/grains.py b/salt/modules/grains.py
index c9f9d1481d53..d623c5c10475 100644
--- a/salt/modules/grains.py
+++ b/salt/modules/grains.py
@@ -17,6 +17,8 @@
 from collections.abc import Mapping
 from functools import reduce  # pylint: disable=redefined-builtin
 
+import yaml
+
 import salt.utils.compat
 import salt.utils.data
 import salt.utils.files
@@ -262,8 +264,9 @@ def setvals(grains, destructive=False, refresh_pillar=True):
     if os.path.isfile(gfn):
         with salt.utils.files.fopen(gfn, "rb") as fp_:
             try:
-                grains = salt.utils.yaml.safe_load(fp_)
-            except salt.utils.yaml.YAMLError as exc:
+                # DO NOT USE salt.utils.yaml.safe_load HERE
+                grains = yaml.safe_load(fp_)
+            except Exception as exc:  # pylint: disable=broad-except
                 return f"Unable to read existing grains file: {exc}"
         if not isinstance(grains, dict):
             grains = {}
diff --git a/tests/pytests/unit/grains/test_extra.py b/tests/pytests/unit/grains/test_extra.py
new file mode 100644
index 000000000000..8ab88a00eecf
--- /dev/null
+++ b/tests/pytests/unit/grains/test_extra.py
@@ -0,0 +1,30 @@
+"""
+tests.pytests.unit.grains.test_extra
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+"""
+
+import pytest
+
+import salt.grains.extra as extra
+import salt.utils.platform
+from tests.support.mock import MagicMock, patch
+
+
+@pytest.fixture
+def configure_loader_modules():
+    return {extra: {}}
+
+
+def test_static_grains_file_conflicting_keys(tmp_path):
+    """
+    Test that static grains files with conflicting keys don't result in failure
+    to load all grains from that file.
+    """
+    with (tmp_path / "grains").open("w") as fh:
+        fh.write('foo: bar\nfoo: baz\nno_conflict_here: "yay"\n')
+
+    with patch.object(
+        salt.utils.platform, "is_proxy", MagicMock(return_value=False)
+    ), patch("salt.grains.extra.__opts__", {"conf_file": str(tmp_path / "minion")}):
+        static_grains = extra.config()
+        assert "no_conflict_here" in static_grains
diff --git a/tests/pytests/unit/modules/test_grains.py b/tests/pytests/unit/modules/test_grains.py
index 24a78ee4a2d2..133dff55ddcb 100644
--- a/tests/pytests/unit/modules/test_grains.py
+++ b/tests/pytests/unit/modules/test_grains.py
@@ -2,6 +2,7 @@
 import os
 
 import pytest
+import yaml
 
 import salt.modules.grains as grainsmod
 import salt.utils.dictupdate as dictupdate
@@ -689,6 +690,34 @@ def test_setval_unicode():
         assert ret[key] == value
 
 
+def test_setvals_conflicting_ids(tmp_path):
+    """
+    Test that conflicting top-level keys in static grains file does not break
+    managing static grains
+    """
+    static_grains = tmp_path / "grains"
+    with static_grains.open("w") as fh:
+        fh.write('foo: bar\nfoo: baz\nno_conflict_here: "yay"\n')
+
+    key = "hello"
+    value = "world"
+
+    with patch.dict(grainsmod.__opts__, {"conf_file": str(tmp_path / "minion")}):
+        # Update the grains file in the temp dir
+        grainsmod.setvals({key: value})
+
+    # Load the contents of the file as text for the below asserts
+    updated_grains = static_grains.read_text()
+
+    top_level_keys = [
+        line.split()[0].rstrip(":") for line in updated_grains.splitlines()
+    ]
+    # Confirm that the conflicting IDs were collapsed into a single value
+    assert top_level_keys.count("foo") == 1
+    # Confirm that the new static grain made it into the file
+    assert yaml.safe_load(updated_grains)[key] == value
+
+
 def test_delval_single():
     with patch.dict(
         grainsmod.__grains__, {"a": "aval", "b": {"nested": "val"}, "c": 8}