Updater: resource compression (#3716)

* toolbox: compress: moved decompressor implementation to separate func
* toolbox: compress: callback-based api; cli: storage unpack command
* toolbox: compress: separate r/w contexts for stream api
* targets: f18: sync API
* compress: naming fixes & cleanup
* toolbox: compress: using hs buffer size for stream buffers
* toolbox: tar: heatshrink stream mode
* toolbox: compress: docs & small cleanup
* toolbox: tar: header support for .hs; updater: now uses .hs for resources; .hs.tar: now rewindable
* toolbox: compress: fixed hs stream tail handling
* updater: reworked progress for resources cleanup; rebalanced stage weights
* updater: single-pass decompression; scripts: print resources compression ratio
* updater: fixed warnings
* toolbox: tar: doxygen
* docs: update
* docs: info or tarhs format; scripts: added standalone compression/decompression tool for heatshrink-formatted streams
* scripts: tarhs: fixed parameter handling
* cli: storage extract command; toolbox: tar: guess type based on extension
* unit_tests: added test for streamed raw hs decompressor `compress_decode_streamed`
* unit_tests: compress: added extraction test for .tar.hs
* rpc: autodetect compressed archives
* scripts: minor cleanup of common parts
* scripts: update: now using in-memory intermediate tar stream
* scripts: added hs.py wrapper for heatshrink-related ops (single object and directory-as-tar compression)
* scripts: naming fixes
* Toolbox: export compress_config_heatshrink_default as const symbol
* Toolbox: fix various types naming
* Toolbox: more of types naming fixes
* Toolbox: use size_t in compress io callbacks and structures
* UnitTests: update to match new compress API
* Toolbox: proper path_extract_extension usage

Co-authored-by: あく <alleteam@gmail.com>

scripts/flipper/assets/heatshrink_stream.py

@@ -0,0 +1,26 @@
+import struct
+
+
+class HeatshrinkDataStreamHeader:
+    MAGIC = 0x53445348
+    VERSION = 1
+
+    def __init__(self, window_size, lookahead_size):
+        self.window_size = window_size
+        self.lookahead_size = lookahead_size
+
+    def pack(self):
+        return struct.pack(
+            "<IBBB", self.MAGIC, self.VERSION, self.window_size, self.lookahead_size
+        )
+
+    @staticmethod
+    def unpack(data):
+        if len(data) != 7:
+            raise ValueError("Invalid header length")
+        magic, version, window_size, lookahead_size = struct.unpack("<IBBB", data)
+        if magic != HeatshrinkDataStreamHeader.MAGIC:
+            raise ValueError("Invalid magic number")
+        if version != HeatshrinkDataStreamHeader.VERSION:
+            raise ValueError("Invalid version")
+        return HeatshrinkDataStreamHeader(window_size, lookahead_size)

scripts/flipper/assets/tarball.py

@@ -0,0 +1,41 @@
+import io
+import tarfile
+
+import heatshrink2
+
+from .heatshrink_stream import HeatshrinkDataStreamHeader
+
+FLIPPER_TAR_FORMAT = tarfile.USTAR_FORMAT
+TAR_HEATSRINK_EXTENSION = ".ths"
+
+
+def tar_sanitizer_filter(tarinfo: tarfile.TarInfo):
+    tarinfo.gid = tarinfo.uid = 0
+    tarinfo.mtime = 0
+    tarinfo.uname = tarinfo.gname = "furippa"
+    return tarinfo
+
+
+def compress_tree_tarball(
+    src_dir, output_name, filter=tar_sanitizer_filter, hs_window=13, hs_lookahead=6
+):
+    plain_tar = io.BytesIO()
+    with tarfile.open(
+        fileobj=plain_tar,
+        mode="w:",
+        format=FLIPPER_TAR_FORMAT,
+    ) as tarball:
+        tarball.add(src_dir, arcname="", filter=filter)
+    plain_tar.seek(0)
+
+    src_data = plain_tar.read()
+    compressed = heatshrink2.compress(
+        src_data, window_sz2=hs_window, lookahead_sz2=hs_lookahead
+    )
+
+    header = HeatshrinkDataStreamHeader(hs_window, hs_lookahead)
+    with open(output_name, "wb") as f:
+        f.write(header.pack())
+        f.write(compressed)
+
+    return len(src_data), len(compressed)

scripts/hs.py

@@ -0,0 +1,145 @@
+#!/usr/bin/env python3
+
+import heatshrink2 as hs
+from flipper.app import App
+from flipper.assets.heatshrink_stream import HeatshrinkDataStreamHeader
+from flipper.assets.tarball import compress_tree_tarball
+
+
+class HSWrapper(App):
+    DEFAULT_WINDOW = 13
+    DEFAULT_LOOKAHEAD = 6
+
+    def init(self):
+        self.subparsers = self.parser.add_subparsers(
+            title="subcommands", dest="subcommand"
+        )
+
+        self.parser_compress = self.subparsers.add_parser(
+            "compress", help="compress file using heatshrink"
+        )
+        self.parser_compress.add_argument(
+            "-w", "--window", help="window size", type=int, default=self.DEFAULT_WINDOW
+        )
+        self.parser_compress.add_argument(
+            "-l",
+            "--lookahead",
+            help="lookahead size",
+            type=int,
+            default=self.DEFAULT_LOOKAHEAD,
+        )
+        self.parser_compress.add_argument("file", help="file to compress")
+        self.parser_compress.add_argument(
+            "-o", "--output", help="output file", required=True
+        )
+        self.parser_compress.set_defaults(func=self.compress)
+
+        self.parser_decompress = self.subparsers.add_parser(
+            "decompress", help="decompress file using heatshrink"
+        )
+        self.parser_decompress.add_argument("file", help="file to decompress")
+        self.parser_decompress.add_argument(
+            "-o", "--output", help="output file", required=True
+        )
+        self.parser_decompress.set_defaults(func=self.decompress)
+
+        self.parser_info = self.subparsers.add_parser("info", help="show file info")
+        self.parser_info.add_argument("file", help="file to show info for")
+        self.parser_info.set_defaults(func=self.info)
+
+        self.parser_tar = self.subparsers.add_parser(
+            "tar", help="create a tarball and compress it"
+        )
+        self.parser_tar.add_argument("dir", help="directory to tar")
+        self.parser_tar.add_argument(
+            "-o", "--output", help="output file", required=True
+        )
+        self.parser_tar.add_argument(
+            "-w", "--window", help="window size", type=int, default=self.DEFAULT_WINDOW
+        )
+        self.parser_tar.add_argument(
+            "-l",
+            "--lookahead",
+            help="lookahead size",
+            type=int,
+            default=self.DEFAULT_LOOKAHEAD,
+        )
+        self.parser_tar.set_defaults(func=self.tar)
+
+    def compress(self):
+        args = self.args
+
+        with open(args.file, "rb") as f:
+            data = f.read()
+
+        compressed = hs.compress(
+            data, window_sz2=args.window, lookahead_sz2=args.lookahead
+        )
+
+        with open(args.output, "wb") as f:
+            header = HeatshrinkDataStreamHeader(args.window, args.lookahead)
+            f.write(header.pack())
+            f.write(compressed)
+
+        self.logger.info(
+            f"Compressed {len(data)} bytes to {len(compressed)} bytes, "
+            f"compression ratio: {len(compressed) * 100 / len(data):.2f}%"
+        )
+        return 0
+
+    def decompress(self):
+        args = self.args
+
+        with open(args.file, "rb") as f:
+            header = HeatshrinkDataStreamHeader.unpack(f.read(7))
+            compressed = f.read()
+
+        self.logger.info(
+            f"Decompressing with window size {header.window_size} and lookahead size {header.lookahead_size}"
+        )
+
+        data = hs.decompress(
+            compressed,
+            window_sz2=header.window_size,
+            lookahead_sz2=header.lookahead_size,
+        )
+
+        with open(args.output, "wb") as f:
+            f.write(data)
+
+        self.logger.info(f"Decompressed {len(compressed)} bytes to {len(data)} bytes")
+        return 0
+
+    def info(self):
+        args = self.args
+
+        try:
+            with open(args.file, "rb") as f:
+                header = HeatshrinkDataStreamHeader.unpack(f.read(7))
+        except Exception as e:
+            self.logger.error(f"Error: {e}")
+            return 1
+
+        self.logger.info(
+            f"Window size: {header.window_size}, lookahead size: {header.lookahead_size}"
+        )
+
+        return 0
+
+    def tar(self):
+        args = self.args
+
+        orig_size, compressed_size = compress_tree_tarball(
+            args.dir, args.output, hs_window=args.window, hs_lookahead=args.lookahead
+        )
+
+        self.logger.info(
+            f"Tarred and compressed {orig_size} bytes to {compressed_size} bytes, "
+            f"compression ratio: {compressed_size * 100 / orig_size:.2f}%"
+        )
+
+        return 0
+
+
+if __name__ == "__main__":
+    HSWrapper()()

scripts/sconsdist.py

@@ -9,6 +9,7 @@ from os.path import basename, exists, join, relpath
 
 from ansi.color import fg
 from flipper.app import App
+from flipper.assets.tarball import FLIPPER_TAR_FORMAT, tar_sanitizer_filter
 from update import Main as UpdateMain
 
 
@@ -266,20 +267,15 @@ class Main(App):
                 ),
                 "w:gz",
                 compresslevel=9,
