Add capability to compress dt overlays

For DTBOs using a header of version 1, mkdtboimg.py
supports zlib/gzip compression as per the
compression format information present in the
4 least significant bits of the flag
field in dt_table_entry.

A new argument "decompress" has been added to the
"dump" command to decompress any compressed
entries before dumping.

Bug: 109892148
Test: sh mkdtboimg_testdata.sh
Test: make vts -j32 && vts-tradefed run vts -m
      VtsFirmwareDtboVerification
Change-Id: Id9524f986506ee8eb91d0ed15d4b91a2d2ed85af

4 files changed, 145 insertions, 46 deletions

diff --git a/utils/src/mkdtboimg.py b/utils/src/mkdtboimg.py
index c2dbb08..e52975d 100755
--- a/utils/src/mkdtboimg.py
+++ b/utils/src/mkdtboimg.py
@@ -23,7 +23,14 @@ from array import array
 from collections import namedtuple
 import struct
 from sys import stdout
+import zlib
 
+class CompressionFormat(object):
+    """Enum representing DT compression format for a DT entry.
+    """
+    NO_COMPRESSION = 0x00
+    ZLIB_COMPRESSION = 0x01
+    GZIP_COMPRESSION = 0x02
 
 class DtEntry(object):
     """Provides individual DT image file arguments to be added to a DTBO.
@@ -31,10 +38,13 @@ class DtEntry(object):
     Attributes:
         _REQUIRED_KEYS: 'keys' needed to be present in the dictionary passed to instantiate
             an object of this class.
+        _COMPRESSION_FORMAT_MASK: Mask to retrieve compression info for DT entry from flags field
+            when a DTBO header of version 1 is used.
     """
+    _COMPRESSION_FORMAT_MASK = 0x0f
+    REQUIRED_KEYS = ('dt_file', 'dt_size', 'dt_offset', 'id', 'rev', 'flags',
+                     'custom0', 'custom1', 'custom2')
 
-    REQUIRED_KEYS = ('dt_file', 'dt_size', 'dt_offset', 'id', 'rev', 'custom0',
-                     'custom1', 'custom2', 'custom3')
     @staticmethod
     def __get_number_or_prop(arg):
         """Converts string to integer or reads the property from DT image.
@@ -82,10 +92,10 @@ class DtEntry(object):
         self.__dt_size = kwargs['dt_size']
         self.__id = self.__get_number_or_prop(kwargs['id'])
         self.__rev = self.__get_number_or_prop(kwargs['rev'])
+        self.__flags = self.__get_number_or_prop(kwargs['flags'])
         self.__custom0 = self.__get_number_or_prop(kwargs['custom0'])
         self.__custom1 = self.__get_number_or_prop(kwargs['custom1'])
         self.__custom2 = self.__get_number_or_prop(kwargs['custom2'])
-        self.__custom3 = self.__get_number_or_prop(kwargs['custom3'])
 
     def __str__(self):
         sb = []
@@ -98,15 +108,26 @@ class DtEntry(object):
         sb.append('{key:>20} = {value:08x}'.format(key='rev',
                                                    value=self.__rev))
         sb.append('{key:>20} = {value:08x}'.format(key='custom[0]',
-                                                   value=self.__custom0))
+                                                   value=self.__flags))
         sb.append('{key:>20} = {value:08x}'.format(key='custom[1]',
-                                                   value=self.__custom1))
+                                                   value=self.__custom0))
         sb.append('{key:>20} = {value:08x}'.format(key='custom[2]',
-                                                   value=self.__custom2))
+                                                   value=self.__custom1))
         sb.append('{key:>20} = {value:08x}'.format(key='custom[3]',
-                                                   value=self.__custom3))
+                                                   value=self.__custom2))
         return '\n'.join(sb)
 
+    def compression_info(self, version):
+        """CompressionFormat: compression format for DT image file.
+
+           Args:
+                version: Version of DTBO header, compression is only
+                         supported from version 1.
+        """
+        if version is 0:
+            return CompressionFormat.NO_COMPRESSION
+        return self.flags & self._COMPRESSION_FORMAT_MASK
+
     @property
     def dt_file(self):
         """file: File handle to the DT image file."""
@@ -117,6 +138,10 @@ class DtEntry(object):
         """int: size in bytes of the DT image file."""
         return self.__dt_size
 
+    @size.setter
+    def size(self, value):
+        self.__dt_size = value
+
     @property
     def dt_offset(self):
         """int: offset in DTBO file for this DT image."""
@@ -137,6 +162,11 @@ class DtEntry(object):
         return self.__rev
 
     @property
+    def flags(self):
+        """int: DT entry _flags for this DT image."""
+        return self.__flags
+
+    @property
     def custom0(self):
         """int: DT entry _custom0 for this DT image."""
         return self.__custom0
@@ -148,14 +178,9 @@ class DtEntry(object):
 
     @property
     def custom2(self):
-        """int: DT entry _custom2 for this DT image."""
+        """int: DT entry custom2 for this DT image."""
         return self.__custom2
 
-    @property
-    def custom3(self):
-        """int: DT entry custom3 for this DT image."""
-        return self.__custom3
-
 
 class Dtbo(object):
     """
