diff --git a/CMakeLists.txt b/CMakeLists.txt index 5233cb1..6a99efb 100644 --- a/CMakeLists.txt +++ b/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) diff --git a/README.md b/README.md index 4d4b1c0..ba0f958 100644 --- a/README.md +++ b/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. diff --git a/bmap-writer-test.sh b/bmap-writer-test.sh index aaa1328..85b29b2 100755 --- a/bmap-writer-test.sh +++ b/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 diff --git a/bmap-writer.cpp b/bmap-writer.cpp index 86e5b2e..7ce73b4 100644 --- a/bmap-writer.cpp +++ b/bmap-writer.cpp @@ -36,6 +36,7 @@ #include #include #include +#include #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 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(decBufferIn.size())); + if (imgFile.gcount() == 0 && imgFile.fail()) { + throw std::string("Failed to read from zstd image file"); + } else { + zstdIn.size = static_cast(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(outStart - decHead), + buffer.begin() + static_cast(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(startBlock * bmap.blockSize), std::ios::beg); imgFile.read(buffer.data(), static_cast(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;