Refactor spec_loader to use yaml.load with SafeLoader

CHANGELOG.rst

@@ -4,6 +4,12 @@ Changelog
 in development
 --------------
 
+Fixed
+~~~~~
+
+* Refactor spec_loader util to use yaml.load with SafeLoader. (security)
+  Contributed by @ashwini-orchestral
+
 Changed
 ~~~~~~~
 

st2common/st2common/cmd/validate_api_spec.py

@@ -114,13 +114,10 @@ def main():
     setup()
 
     try:
-        # 1. Validate there are no duplicates keys in the YAML file
-        spec_loader.load_spec(
-            "st2common", "openapi.yaml.j2", allow_duplicate_keys=False
-        )
+        # Validate there are no duplicates keys in the YAML file.
+        # The spec loader do not allow duplicate keys.
+        spec_loader.load_spec("st2common", "openapi.yaml.j2")
 
-        # 2. Validate schema (currently broken)
-        # validate_spec()
         ret = 0
     except Exception:
         LOG.error("Failed to validate openapi.yaml file", exc_info=True)

st2common/st2common/util/spec_loader.py

@@ -19,13 +19,13 @@
 import jinja2
 import yaml
 
-from yaml.constructor import ConstructorError
-from yaml.nodes import MappingNode
-
 try:
-    from yaml import CLoader as Loader
+    from yaml import CSafeLoader as SafeLoader
 except ImportError:
-    from yaml import Loader
+    from yaml import SafeLoader
+
+from yaml import constructor
+from yaml import nodes
 
 import st2common.constants.pack
 import st2common.constants.action
@@ -43,14 +43,12 @@
 }
 
 
-class UniqueKeyLoader(Loader):
-    """
-    YAML loader which throws on a duplicate key.
-    """
-
+# Custom YAML loader that throw an exception on duplicate key.
+# Credit: https://gist.github.com/pypt/94d747fe5180851196eb
+class UniqueKeyLoader(SafeLoader):
     def construct_mapping(self, node, deep=False):
-        if not isinstance(node, MappingNode):
-            raise ConstructorError(
+        if not isinstance(node, nodes.MappingNode):
+            raise constructor.ConstructorError(
                 None,
                 None,
                 "expected a mapping node, but found %s" % node.id,
@@ -62,35 +60,37 @@ def construct_mapping(self, node, deep=False):
             try:
                 hash(key)
             except TypeError as exc:
-                raise ConstructorError(
+                raise constructor.ConstructorError(
                     "while constructing a mapping",
                     node.start_mark,
                     "found unacceptable key (%s)" % exc,
                     key_node.start_mark,
                 )
             # check for duplicate keys
             if key in mapping:
-                raise ConstructorError(
-                    "while constructing a mapping",
-                    node.start_mark,
-                    "found duplicate key",
-                    key_node.start_mark,
+                raise constructor.ConstructorError(
+                    'found duplicate key "%s"' % key_node.value
                 )
             value = self.construct_object(value_node, deep=deep)
             mapping[key] = value
         return mapping
 
 
-def load_spec(module_name, spec_file, allow_duplicate_keys=False):
-    spec_string = generate_spec(module_name, spec_file)
+# Add UniqueKeyLoader to the yaml SafeLoader so it is invoked by safe_load.
+yaml.add_constructor(
+    yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG,
+    UniqueKeyLoader.construct_mapping,
+    Loader=SafeLoader,
+)
 
-    # 1. Check for duplicate keys
-    if not allow_duplicate_keys:
-        yaml.load(spec_string, UniqueKeyLoader)
 
-    # 2. Generate actual spec
-    spec = yaml.safe_load(spec_string)
-    return spec
+def load_spec(module_name, spec_file):
+    spec_string = generate_spec(module_name, spec_file)
+
+    # The use of yaml.load and passing SafeLoader is the same as yaml.safe_load which
+    # makes the same call. The difference here is that we use CSafeLoader where possible
+    # to improve performance and yaml.safe_load uses the python implementation by default.
+    return yaml.load(spec_string, SafeLoader)
 
 
 def generate_spec(module_name, spec_file):

st2common/tests/unit/test_spec_loader.py

@@ -0,0 +1,36 @@
+# Copyright 2020 The StackStorm Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import absolute_import
+
+import unittest2
+import yaml
+
+from st2common.util import spec_loader
+
+
+class SpecLoaderTest(unittest2.TestCase):
+    def test_spec_loader(self):
+        self.assertTrue(
+            isinstance(spec_loader.load_spec("st2common", "openapi.yaml.j2"), dict)
+        )
+
+    def test_bad_spec_duplicate_keys(self):
+        self.assertRaisesRegex(
+            yaml.constructor.ConstructorError,
+            'found duplicate key "swagger"',
+            spec_loader.load_spec,
+            "st2tests.fixtures.specs",
+            "openapi.yaml.j2",
+        )

st2tests/st2tests/fixtures/specs/openapi.yaml.j2

@@ -0,0 +1,6 @@
+swagger: '2.0'
+swagger: '3.0'
+
+info:
+  version: "1.0.0"
+  title: StackStorm API