@@ -544,16 +544,22 @@ def verify(imgfile, key):
544544 if magic != IMAGE_MAGIC :
545545 return VerifyResult .INVALID_MAGIC , None , None
546546
547- tlv_info = b [header_size + img_size :header_size + img_size + TLV_INFO_SIZE ]
547+ tlv_off = header_size + img_size
548+ tlv_info = b [tlv_off :tlv_off + TLV_INFO_SIZE ]
548549 magic , tlv_tot = struct .unpack ('HH' , tlv_info )
550+ if magic == TLV_PROT_INFO_MAGIC :
551+ tlv_off += tlv_tot
552+ tlv_info = b [tlv_off :tlv_off + TLV_INFO_SIZE ]
553+ magic , tlv_tot = struct .unpack ('HH' , tlv_info )
554+
549555 if magic != TLV_INFO_MAGIC :
550556 return VerifyResult .INVALID_TLV_INFO_MAGIC , None , None
551557
552558 sha = hashlib .sha256 ()
553- sha .update (b [:header_size + img_size ])
559+ prot_tlv_size = tlv_off
560+ sha .update (b [:prot_tlv_size ])
554561 digest = sha .digest ()
555562
556- tlv_off = header_size + img_size
557563 tlv_end = tlv_off + tlv_tot
558564 tlv_off += TLV_INFO_SIZE # skip tlv info
559565 while tlv_off < tlv_end :
@@ -569,7 +575,7 @@ def verify(imgfile, key):
569575 elif key is not None and tlv_type == TLV_VALUES [key .sig_tlv ()]:
570576 off = tlv_off + TLV_SIZE
571577 tlv_sig = b [off :off + tlv_len ]
572- payload = b [:header_size + img_size ]
578+ payload = b [:prot_tlv_size ]
573579 try :
574580 if hasattr (key , 'verify' ):
575581 key .verify (tlv_sig , payload )
0 commit comments