-                format=tarfile.USTAR_FORMAT,
+                format=FLIPPER_TAR_FORMAT,
             ) as tar:
                 self.note_dist_component(
                     "update", "tgz", self.get_dist_path(bundle_tgz)
                 )
 
-                # Strip uid and gid in case of overflow
-                def tar_filter(tarinfo):
-                    tarinfo.uid = tarinfo.gid = 0
-                    tarinfo.mtime = 0
-                    tarinfo.uname = tarinfo.gname = "furippa"
-                    return tarinfo
-
-                tar.add(bundle_dir, arcname=bundle_dir_name, filter=tar_filter)
+                tar.add(
+                    bundle_dir, arcname=bundle_dir_name, filter=tar_sanitizer_filter
+                )
         return bundle_result
 
 

scripts/update.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 
+import io
 import math
 import os
 import shutil
@@ -7,9 +8,12 @@ import tarfile
 import zlib
 from os.path import exists, join
 
+import heatshrink2
 from flipper.app import App
 from flipper.assets.coprobin import CoproBinary, get_stack_type
+from flipper.assets.heatshrink_stream import HeatshrinkDataStreamHeader
 from flipper.assets.obdata import ObReferenceValues, OptionBytesData
+from flipper.assets.tarball import compress_tree_tarball, tar_sanitizer_filter
 from flipper.utils.fff import FlipperFormatFile
 from slideshow import Main as SlideshowMain
 
