Merge pull request #6 from WallaceIT/zstd_and_sync

Add support for Zstandard and write with O_SYNC

CMakeLists.txt

@@ -39,12 +39,21 @@ else()
     message(FATAL_ERROR "liblzma not found")
 endif()
 
+# Find libzstd
+find_package(PkgConfig REQUIRED)
+pkg_check_modules(ZSTD REQUIRED IMPORTED_TARGET libzstd)
+if (ZSTD_FOUND)
+    include_directories(${ZSTD_INCLUDE_DIRS})
+else()
+    message(FATAL_ERROR "libzstd not found")
+endif()
+
 # Add the executable
 add_executable(bmap-writer bmap-writer.cpp)
 target_compile_options(bmap-writer PUBLIC -Wformat -Wformat-security -Wconversion -Wsign-conversion -pedantic -Werror)
 
 # Link the libraries
-target_link_libraries(bmap-writer ${LIBXML2_LIBRARIES} ${OPENSSL_LIBRARIES} ${ZLIB_LIBRARIES} ${LIBLZMA_LIBRARIES})
+target_link_libraries(bmap-writer ${LIBXML2_LIBRARIES} ${OPENSSL_LIBRARIES} ${ZLIB_LIBRARIES} ${LIBLZMA_LIBRARIES} ${ZSTD_LIBRARIES})
 
 # Specify the install rules
 install(TARGETS bmap-writer DESTINATION bin)

README.md

@@ -11,7 +11,7 @@ Unlike the Yocto BMAP tool, `bmap-writer` is C++ based does not require Python a
 
 - **Alternative to Yocto BMAP Tool**: Provides a lightweight alternative specifically for embedded systems.
 - **No Python Required**: Does not require Python, making it easier to integrate into various environments.
-- **Support for Compressed Images**: Handles gzip and xz compressed images, decompressing them on-the-fly during the writing process.
+- **Support for Compressed Images**: Handles gzip, xz anz zstd compressed images, decompressing them on-the-fly during the writing process.
 - **Checksum Verification**: Ensures data integrity by verifying checksums for each block.
 - **Efficient Writing**: Writes only the necessary blocks, reducing the overall write time and wear on storage devices.
 

bmap-writer-test.sh

@@ -10,12 +10,21 @@ if [ ! -f test.img ]; then
     dd if=/dev/urandom of=test.img bs=4k count=1 seek=131072 conv=notrunc  > /dev/null 2>&1
 fi
 
-if [ ! -f test.img.gz ] || [ ! -f test.img.xz ]; then
-    echo "## Compress the file with xz and gzip"
-    xz   -z test.img -c   > test.img.xz
+if [ ! -f test.img.gz ]; then
+    echo "## Compress the file with gzip"
     gzip -9 test.img -c   > test.img.gz
 fi
 
+if [ ! -f test.img.xz ]; then
+    echo "## Compress the file with xz"
+    xz   -z test.img -c   > test.img.xz
+fi
+
+if [ ! -f test.img.zst ]; then
+    echo "## Compress the file with zstd"
+    zstd -f -k -c -3 --threads=8 test.img > test.img.zst
+fi
+
 if [ ! -f test.img.bmap ] ; then
     echo "## Create a bmap file"
     bmaptool create test.img -o test.img.bmap
@@ -35,3 +44,7 @@ cmp test.img.out test.gz.img.out
 echo "## Write the file with bmap-writer and xz"
 ./bmap-writer test.img.xz test.img.bmap test.xz.img.out
 cmp test.img.out test.xz.img.out
+
+echo "## Write the file with bmap-writer and zstd"
+./bmap-writer test.img.zst test.img.bmap test.zst.img.out
+cmp test.img.out test.zst.img.out

bmap-writer.cpp

@@ -36,6 +36,7 @@
 #include <openssl/evp.h>
 #include <zlib.h>
 #include <lzma.h>
+#include <zstd.h>
 
 #define CHECKSUM_LENGTH 64
 #define RANGE_LENGTH    19
@@ -48,6 +49,10 @@
 #define XZ_MAGIC_3   'X'
 #define XZ_MAGIC_4   'Z'
 #define XZ_MAGIC_5  0x00
+#define ZSTD_MAGIC_0 0x28
+#define ZSTD_MAGIC_1 0xb5
+#define ZSTD_MAGIC_2 0x2f
+#define ZSTD_MAGIC_3 0xfd
 
 #define DEC_BUFFER_SIZE (1024 * 16)
 
@@ -142,6 +147,9 @@ int getCompressionType(const std::string &imageFile, std::string &compressionTyp
     } else if (buffer[0] == XZ_MAGIC_0 && buffer[1] == XZ_MAGIC_1 && buffer[2] == XZ_MAGIC_2 &&
                buffer[3] == XZ_MAGIC_3 && buffer[4] == XZ_MAGIC_4 && buffer[5] == XZ_MAGIC_5) {
         compressionType = "xz";
+    } else if (buffer[0] == ZSTD_MAGIC_0 && buffer[1] == ZSTD_MAGIC_1 &&
+               buffer[2] == ZSTD_MAGIC_2 && buffer[3] == ZSTD_MAGIC_3) {
+        compressionType = "zstd";
     } else {
         compressionType = "none";
     }
