From 5f5d263df6486a587abbac108c85ca07eb369681 Mon Sep 17 00:00:00 2001 From: Francesco Valla Date: Sat, 14 Dec 2024 22:48:22 +0100 Subject: [PATCH 1/2] Use std::string to host sha256 hash Instead of using a C-style buffer, rely on std:string to host the calculated sha256 hash and to compare it with the expected one. Signed-off-by: Francesco Valla --- bmap-writer.cpp | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/bmap-writer.cpp b/bmap-writer.cpp index 17aea3a..b698c5c 100644 --- a/bmap-writer.cpp +++ b/bmap-writer.cpp @@ -20,6 +20,7 @@ #include #include +#include #include #include #include @@ -104,21 +105,24 @@ bmap_t parseBMap(const std::string &filename) { return bmapData; } -void computeSHA256(const char *buffer, size_t size, char *output) { +std::string computeSHA256(const std::vector& buffer, size_t size) { EVP_MD_CTX *mdctx; unsigned char hash[EVP_MAX_MD_SIZE]; unsigned int hash_len; mdctx = EVP_MD_CTX_new(); EVP_DigestInit_ex(mdctx, EVP_sha256(), NULL); - EVP_DigestUpdate(mdctx, buffer, size); + EVP_DigestUpdate(mdctx, buffer.data(), size); EVP_DigestFinal_ex(mdctx, hash, &hash_len); EVP_MD_CTX_free(mdctx); + std::ostringstream output; + output << std::hex; for (unsigned int i = 0; i < hash_len; ++i) { - sprintf(output + (i * 2), "%02x", hash[i]); + output << std::setfill('0') << std::setw(2) << static_cast(hash[i]); } - output[CHECKSUM_LENGTH] = 0; + + return output.str(); } int getCompressionType(const std::string &imageFile, std::string &compressionType) { @@ -318,9 +322,8 @@ int BmapWriteImage(const std::string &imageFile, const bmap_t &bmap, const std:: } // Compute and verify the checksum - char computedChecksum[CHECKSUM_LENGTH + 1]; - computeSHA256(buffer.data(), outBytes, computedChecksum); - if (strcmp(computedChecksum, range.checksum.c_str()) != 0) { + std::string computedChecksum = computeSHA256(buffer, outBytes); + if (computedChecksum != range.checksum) { std::cerr << "Checksum verification failed for range: " << range.range << std::endl; std::cerr << "Computed Checksum: " << computedChecksum << std::endl; std::cerr << "Expected Checksum: " << range.checksum << std::endl; From f33b5b798475c91a2024ebb81016a0c86ccc57af Mon Sep 17 00:00:00 2001 From: Francesco Valla Date: Sat, 14 Dec 2024 22:50:54 +0100 Subject: [PATCH 2/2] Simplify error path in image writing To avoid code duplication and have a simpler error path, use try-catch instead of multiple return points. Signed-off-by: Francesco Valla --- bmap-writer.cpp | 293 ++++++++++++++++++++++-------------------------- 1 file changed, 137 insertions(+), 156 deletions(-) diff --git a/bmap-writer.cpp b/bmap-writer.cpp index b698c5c..4f96c47 100644 --- a/bmap-writer.cpp +++ b/bmap-writer.cpp @@ -173,196 +173,177 @@ void printBufferHex(const char *buffer, size_t size) { } int BmapWriteImage(const std::string &imageFile, const bmap_t &bmap, const std::string &device, const std::string &compressionType) { - int dev_fd = open(device.c_str(), O_WRONLY | O_CREAT, S_IRUSR | S_IWUSR); - if (dev_fd < 0) { - std::cerr << "Unable to open or create target device" << std::endl; - return 1; - } - gzFile gzImg = nullptr; lzma_stream lzmaStream = LZMA_STREAM_INIT; std::vector decBufferIn(DEC_BUFFER_SIZE); size_t decHead = 0; std::ifstream imgFile; + int dev_fd = -1; + int ret = 0; - if (compressionType == "gzip") { - gzImg = gzopen(imageFile.c_str(), "rb"); - if (!gzImg) { - std::cerr << "Unable to open gzip image file" << std::endl; - close(dev_fd); - return 1; - } - } else if (compressionType == "xz") { - imgFile.open(imageFile, std::ios::binary); - if (!imgFile) { - std::cerr << "Unable to open xz image file" << std::endl; - close(dev_fd); - return 1; - } - lzma_ret ret = lzma_stream_decoder(&lzmaStream, UINT64_MAX, 0); - if (ret != LZMA_OK) { - std::cerr << "Failed to initialize lzma decoder: " << ret << std::endl; - close(dev_fd); - return 1; + try { + dev_fd = open(device.c_str(), O_WRONLY | O_CREAT, S_IRUSR | S_IWUSR); + if (dev_fd < 0) { + throw std::string("Unable to open or create target device"); } - lzmaStream.avail_in = 0; - } else if (compressionType == "none") { - imgFile.open(imageFile, std::ios::binary); - if (!imgFile) { - std::cerr << "Unable to open image file" << std::endl; - close(dev_fd); - return 1; - } - } else { - std::cerr << "Unsupported compression type" << std::endl; - close(dev_fd); - return 1; - } + if (compressionType == "gzip") { + gzImg = gzopen(imageFile.c_str(), "rb"); + if (!gzImg) { + throw std::string("Unable to open gzip image file"); + } + } else if (compressionType == "xz") { + imgFile.open(imageFile, std::ios::binary); + if (!imgFile) { + throw std::string("Unable to open xz image file"); + } + lzma_ret ret = lzma_stream_decoder(&lzmaStream, UINT64_MAX, 0); + if (ret != LZMA_OK) { + throw std::string("Failed to initialize lzma decoder: ") + std::to_string(static_cast(ret)); + } - for (const auto &range : bmap.ranges) { - size_t startBlock, endBlock; - if (sscanf(range.range.c_str(), "%zu-%zu", &startBlock, &endBlock) == 1) { - endBlock = startBlock; // Handle single block range + lzmaStream.avail_in = 0; + } else if (compressionType == "none") { + imgFile.open(imageFile, std::ios::binary); + if (!imgFile) { + throw std::string("Unable to open image file"); + } + } else { + throw std::string("Unsupported compression type ") + compressionType; } - std::cout << "Processing Range: startBlock=" << startBlock << ", endBlock=" << endBlock << std::endl; - - size_t bufferSize = (endBlock - startBlock + 1) * bmap.blockSize; - std::vector buffer(bufferSize); - size_t outBytes = 0; - if (compressionType == "gzip") { - gzseek(gzImg, static_cast(startBlock * bmap.blockSize), SEEK_SET); - int readBytes = gzread(gzImg, buffer.data(), static_cast(bufferSize)); - if (readBytes < 0) { - std::cerr << "Failed to read from gzip image file" << std::endl; - close(dev_fd); - gzclose(gzImg); - return 1; + for (const auto &range : bmap.ranges) { + size_t startBlock, endBlock; + if (sscanf(range.range.c_str(), "%zu-%zu", &startBlock, &endBlock) == 1) { + endBlock = startBlock; // Handle single block range } - outBytes = static_cast(readBytes); - } else if (compressionType == "xz") { - const size_t outStart = startBlock * bmap.blockSize; - const size_t outEnd = ((endBlock + 1) * bmap.blockSize); - - // Initialize the output buffer for the decompressor - lzmaStream.next_out = reinterpret_cast(buffer.data()); - lzmaStream.avail_out = static_cast(buffer.size()); - - while (outBytes < bufferSize) { - size_t chunkSize = 0; - - // Whenever no more input data is available, read some from the compressed file - // and reset the input parameters for the decompressor - if (lzmaStream.avail_in == 0) { - imgFile.read(decBufferIn.data(), static_cast(decBufferIn.size())); - if (imgFile.gcount() == 0 && imgFile.fail()) { - std::cerr << "Failed to read from xz image file" << std::endl; - close(dev_fd); - imgFile.close(); - return 1; - } else { - lzmaStream.next_in = reinterpret_cast(decBufferIn.data()); - lzmaStream.avail_in = static_cast(imgFile.gcount()); - } - } + std::cout << "Processing Range: startBlock=" << startBlock << ", endBlock=" << endBlock << std::endl; - // Save the current status of the output buffer... - chunkSize = lzmaStream.avail_out; + size_t bufferSize = (endBlock - startBlock + 1) * bmap.blockSize; + std::vector buffer(bufferSize); + size_t outBytes = 0; - lzma_ret ret = lzma_code(&lzmaStream, LZMA_RUN); - if (ret != LZMA_OK && ret != LZMA_STREAM_END) { - std::cerr << "Failed to decompress xz image file: " << ret << std::endl; - close(dev_fd); - imgFile.close(); - return 1; + if (compressionType == "gzip") { + gzseek(gzImg, static_cast(startBlock * bmap.blockSize), SEEK_SET); + int readBytes = gzread(gzImg, buffer.data(), static_cast(bufferSize)); + if (readBytes < 0) { + throw std::string("Failed to read from gzip image file"); } + outBytes = static_cast(readBytes); + } else if (compressionType == "xz") { + const size_t outStart = startBlock * bmap.blockSize; + const size_t outEnd = ((endBlock + 1) * bmap.blockSize); + + // Initialize the output buffer for the decompressor + lzmaStream.next_out = reinterpret_cast(buffer.data()); + lzmaStream.avail_out = static_cast(buffer.size()); + + while (outBytes < bufferSize) { + size_t chunkSize = 0; + + // Whenever no more input data is available, read some from the compressed file + // and reset the input parameters for the decompressor + if (lzmaStream.avail_in == 0) { + imgFile.read(decBufferIn.data(), static_cast(decBufferIn.size())); + if (imgFile.gcount() == 0 && imgFile.fail()) { + throw std::string("Failed to read from xz image file"); + } else { + lzmaStream.next_in = reinterpret_cast(decBufferIn.data()); + lzmaStream.avail_in = static_cast(imgFile.gcount()); + } + } - // ...and then extract the size of the decompressed chunk - chunkSize -= lzmaStream.avail_out; - - 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 - lzmaStream.next_out = reinterpret_cast(buffer.data()); - lzmaStream.avail_out = static_cast(buffer.size()); - } 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; - lzmaStream.next_out = reinterpret_cast(buffer.data()) + validData; - lzmaStream.avail_out = buffer.size() - validData; - } + // Save the current status of the output buffer... + chunkSize = lzmaStream.avail_out; + + lzma_ret ret = lzma_code(&lzmaStream, LZMA_RUN); + if (ret != LZMA_OK && ret != LZMA_STREAM_END) { + throw std::string("Failed to decompress xz image file: ") + std::to_string(static_cast(ret)); + } + + // ...and then extract the size of the decompressed chunk + chunkSize -= lzmaStream.avail_out; + + 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 + lzmaStream.next_out = reinterpret_cast(buffer.data()); + lzmaStream.avail_out = static_cast(buffer.size()); + } 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; + lzmaStream.next_out = reinterpret_cast(buffer.data()) + validData; + lzmaStream.avail_out = buffer.size() - validData; + } - // Advance the head of the decompressed data - decHead += chunkSize; + // Advance the head of the decompressed data + decHead += chunkSize; - // In case all the required data has been decompressed OR the XZ stream is ended - // OR the input file has been read completely, stop this decompression loop - if ((lzmaStream.avail_out == 0) || (ret == LZMA_STREAM_END) || - (lzmaStream.avail_in == 0 && imgFile.eof())) { - break; + // In case all the required data has been decompressed OR the XZ stream is ended + // OR the input file has been read completely, stop this decompression loop + if ((lzmaStream.avail_out == 0) || (ret == LZMA_STREAM_END) || + (lzmaStream.avail_in == 0 && imgFile.eof())) { + break; + } + } + } else if (compressionType == "none") { + imgFile.seekg(static_cast(startBlock * bmap.blockSize), std::ios::beg); + imgFile.read(buffer.data(), static_cast(bufferSize)); + outBytes = static_cast(imgFile.gcount()); + if (outBytes == 0 && imgFile.fail()) { + throw std::string("Failed to read from image file"); } } - } else if (compressionType == "none") { - imgFile.seekg(static_cast(startBlock * bmap.blockSize), std::ios::beg); - imgFile.read(buffer.data(), static_cast(bufferSize)); - outBytes = static_cast(imgFile.gcount()); - if (outBytes == 0 && imgFile.fail()) { - std::cerr << "Failed to read from image file" << std::endl; - close(dev_fd); - imgFile.close(); - return 1; + + // Compute and verify the checksum + std::string computedChecksum = computeSHA256(buffer, outBytes); + if (computedChecksum != range.checksum) { + std::stringstream err; + err << "Checksum verification failed for range: " << range.range << std::endl; + err << "Computed Checksum: " << computedChecksum << std::endl; + err << "Expected Checksum: " << range.checksum; + //std::cerr << "Buffer content (hex):" << std::endl; + //printBufferHex(buffer.data(), outBytes); + throw std::string(err.str()); } - } - // Compute and verify the checksum - std::string computedChecksum = computeSHA256(buffer, outBytes); - if (computedChecksum != range.checksum) { - std::cerr << "Checksum verification failed for range: " << range.range << std::endl; - std::cerr << "Computed Checksum: " << computedChecksum << std::endl; - std::cerr << "Expected Checksum: " << range.checksum << std::endl; - //std::cerr << "Buffer content (hex):" << std::endl; - //printBufferHex(buffer.data(), outBytes); - close(dev_fd); - if (compressionType == "gzip") { - gzclose(gzImg); - } else if (compressionType == "xz" || compressionType == "none") { - imgFile.close(); + if (pwrite(dev_fd, buffer.data(), outBytes, static_cast(startBlock * bmap.blockSize)) < 0) { + throw std::string("Write to device failed"); } - return 1; } - if (pwrite(dev_fd, buffer.data(), outBytes, static_cast(startBlock * bmap.blockSize)) < 0) { - std::cerr << "Write to device failed"<< std::endl; - close(dev_fd); - if (compressionType == "gzip") { - gzclose(gzImg); - } else if (compressionType == "xz" || compressionType == "none") { - imgFile.close(); - } - return 1; + if (fsync(dev_fd) != 0) { + throw std::string("fsync failed after all writes"); } + + std::cout << "Finished writing image to device." << std::endl; + } + catch (std::string& err) { + std::cerr << err << std::endl; + ret = -1; } - if (fsync(dev_fd) != 0) { - std::cerr << "fsync failed after all writes"<< std::endl; + if (dev_fd >= 0) { + close(dev_fd); + } + if (imgFile.is_open()) { + imgFile.close(); } - close(dev_fd); if (compressionType == "gzip") { gzclose(gzImg); - } else if (compressionType == "xz" || compressionType == "none") { - imgFile.close(); + } else if (compressionType == "xz") { + lzma_end(&lzmaStream); } - std::cout << "Finished writing image to device." << std::endl; - return 0; + + return ret; } int main(int argc, char *argv[]) {