From 13c4bb72b1b05758181705507343467975ab2cff Mon Sep 17 00:00:00 2001 From: Mateusz Michalek Date: Thu, 2 Jan 2025 14:48:28 +0100 Subject: [PATCH] scripts: imgtool: fix compression with encryption adds TLV which stores compressed image size in case encryption is on. This is to avoid wrong streaam size due to encryption padding. Signed-off-by: Mateusz Michalek --- scripts/imgtool/image.py | 11 ++++++++--- scripts/imgtool/main.py | 17 ++++++++++++----- 2 files changed, 20 insertions(+), 8 deletions(-) diff --git a/scripts/imgtool/image.py b/scripts/imgtool/image.py index 88ef7b90c9..3e2c11073f 100644 --- a/scripts/imgtool/image.py +++ b/scripts/imgtool/image.py @@ -94,6 +94,7 @@ 'DECOMP_SIZE': 0x70, 'DECOMP_SHA': 0x71, 'DECOMP_SIGNATURE': 0x72, + 'COMP_DEC_SIZE' : 0x73, } TLV_SIZE = 4 @@ -460,7 +461,7 @@ def create(self, key, public_key_format, enckey, dependencies=None, sw_type=None, custom_tlvs=None, compression_tlvs=None, compression_type=None, encrypt_keylen=128, clear=False, fixed_sig=None, pub_key=None, vector_to_sign=None, - user_sha='auto', is_pure=False): + user_sha='auto', is_pure=False, keep_comp_size=False, dont_encrypt=False): self.enckey = enckey # key decides on sha, then pub_key; of both are none default is used @@ -522,6 +523,9 @@ def create(self, key, public_key_format, enckey, dependencies=None, dependencies_num = len(dependencies[DEP_IMAGES_KEY]) protected_tlv_size += (dependencies_num * 16) + if keep_comp_size: + compression_tlvs["COMP_DEC_SIZE"] = struct.pack( + self.get_struct_endian() + 'L', self.image_size) if compression_tlvs is not None: for value in compression_tlvs.values(): protected_tlv_size += TLV_SIZE + len(value) @@ -537,7 +541,7 @@ def create(self, key, public_key_format, enckey, dependencies=None, # # This adds the padding if image is not aligned to the 16 Bytes # in encrypted mode - if self.enckey is not None: + if self.enckey is not None and dont_encrypt is False: pad_len = len(self.payload) % 16 if pad_len > 0: pad = bytes(16 - pad_len) @@ -594,6 +598,7 @@ def create(self, key, public_key_format, enckey, dependencies=None, prot_tlv.add(tag, value) protected_tlv_off = len(self.payload) + self.payload += prot_tlv.get() tlv = TLV(self.endian) @@ -657,7 +662,7 @@ def create(self, key, public_key_format, enckey, dependencies=None, if protected_tlv_off is not None: self.payload = self.payload[:protected_tlv_off] - if enckey is not None: + if enckey is not None and dont_encrypt is False: if encrypt_keylen == 256: plainkey = os.urandom(32) else: diff --git a/scripts/imgtool/main.py b/scripts/imgtool/main.py index 434530c7a4..28caa89478 100755 --- a/scripts/imgtool/main.py +++ b/scripts/imgtool/main.py @@ -521,12 +521,11 @@ def sign(key, public_key_format, align, version, pad_sig, header_size, 'Pure signatures, currently, enforces preferred hash algorithm, ' 'and forbids sha selection by user.') - img.create(key, public_key_format, enckey, dependencies, boot_record, + if compression in ["lzma2", "lzma2armthumb"]: + img.create(key, public_key_format, enckey, dependencies, boot_record, custom_tlvs, compression_tlvs, None, int(encrypt_keylen), clear, baked_signature, pub_key, vector_to_sign, user_sha=user_sha, - is_pure=is_pure) - - if compression in ["lzma2", "lzma2armthumb"]: + is_pure=is_pure, keep_comp_size=False, dont_encrypt=True) compressed_img = image.Image(version=decode_version(version), header_size=header_size, pad_header=pad_header, pad=pad, confirm=confirm, align=int(align), @@ -562,12 +561,20 @@ def sign(key, public_key_format, align, version, pad_sig, header_size, lc = comp_default_lc, lp = comp_default_lp) compressed_img.load_compressed(compressed_data, compression_header) compressed_img.base_addr = img.base_addr + keep_comp_size = False; + if enckey: + keep_comp_size = True compressed_img.create(key, public_key_format, enckey, dependencies, boot_record, custom_tlvs, compression_tlvs, compression, int(encrypt_keylen), clear, baked_signature, pub_key, vector_to_sign, user_sha=user_sha, - is_pure=is_pure) + is_pure=is_pure, keep_comp_size=keep_comp_size) img = compressed_img + else: + img.create(key, public_key_format, enckey, dependencies, boot_record, + custom_tlvs, compression_tlvs, None, int(encrypt_keylen), clear, + baked_signature, pub_key, vector_to_sign, user_sha=user_sha, + is_pure=is_pure) img.save(outfile, hex_addr) if sig_out is not None: new_signature = img.get_signature()