@@ -176,13 +184,16 @@ int BmapWriteImage(const std::string &imageFile, const bmap_t &bmap, const std::
     gzFile gzImg = nullptr;
     lzma_stream lzmaStream = LZMA_STREAM_INIT;
     std::vector<char> decBufferIn(DEC_BUFFER_SIZE);
+    ZSTD_DStream* zstdStream = nullptr;
+    ZSTD_inBuffer zstdIn = { nullptr, 0, 0 };
+    ZSTD_outBuffer zstdOut = { nullptr, 0, 0 };
     size_t decHead = 0;
     std::ifstream imgFile;
     int dev_fd = -1;
     int ret = 0;
 
     try {
-        dev_fd = open(device.c_str(), O_WRONLY | O_CREAT, S_IRUSR | S_IWUSR);
+        dev_fd = open(device.c_str(), O_WRONLY | O_CREAT | O_SYNC, S_IRUSR | S_IWUSR);
         if (dev_fd < 0) {
             throw std::string("Unable to open or create target device");
         }
@@ -203,6 +214,20 @@ int BmapWriteImage(const std::string &imageFile, const bmap_t &bmap, const std::
             }
 
             lzmaStream.avail_in = 0;
+        } else if (compressionType == "zstd") {
+            imgFile.open(imageFile, std::ios::binary);
+            if (!imgFile) {
+                throw std::string("Unable to open image file");
+            }
+
+            zstdStream = ZSTD_createDStream();
+            if (zstdStream == nullptr) {
+                throw std::string("Failed to initialize zstd decoder");
+            }
+
+            zstdIn.src = decBufferIn.data();
+            zstdIn.size = decBufferIn.size();
+            zstdIn.pos = zstdIn.size;
         } else if (compressionType == "none") {
             imgFile.open(imageFile, std::ios::binary);
             if (!imgFile) {
@@ -292,6 +317,61 @@ int BmapWriteImage(const std::string &imageFile, const bmap_t &bmap, const std::
                         break;
                     }
                 }
+            } else if (compressionType == "zstd") {
+                const size_t outStart = startBlock * bmap.blockSize;
+                const size_t outEnd = ((endBlock + 1) * bmap.blockSize);
+
+                // Init output buffer
+                zstdOut.dst = buffer.data();
+                zstdOut.size = buffer.size();
+                zstdOut.pos = 0;
+
+                while (outBytes < bufferSize) {
+                    size_t chunkSize = 0;
+                    size_t zrc;
+
+                    if (zstdIn.pos == zstdIn.size) {
+                        imgFile.read(decBufferIn.data(), static_cast<ssize_t>(decBufferIn.size()));
+                        if (imgFile.gcount() == 0 && imgFile.fail()) {
+                            throw std::string("Failed to read from zstd image file");
+                        } else {
+                            zstdIn.size = static_cast<size_t>(imgFile.gcount());
+                            zstdIn.pos = 0;
+                        }
+                    }
+
+                    zrc = ZSTD_decompressStream(zstdStream, &zstdOut, &zstdIn);
+                    if (ZSTD_isError(zrc)) {
+                        throw std::string("Failed to decompress zstd image file: ") + std::string(ZSTD_getErrorName(zrc));
+                    }
+
+                    chunkSize = zstdOut.pos - outBytes;
+
+                    if (decHead >= outStart && (decHead + chunkSize) <= outEnd) {
+                        // Case 1: all decoded data can be used
+                        outBytes += chunkSize;
+                    } else if (decHead < outStart && (decHead + chunkSize) <= outStart) {
+                        // Case 2: all decoded data shall be discarded
+                        zstdOut.pos = 0;
+                    } else if (decHead < outStart && (decHead + chunkSize) > outStart) {
+                        // Case 3: only the last portion of the decoded data can be used
+                        std::move(buffer.begin() + static_cast<long int>(outStart - decHead),
+                                  buffer.begin() + static_cast<long int>(chunkSize),
+                                  buffer.begin());
+                        size_t validData = chunkSize - (outStart - decHead);
+                        outBytes += validData;
+                        zstdOut.pos = validData;
+                    }
+
+                    // Advance the head of the decompressed data
+                    decHead += chunkSize;
+
+                    // If no more data is available in the input buffer and the input file has been
+                    // read completely, stop this decompression loop
+                    if ((zstdIn.pos == zstdIn.size) && imgFile.eof()) {
+                        break;
+                    }
+                }
             } else if (compressionType == "none") {
                 imgFile.seekg(static_cast<std::streamoff>(startBlock * bmap.blockSize), std::ios::beg);
                 imgFile.read(buffer.data(), static_cast<std::streamsize>(bufferSize));
@@ -341,6 +421,8 @@ int BmapWriteImage(const std::string &imageFile, const bmap_t &bmap, const std::
         gzclose(gzImg);
     } else if (compressionType == "xz") {
         lzma_end(&lzmaStream);
+    } else if (compressionType == "zstd") {
+        ZSTD_freeDStream(zstdStream);
     }
 
     return ret;