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 <[email protected]>

scripts/imgtool/image.py

@@ -94,6 +94,7 @@
         'DECOMP_SIZE': 0x70,
         'DECOMP_SHA': 0x71,
         'DECOMP_SIGNATURE': 0x72,
+        'COMP_DEC_SIZE' : 0x73,
 }
 
 TLV_SIZE = 4
@@ -264,6 +265,7 @@ def __init__(self, version=None, header_size=IMAGE_HEADER_SIZE,
         if load_addr and rom_fixed:
             raise click.UsageError("Can not set rom_fixed and load_addr at the same time")
 
+        self.img_only_size = None
         self.image_hash = None
         self.image_size = None
         self.signature = None
@@ -362,6 +364,7 @@ def load_compressed(self, data, compression_header):
         """Load an image from buffer"""
         self.payload = compression_header + data
         self.image_size = len(self.payload)
+        self.img_only_size = len(data)
 
         # Add the image header if needed.
         if self.pad_header and self.header_size > 0:
@@ -459,7 +462,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):
         self.enckey = enckey
 
         # key decides on sha, then pub_key; of both are none default is used
@@ -593,6 +596,12 @@ def create(self, key, public_key_format, enckey, dependencies=None,
                     prot_tlv.add(tag, value)
 
             protected_tlv_off = len(self.payload)
+            if keep_comp_size:
+                if not self.img_only_size:
+                    sys.exit("Error. Compressed image size should be known at this point.");
+                prot_tlv.add("COMP_DEC_SIZE", struct.pack(
+                    self.get_struct_endian() + 'L', self.img_only_size))
+
             self.payload += prot_tlv.get()
 
         tlv = TLV(self.endian)

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, None, 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"]:
         compressed_img = image.Image(version=decode_version(version),
                   header_size=header_size, pad_header=pad_header,
                   pad=pad, confirm=confirm, align=int(align),
@@ -562,11 +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)
+               pub_key, vector_to_sign, user_sha=user_sha,
+               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()