@@ -168,6 +193,8 @@ class Dtbo(object):
         _DT_TABLE_HEADER_INTS: Number of integers in DT table header.
         _DT_ENTRY_HEADER_SIZE: Size of Device tree entry header within a DTBO.
         _DT_ENTRY_HEADER_INTS: Number of integers in DT entry header.
+        _GZIP_COMPRESSION_WBITS: Argument 'wbits' for gzip compression
+        _ZLIB_DECOMPRESSION_WBITS: Argument 'wbits' for zlib/gzip compression
     """
 
     _DTBO_MAGIC = 0xd7b7ab1e
@@ -175,6 +202,8 @@ class Dtbo(object):
     _DT_TABLE_HEADER_INTS = 8
     _DT_ENTRY_HEADER_SIZE = struct.calcsize('>8I')
     _DT_ENTRY_HEADER_INTS = 8
+    _GZIP_COMPRESSION_WBITS = 31
+    _ZLIB_DECOMPRESSION_WBITS = 47
 
     def _update_dt_table_header(self):
         """Converts header entries into binary data for DTBO header.
@@ -202,8 +231,8 @@ class Dtbo(object):
         """
         struct.pack_into('>8I', self.__metadata, metadata_offset, dt_entry.size,
                          dt_entry.dt_offset, dt_entry.image_id, dt_entry.rev,
-                         dt_entry.custom0, dt_entry.custom1, dt_entry.custom2,
-                         dt_entry.custom3)
+                         dt_entry.flags, dt_entry.custom0, dt_entry.custom1,
+                         dt_entry.custom2)
 
     def _update_metadata(self):
         """Updates the DTBO metadata.
@@ -370,6 +399,40 @@ class Dtbo(object):
         """Returns a list of DtEntry objects found in DTBO file."""
         return self.__dt_entries
 
