hwid: Patch YAML mapping constructor to fail on duplicated keys.

The default PyYAML's mapping constructor does not warn nor raise
exception when there are duplicated keys in a mapping.

BUG=chrome-os-partner:37618
TEST=unit tests

Change-Id: I7d603eca6b77a485657d4e757146d553c6eaf947
Signed-off-by: Ricky Liang <jcliang@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/275870
diff --git a/py/hwid/database.py b/py/hwid/database.py
index c194362..ab1f655 100644
--- a/py/hwid/database.py
+++ b/py/hwid/database.py
@@ -26,6 +26,37 @@
 from cros.factory.utils import file_utils
 
 
+def PatchYAMLMappingConstructor():
+  """Patch the mapping constructor of PyYAML to fail on duplicated keys."""
+  def ConstructMapping(self, node, deep=False):
+    if not isinstance(node, yaml.nodes.MappingNode):
+      raise yaml.constructor.ConstructorError(
+          None, None, 'expected a mapping node, but found %s' % node.id,
+          node.start_mark)
+    mapping = {}
+    for key_node, value_node in node.value:
+      key = self.construct_object(key_node, deep=deep)
+      try:
+        hash(key)
+      except TypeError, exc:
+        raise yaml.constructor.ConstructorError(
+            'while constructing a mapping', node.start_mark,
+            'found unacceptable key (%s)' % exc, key_node.start_mark)
+      value = self.construct_object(value_node, deep=deep)
+      if key in mapping:
+        raise yaml.constructor.ConstructorError(
+            'while constructing a mapping', node.start_mark,
+            'found duplicated key (%s)' % key, key_node.start_mark)
+      mapping[key] = value
+    return mapping
+
+  yaml.add_constructor(yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG,
+                       ConstructMapping)
+
+
+PatchYAMLMappingConstructor()
+
+
 class Database(object):
   """A class for reading in, parsing, and obtaining information of the given
   device-specific component database.
diff --git a/py/hwid/database_unittest.py b/py/hwid/database_unittest.py
index 1143f21..faa760f 100755
--- a/py/hwid/database_unittest.py
+++ b/py/hwid/database_unittest.py
@@ -37,6 +37,11 @@
         Database.LoadFile,
         os.path.join(_TEST_DATA_PATH, 'test_db_wrong_checksum_field.yaml'),
         verify_checksum=True)
+    self.assertRaisesRegexp(
+        yaml.constructor.ConstructorError,
+        r'while constructing a mapping\n.*\nfound duplicated key \(1\)',
+        Database.LoadFile,
+        os.path.join(_TEST_DATA_PATH, 'test_db_duplicated_keys.yaml'))
 
   def testDatabaseChecksum(self):
     self.assertEquals(
diff --git a/py/hwid/testdata/test_db_duplicated_keys.yaml b/py/hwid/testdata/test_db_duplicated_keys.yaml
new file mode 100644
index 0000000..7aea700
--- /dev/null
+++ b/py/hwid/testdata/test_db_duplicated_keys.yaml
@@ -0,0 +1,36 @@
+# Dummy line before checksum
+checksum: c83402d5afa4a70f16cba436b61c4e4ad52ff138
+# Dummy line after checksum
+board: CHROMEBOOK
+
+encoding_patterns:
+  0: default
+
+image_id:
+  0: MP
+
+pattern:
+  - image_ids: [0]
+    encoding_scheme: base8192
+    fields:
+    - cpu_field: 2
+
+encoded_fields:
+  cpu_field:
+    0: { cpu: fast_cpu }
+    1: { cpu: really_fast_cpu }
+    1: { cpu: super_fast_cpu }
+
+components:
+  cpu:
+    items:
+      fast_cpu:
+        values: { compact_str: fast CPU }
+      really_fast_cpu:
+        values: { compact_str: really fast CPU }
+      super_fast_cpu:
+        values: { compact_str: super fast CPU }
+
+rules:
+- name: device_info.set_image_id
+  evaluate: SetImageId('MP')