Merge pull request #54 from vrahane/fix_tlv_parsing

img_mgmt: TLVs were not getting parsed correctly
diff --git a/cmd/img_mgmt/include/img_mgmt/image.h b/cmd/img_mgmt/include/img_mgmt/image.h
index d390783..53ca7f6 100644
--- a/cmd/img_mgmt/include/img_mgmt/image.h
+++ b/cmd/img_mgmt/include/img_mgmt/image.h
@@ -28,6 +28,7 @@
 
 #define IMAGE_MAGIC                 0x96f3b83d
 #define IMAGE_TLV_INFO_MAGIC        0x6907
+#define IMAGE_TLV_PROT_INFO_MAGIC   0x6908
 
 #define IMAGE_HEADER_SIZE           32
 
diff --git a/cmd/img_mgmt/src/img_mgmt.c b/cmd/img_mgmt/src/img_mgmt.c
index edaa1f8..b994cae 100644
--- a/cmd/img_mgmt/src/img_mgmt.c
+++ b/cmd/img_mgmt/src/img_mgmt.c
@@ -78,8 +78,8 @@
  * Finds the TLVs in the specified image slot, if any.
  */
 static int
-img_mgmt_find_tlvs(const struct image_header *hdr,
-                   int slot, size_t *start_off, size_t *end_off)
+img_mgmt_find_tlvs(int slot, size_t *start_off, size_t *end_off,
+                   uint16_t magic)
 {
     struct image_tlv_info tlv_info;
     int rc;
@@ -90,7 +90,7 @@
         return MGMT_ERR_EUNKNOWN;
     }
 
-    if (tlv_info.it_magic != IMAGE_TLV_INFO_MAGIC) {
+    if (tlv_info.it_magic != magic) {
         /* No TLVs. */
         return MGMT_ERR_ENOENT;
     }
@@ -161,11 +161,22 @@
         *flags = hdr.ih_flags;
     }
 
-    /* Read the image's TLVs.  All images are required to have a hash TLV.  If
-     * the hash is missing, the image is considered invalid.
+    /* Read the image's TLVs. We first try to find the protected TLVs, if the protected
+     * TLV does not exist, we try to find non-protected TLV which also contains the hash
+     * TLV. All images are required to have a hash TLV.  If the hash is missing, the image
+     * is considered invalid.
      */
     data_off = hdr.ih_hdr_size + hdr.ih_img_size;
-    rc = img_mgmt_find_tlvs(&hdr, image_slot, &data_off, &data_end);
+
+    rc = img_mgmt_find_tlvs(image_slot, &data_off, &data_end, IMAGE_TLV_PROT_INFO_MAGIC);
+    if (!rc) {
+        /* The data offset should start after the header bytes after the end of the protected TLV,
+         * if one exists.
+         */
+        data_off = data_end - sizeof(struct image_tlv_info);
+    }
+
+    rc = img_mgmt_find_tlvs(image_slot, &data_off, &data_end, IMAGE_TLV_INFO_MAGIC);
     if (rc != 0) {
         return MGMT_ERR_EUNKNOWN;
     }