+    def compress_dt_entry(self, compression_format, dt_entry_file):
+        """Compresses a DT entry.
+
+        Args:
+            compression_format: Compression format for DT Entry
+            dt_entry_file: File handle to read DT entry from.
+
+        Returns:
+            Compressed DT entry and its length.
+
+        Raises:
+            ValueError if unrecognized compression format is found.
+        """
+        compress_zlib = zlib.compressobj()  #  zlib
+        compress_gzip = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION,
+                                         zlib.DEFLATED, self._GZIP_COMPRESSION_WBITS)  #  gzip
+        compression_obj_dict = {
+            CompressionFormat.NO_COMPRESSION: None,
+            CompressionFormat.ZLIB_COMPRESSION: compress_zlib,
+            CompressionFormat.GZIP_COMPRESSION: compress_gzip,
+        }
+
+        if compression_format not in compression_obj_dict:
+            ValueError("Bad compression format %d" % compression_format)
+
+        if compression_format is CompressionFormat.NO_COMPRESSION:
+            dt_entry = dt_entry_file.read()
+        else:
+            compression_object = compression_obj_dict[compression_format]
+            dt_entry_file.seek(0)
+            dt_entry = compression_object.compress(dt_entry_file.read())
+            dt_entry += compression_object.flush()
+        return dt_entry, len(dt_entry)
+
     def add_dt_entries(self, dt_entries):
         """Adds DT image files to the DTBO object.
 
@@ -379,6 +442,12 @@ class Dtbo(object):
         Args:
             dt_entries: List of DtEntry object to be added.
 
+        Returns:
+            A buffer containing all DT entries.
+
+        Raises:
+            ValueError: if the list of DT entries is empty or if a list of DT entries
+                has already been added to the DTBO.
         """
         if not dt_entries:
             raise ValueError('Attempted to add empty list of DT entries')
@@ -390,14 +459,21 @@ class Dtbo(object):
         dt_offset = (self.header_size +
                      dt_entry_count * self.dt_entry_size)
 
+        dt_entry_buf = ""
         for dt_entry in dt_entries:
             if not isinstance(dt_entry, DtEntry):
                 raise ValueError('Adding invalid DT entry object to DTBO')
             entry = self._find_dt_entry_with_same_file(dt_entry)
-            if entry:
+            dt_entry_compression_info = dt_entry.compression_info(self.version)
+            if entry and (entry.compression_info(self.version)
+                          == dt_entry_compression_info):
                 dt_entry.dt_offset = entry.dt_offset
+                dt_entry.size = entry.size
             else:
                 dt_entry.dt_offset = dt_offset
+                compressed_entry, dt_entry.size = self.compress_dt_entry(dt_entry_compression_info,
+                                                                         dt_entry.dt_file)
+                dt_entry_buf += compressed_entry
                 dt_offset += dt_entry.size
                 self.total_size += dt_entry.size
             self.__dt_entries.append(dt_entry)
@@ -405,7 +481,9 @@ class Dtbo(object):
             self.__metadata_size += self.dt_entry_size
             self.total_size += self.dt_entry_size
 
-    def extract_dt_file(self, idx, fout):
+        return dt_entry_buf
+
+    def extract_dt_file(self, idx, fout, decompress):
         """Extract DT Image files embedded in the DTBO file.
 
         Extracts Device Tree blob image file at given index into a file handle.
@@ -413,6 +491,11 @@ class Dtbo(object):
         Args:
             idx: Index of the DT entry in the DTBO file.
             fout: File handle where the DTB at index idx to be extracted into.
+            decompress: If a DT entry is compressed, decompress it before writing
+                it to the file handle.
+
+        Raises:
+            ValueError: if invalid DT entry index or compression format is detected.
         """
         if idx > self.dt_entry_count:
             raise ValueError('Invalid index %d of DtEntry' % idx)
@@ -421,15 +504,26 @@ class Dtbo(object):
         offset = self.dt_entries[idx].dt_offset
         self.__file.seek(offset, 0)
         fout.seek(0)
-        fout.write(self.__file.read(size))
+        compression_format = self.dt_entries[idx].compression_info(self.version)
+        if decompress and compression_format:
+            if (compression_format == CompressionFormat.ZLIB_COMPRESSION or
+                compression_format == CompressionFormat.GZIP_COMPRESSION):
+                fout.write(zlib.decompress(self.__file.read(size), self._ZLIB_DECOMPRESSION_WBITS))
+            else:
+                raise ValueError("Unknown compression format detected")
+        else:
+            fout.write(self.__file.read(size))
 
-    def commit(self):
+    def commit(self, dt_entry_buf):
         """Write out staged changes to the DTBO object to create a DTBO file.
 
         Writes a fully instantiated Dtbo Object into the output file using the
         file handle present in '_file'. No checks are performed on the object
         except for existence of output file handle on the object before writing
         out the file.
+
+        Args:
+            dt_entry_buf: Buffer containing all DT entries.
         """
         if not self.__file:
             raise ValueError('No file given to write to.')
