From cd15b640b0b8d5a7ba5f1c0587e4f9c767e2d8fb Mon Sep 17 00:00:00 2001
From: Stefan Herbrechtsmeier <stefan.herbrechtsmeier@weidmueller.com>
Date: Fri, 19 Aug 2022 16:25:38 +0200
Subject: [PATCH] binman: Add zstd bintool

Add zstd bintool to binman to support on-the-fly compression.

Signed-off-by: Stefan Herbrechtsmeier <stefan.herbrechtsmeier@weidmueller.com>
Reviewed-by: Simon Glass <sjg@chromium.org>
---
 tools/binman/btool/zstd.py                  | 30 +++++++++++++++++++++++++++++
 tools/binman/entry.py                       |  2 +-
 tools/binman/etype/blob_dtb.py              |  4 ++++
 tools/binman/ftest.py                       | 12 ++++++++++--
 tools/binman/test/238_compress_dtb_zstd.dts | 16 +++++++++++++++
 5 files changed, 61 insertions(+), 3 deletions(-)
 create mode 100644 tools/binman/btool/zstd.py
 create mode 100644 tools/binman/test/238_compress_dtb_zstd.dts

diff --git a/tools/binman/btool/zstd.py b/tools/binman/btool/zstd.py
new file mode 100644
index 0000000..299bd37
--- /dev/null
+++ b/tools/binman/btool/zstd.py
@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: GPL-2.0+
+# Copyright (C) 2022 Weidmüller Interface GmbH & Co. KG
+# Stefan Herbrechtsmeier <stefan.herbrechtsmeier@weidmueller.com>
+#
+"""Bintool implementation for zstd
+
+zstd allows compression and decompression of files.
+
+Documentation is available via::
+
+   man zstd
+"""
+
+from binman import bintool
+
+# pylint: disable=C0103
+class Bintoolzstd(bintool.BintoolPacker):
+    """Compression/decompression using the zstd algorithm
+
+    This bintool supports running `zstd` to compress and decompress data, as
+    used by binman.
+
+    It is also possible to fetch the tool, which uses `apt` to install it.
+
+    Documentation is available via::
+
+        man zstd
+    """
+    def __init__(self, name):
+        super().__init__(name)
diff --git a/tools/binman/entry.py b/tools/binman/entry.py
index 770130b..63ec5ce 100644
--- a/tools/binman/entry.py
+++ b/tools/binman/entry.py
@@ -1191,7 +1191,7 @@ features to produce new behaviours.
         """
         algo = self.compress
         if algo != 'none':
-            algos = ['bzip2', 'gzip', 'lz4', 'lzma', 'lzo', 'xz']
+            algos = ['bzip2', 'gzip', 'lz4', 'lzma', 'lzo', 'xz', 'zstd']
             if algo not in algos:
                 raise ValueError("Unknown algorithm '%s'" % algo)
             names = {'lzma': 'lzma_alone', 'lzo': 'lzop'}
diff --git a/tools/binman/etype/blob_dtb.py b/tools/binman/etype/blob_dtb.py
index 5a6a454..6a3fbc4 100644
--- a/tools/binman/etype/blob_dtb.py
+++ b/tools/binman/etype/blob_dtb.py
@@ -47,6 +47,10 @@ class Entry_blob_dtb(Entry_blob):
     def ProcessContents(self):
         """Re-read the DTB contents so that we get any calculated properties"""
         _, indata = state.GetFdtContents(self.GetFdtEtype())
+
+        if self.compress == 'zstd' and self.prepend != 'length':
+            self.Raise('The zstd compression requires a length header')
+
         data = self.CompressData(indata)
         return self.ProcessContentsUpdate(data)
 
diff --git a/tools/binman/ftest.py b/tools/binman/ftest.py
index bfca9c1..0b17740 100644
--- a/tools/binman/ftest.py
+++ b/tools/binman/ftest.py
@@ -107,7 +107,7 @@ BASE_DTB_PROPS = ['offset', 'size', 'image-pos']
 REPACK_DTB_PROPS = ['orig-offset', 'orig-size']
 
 # Supported compression bintools
-COMP_BINTOOLS = ['bzip2', 'gzip', 'lz4', 'lzma_alone', 'lzop', 'xz']
+COMP_BINTOOLS = ['bzip2', 'gzip', 'lz4', 'lzma_alone', 'lzop', 'xz', 'zstd']
 
 class TestFunctional(unittest.TestCase):
     """Functional tests for binman
@@ -5881,7 +5881,8 @@ fdt         fdtmap                Extract the devicetree blob from the fdtmap
 
     def testCompUtilPadding(self):
         """Test padding of compression algorithms"""
-        for bintool in self.comp_bintools.values():
+        # Skip zstd because it doesn't support padding
+        for bintool in [v for k,v in self.comp_bintools.items() if k != 'zstd']:
             self._CheckBintool(bintool)
             data = bintool.compress(COMPRESS_DATA)
             self.assertNotEqual(COMPRESS_DATA, data)
@@ -5889,6 +5890,13 @@ fdt         fdtmap                Extract the devicetree blob from the fdtmap
             orig = bintool.decompress(data)
             self.assertEquals(COMPRESS_DATA, orig)
 
+    def testCompressDtbZstd(self):
+        """Test that zstd compress of device-tree files failed"""
+        with self.assertRaises(ValueError) as e:
+            self._DoTestFile('238_compress_dtb_zstd.dts')
+        self.assertIn("Node '/binman/u-boot-dtb': The zstd compression "
+                      "requires a length header", str(e.exception))
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/tools/binman/test/238_compress_dtb_zstd.dts b/tools/binman/test/238_compress_dtb_zstd.dts
new file mode 100644
index 0000000..90cf85d
--- /dev/null
+++ b/tools/binman/test/238_compress_dtb_zstd.dts
@@ -0,0 +1,16 @@
+// SPDX-License-Identifier: GPL-2.0+
+
+/dts-v1/;
+
+/ {
+	#address-cells = <1>;
+	#size-cells = <1>;
+
+	binman {
+		u-boot {
+		};
+		u-boot-dtb {
+			compress = "zstd";
+		};
+	};
+};
-- 
2.7.4