@@ -20,8 +24,7 @@ class Main(App):
 
     #  No compression, plain tar
     RESOURCE_TAR_MODE = "w:"
-    RESOURCE_TAR_FORMAT = tarfile.USTAR_FORMAT
-    RESOURCE_FILE_NAME = "resources.tar"
+    RESOURCE_FILE_NAME = "resources.ths"  # .Tar.HeatShrink
     RESOURCE_ENTRY_NAME_MAX_LENGTH = 100
 
     WHITELISTED_STACK_TYPES = set(
@@ -34,6 +37,9 @@ class Main(App):
     FLASH_BASE = 0x8000000
     MIN_LFS_PAGES = 6
 
+    HEATSHRINK_WINDOW_SIZE = 13
+    HEATSHRINK_LOOKAHEAD_SIZE = 6
+
     # Post-update slideshow
     SPLASH_BIN_NAME = "splash.bin"
 
@@ -221,23 +227,19 @@ class Main(App):
                 f"Cannot package resource: name '{tarinfo.name}' too long"
             )
             raise ValueError("Resource name too long")
-        tarinfo.gid = tarinfo.uid = 0
-        tarinfo.mtime = 0
-        tarinfo.uname = tarinfo.gname = "furippa"
-        return tarinfo
+        return tar_sanitizer_filter(tarinfo)
 
     def package_resources(self, srcdir: str, dst_name: str):
         try:
-            with tarfile.open(
-                dst_name, self.RESOURCE_TAR_MODE, format=self.RESOURCE_TAR_FORMAT
-            ) as tarball:
-                tarball.add(
-                    srcdir,
-                    arcname="",
-                    filter=self._tar_filter,
-                )
+            src_size, compressed_size = compress_tree_tarball(
+                srcdir, dst_name, filter=self._tar_filter
+            )
+
+            self.logger.info(
+                f"Resources compression ratio: {compressed_size * 100 / src_size:.2f}%"
+            )
             return True
-        except ValueError as e:
+        except Exception as e:
             self.logger.error(f"Cannot package resources: {e}")
             return False