@@ -438,10 +532,10 @@ class Dtbo(object):
             raise ValueError('No DT image files to embed into DTBO image given.')
 
         self._update_metadata()
+
         self.__file.seek(0)
         self.__file.write(self.__metadata)
-        for dt_entry in self.__dt_entries:
-            self.__file.write(dt_entry.dt_file.read())
+        self.__file.write(dt_entry_buf)
         self.__file.flush()
 
 
@@ -468,6 +562,9 @@ def parse_dt_entry(global_args, arglist):
                         default=global_args.global_id)
     parser.add_argument('--rev', type=str, dest='rev',
                         action='store', default=global_args.global_rev)
+    parser.add_argument('--flags', type=str, dest='flags',
+                        action='store',
+                        default=global_args.global_flags)
     parser.add_argument('--custom0', type=str, dest='custom0',
                         action='store',
                         default=global_args.global_custom0)
@@ -477,9 +574,6 @@ def parse_dt_entry(global_args, arglist):
     parser.add_argument('--custom2', type=str, dest='custom2',
                         action='store',
                         default=global_args.global_custom2)
-    parser.add_argument('--custom3', type=str, dest='custom3',
-                        action='store',
-                        default=global_args.global_custom3)
     return parser.parse_args(arglist)
 
 
@@ -640,14 +734,14 @@ def parse_create_args(arg_list):
                         action='store', default='0')
     parser.add_argument('--rev', type=str, dest='global_rev',
                         action='store', default='0')
+    parser.add_argument('--flags', type=str, dest='global_flags',
+                        action='store', default='0')
     parser.add_argument('--custom0', type=str, dest='global_custom0',
                         action='store', default='0')
     parser.add_argument('--custom1', type=str, dest='global_custom1',
                         action='store', default='0')
     parser.add_argument('--custom2', type=str, dest='global_custom2',
                         action='store', default='0')
-    parser.add_argument('--custom3', type=str, dest='global_custom3',
-                        action='store', default='0')
     args = parser.parse_args(argv)
     return args, remainder
 
@@ -669,6 +763,7 @@ def parse_dump_cmd_args(arglist):
                         default=stdout)
     parser.add_argument('--dtb', '-b', nargs='?', type=str,
                         dest='dtfilename')
+    parser.add_argument('--decompress', action='store_true', dest='decompress')
     return parser.parse_args(arglist)
 
 def parse_config_create_cmd_args(arglist):
@@ -703,8 +798,8 @@ def create_dtbo_image(fout, argv):
         raise ValueError('List of dtimages to add to DTBO not provided')
     dt_entries = parse_dt_entries(global_args, remainder)
     dtbo = Dtbo(fout, global_args.page_size, global_args.version)
-    dtbo.add_dt_entries(dt_entries)
-    dtbo.commit()
+    dt_entry_buf = dtbo.add_dt_entries(dt_entries)
+    dtbo.commit(dt_entry_buf)
     fout.close()
 
 def dump_dtbo_image(fin, argv):
@@ -724,7 +819,7 @@ def dump_dtbo_image(fin, argv):
         num_entries = len(dtbo.dt_entries)
         for idx in range(0, num_entries):
             with open(args.dtfilename + '.{:d}'.format(idx), 'wb') as fout:
-                dtbo.extract_dt_file(idx, fout)
+                dtbo.extract_dt_file(idx, fout, args.decompress)
     args.outfile.write(str(dtbo) + '\n')
     args.outfile.close()
 
@@ -739,7 +834,7 @@ def create_dtbo_image_from_config(fout, argv):
     if not args.conf_file:
         raise ValueError('Configuration file must be provided')
 
-    _DT_KEYS = ('id', 'rev', 'custom0', 'custom1', 'custom2', 'custom3')
+    _DT_KEYS = ('id', 'rev', 'flags', 'custom0', 'custom1', 'custom2')
     _GLOBAL_KEYS = ('page_size', 'version')
 
     global_args, dt_args = parse_config_file(args.conf_file,
@@ -760,8 +855,8 @@ def create_dtbo_image_from_config(fout, argv):
 
     # Create and write DTBO file
     dtbo = Dtbo(fout, global_args['page_size'], global_args['version'])
-    dtbo.add_dt_entries(dt_entries)
-    dtbo.commit()
+    dt_entry_buf = dtbo.add_dt_entries(dt_entries)
+    dtbo.commit(dt_entry_buf)
     fout.close()
 
 def print_default_usage(progname):
@@ -805,10 +900,10 @@ def print_create_usage(progname):
     sb.append('      --version <number>       DTBO version. Default: 0')
     sb.append('      --id <number>       The default value to set property id in dt_table_entry. Default: 0')
     sb.append('      --rev <number>')
+    sb.append('      --flags=<number>')
     sb.append('      --custom0=<number>')
-    sb.append('      --custom2=<number>')
-    sb.append('      --custom3=<number>')
-    sb.append('      --custom4=<number>\n')
+    sb.append('      --custom1=<number>')
+    sb.append('      --custom2=<number>\n')
 
     sb.append('      The value could be a number or a DT node path.')
     sb.append('      <number> could be a 32-bits digit or hex value, ex. 68000, 0x6800.')
diff --git a/utils/tests/data/mkdtboimg.cfg b/utils/tests/data/mkdtboimg.cfg
index 603292c..9c3a8e8 100644
--- a/utils/tests/data/mkdtboimg.cfg
+++ b/utils/tests/data/mkdtboimg.cfg
@@ -2,7 +2,7 @@
   page_size=4096
   id=0x100                # Read the value form the path in dtb
   rev=0x100               # Read the value form the path in dtb
-  custom0=0xabc           # default value
+  flags=0xab0             # default value
   version=1
 
 board1v1.dts.dtb
@@ -16,4 +16,7 @@ board2v1.dts.dtb
   rev=0x201               # override with another value
 
 board1v1.dts.dtb
-  custom0=0xdef           # override with another value
+  flags=0xd01           # override with another value, do zlib compression
+
+board1v1.dts.dtb
+  flags=0xd02           # override with another value, do gzip compression
diff --git a/utils/tests/data/mkdtimg.cfg b/utils/tests/data/mkdtimg.cfg
index e7a464e..ef7d568 100644
--- a/utils/tests/data/mkdtimg.cfg
+++ b/utils/tests/data/mkdtimg.cfg
@@ -15,3 +15,6 @@ board2v1.dts.dtb
 
 board1v1.dts.dtb
   custom0=0xdef           # override with another value
+
+board1v1.dts.dtb
+  custom0=0xdee           # override with another value
diff --git a/utils/tests/mkdtboimg_testdata.sh b/utils/tests/mkdtboimg_testdata.sh
index 1574343..12c2f95 100755
--- a/utils/tests/mkdtboimg_testdata.sh
+++ b/utils/tests/mkdtboimg_testdata.sh
@@ -55,14 +55,16 @@ echo "Creating dtbo image with mkdtbimg"
 mkdtimg create ${MKDTIMG_OUT}/create.img --page_size=4096 --id=0x100 --version=1\
     --rev=0x100 --custom0=0xabc "${OUTDIR}/board1v1.dts.dtb" "${OUTDIR}/board1v1_1.dts.dtb" \
     --id=0xddccbbaa --rev=0x01000100 "${OUTDIR}/board2v1.dts.dtb" --id=0x200 \
-    --rev=0x201 "${OUTDIR}/board1v1.dts.dtb" --custom0=0xdef > /dev/null
+    --rev=0x201 "${OUTDIR}/board1v1.dts.dtb" --custom0=0xdef \
+    "${OUTDIR}/board1v1.dts.dtb" --custom0=0xdef > /dev/null
 
 echo "Creating dtbo image with mkdtboimg"
 ../src/mkdtboimg.py create  ${MKDTBOIMG_OUTCREATE}/create.img --page_size=4096 \
-    --id=0x100 --rev=0x100 --custom0=0xabc --version=1 "${OUTDIR}/board1v1.dts.dtb" \
+    --id=0x100 --rev=0x100 --flags=0xabc0 --version=1 "${OUTDIR}/board1v1.dts.dtb" \
     "${OUTDIR}/board1v1_1.dts.dtb" --id=0xddccbbaa --rev=0x01000100 \
     "${OUTDIR}/board2v1.dts.dtb" --id=0x200 --rev=0x201 \
-    "${OUTDIR}/board1v1.dts.dtb" --custom0=0xdef > /dev/null
+    "${OUTDIR}/board1v1.dts.dtb" --flags=0xd01 \
+    "${OUTDIR}/board1v1.dts.dtb" --flags=0xd02 > /dev/null
 
 echo "Creating dtbo image with ${PYCONFIG} config file"
 ../src/mkdtboimg.py cfg_create ${MKDTBOIMG_OUTCFG}/create.img ${PYCONFIG} --dtb-dir "${OUTDIR}"
@@ -71,10 +73,10 @@ echo "Dumping fragments from mkdtimg tool image"
 mkdtimg dump ${MKDTIMG_OUT}/create.img -b "${MKDTIMG_DUMP}"| grep -v 'FDT' > ${MKDTIMG_OUT}/create.dump
 
 echo "Dumping fragments from mkdtboimg.py tool for image generated with 'create'"
-../src/mkdtboimg.py dump ${MKDTBOIMG_OUTCREATE}/create.img --output ${MKDTBOIMG_OUTCREATE}/create.dump -b "${MKDTBOIMG_CREATEDUMP}"
+../src/mkdtboimg.py dump ${MKDTBOIMG_OUTCREATE}/create.img --output ${MKDTBOIMG_OUTCREATE}/create.dump -b "${MKDTBOIMG_CREATEDUMP}" --decompress
 
 echo "Dumping fragments from mkdtboimg.py tool for image generated with 'cfg_create'"
-../src/mkdtboimg.py dump ${MKDTBOIMG_OUTCFG}/create.img --output ${MKDTBOIMG_OUTCFG}/create.dump -b "${MKDTBOIMG_CFGDUMP}"
+../src/mkdtboimg.py dump ${MKDTBOIMG_OUTCFG}/create.img --output ${MKDTBOIMG_OUTCFG}/create.dump -b "${MKDTBOIMG_CFGDUMP}" --decompress
 
 echo "======================================================================================"
 echo "Testing differences between image created by 'create' for 'mkdtimg' and 'mkdtboimg.py'"
@@ -89,8 +91,6 @@ do
     echo "diff $x vs ${MKDTBOIMG_OUTCREATE}/$file"
     diff $x ${MKDTBOIMG_OUTCREATE}/$file
 done
-echo "diff ${MKDTIMG_OUT}/create.dump vs ${MKDTBOIMG_OUTCREATE}/create.dump";
-diff ${MKDTIMG_OUT}/create.dump ${MKDTBOIMG_OUTCREATE}/create.dump
 echo "=========================================================================================="
 echo "Testing differences between image created by 'cfg_create' for 'mkdtimg' and 'mkdtboimg.py'"
 echo "=========================================================================================="
@@ -104,5 +104,3 @@ do
     echo "diff $x vs ${MKDTBOIMG_OUTCFG}/$file"
     diff $x ${MKDTBOIMG_OUTCFG}/$file
 done
-echo "diff ${MKDTIMG_OUT}/create.dump vs ${MKDTBOIMG_OUTCFG}/create.dump";
-diff ${MKDTIMG_OUT}/create.dump ${MKDTBOIMG_OUTCFG}/create.dump