diff --git a/compression/JFJochCompressor.cpp b/compression/JFJochCompressor.cpp index aa4d06b5..1afa0177 100644 --- a/compression/JFJochCompressor.cpp +++ b/compression/JFJochCompressor.cpp @@ -16,6 +16,20 @@ extern "C" { void bshuf_write_uint64_BE(void* buf, uint64_t num); } +// Worst-case size of one compressed block, including its 4-byte length prefix. Mirrors the +// per-block term of MaxCompressedSize(), so a dest sized to MaxCompressedSize() never fails. +static size_t MaxCompressedBlockSize(CompressionAlgorithm algorithm, size_t src_size) { + switch (algorithm) { + case CompressionAlgorithm::BSHUF_LZ4: + return LZ4_compressBound(src_size) + 4; + case CompressionAlgorithm::BSHUF_ZSTD: + case CompressionAlgorithm::BSHUF_ZSTD_RLE: + return ZSTD_compressBound(src_size) + 4; + default: + return src_size + 4; + } +} + JFJochBitShuffleCompressor::JFJochBitShuffleCompressor(CompressionAlgorithm in_algorithm) { algorithm = in_algorithm; } @@ -59,12 +73,12 @@ size_t JFJochBitShuffleCompressor::CompressBlock(char *dest, const char *source, std::vector JFJochBitShuffleCompressor::Compress(const void *source, size_t nelements, size_t elem_size) { std::vector tmp(MaxCompressedSize(algorithm, nelements, elem_size)); - size_t tmp_size = Compress(tmp.data(), source, nelements, elem_size); + size_t tmp_size = Compress(tmp.data(), tmp.size(), source, nelements, elem_size); tmp.resize(tmp_size); return tmp; } -size_t JFJochBitShuffleCompressor::Compress(void *dest, const void *source, size_t nelements, size_t elem_size) { +int64_t JFJochBitShuffleCompressor::Compress(void *dest, size_t dest_size, const void *source, size_t nelements, size_t elem_size) { auto c_dest = (char *) dest; auto c_source = (char *) source; @@ -72,10 +86,15 @@ size_t JFJochBitShuffleCompressor::Compress(void *dest, const void *source, size if (algorithm == CompressionAlgorithm::NO_COMPRESSION) { // Trivial case if no compression - copy content + if (nelements * elem_size > dest_size) + return -1; memcpy(dest, source, nelements * elem_size); return nelements * elem_size; } + if (dest_size < 12) + return -1; + bshuf_write_uint64_BE(c_dest, nelements * elem_size); bshuf_write_uint32_BE(c_dest + 8, DefaultBlockSize * elem_size); @@ -86,18 +105,27 @@ size_t JFJochBitShuffleCompressor::Compress(void *dest, const void *source, size size_t reminder_size = nelements - num_full_blocks * DefaultBlockSize; size_t compressed_size = 12; - - for (int i = 0; i < num_full_blocks; i++) + // Blocks are small relative to the image, so before each one we just check that the + // remaining space still covers that block's worst case, and bail out (-1) if not. + for (int i = 0; i < num_full_blocks; i++) { + if (compressed_size + MaxCompressedBlockSize(algorithm, DefaultBlockSize * elem_size) > dest_size) + return -1; compressed_size += CompressBlock(c_dest + compressed_size, c_source + i * DefaultBlockSize * elem_size, DefaultBlockSize, elem_size); + } size_t last_block_size = reminder_size - reminder_size % BSHUF_BLOCKED_MULT; - if (last_block_size > 0) + if (last_block_size > 0) { + if (compressed_size + MaxCompressedBlockSize(algorithm, last_block_size * elem_size) > dest_size) + return -1; compressed_size += CompressBlock(c_dest + compressed_size, c_source + num_full_blocks * DefaultBlockSize * elem_size, last_block_size, elem_size); + } size_t leftover_bytes = (reminder_size % BSHUF_BLOCKED_MULT) * elem_size; if (leftover_bytes > 0) { + if (compressed_size + leftover_bytes > dest_size) + return -1; memcpy(c_dest + compressed_size, c_source + (num_full_blocks * DefaultBlockSize + last_block_size) * elem_size, leftover_bytes); compressed_size += leftover_bytes; } diff --git a/compression/JFJochCompressor.h b/compression/JFJochCompressor.h index 0da07ce6..255a49cb 100644 --- a/compression/JFJochCompressor.h +++ b/compression/JFJochCompressor.h @@ -25,8 +25,8 @@ public: explicit JFJochBitShuffleCompressor(CompressionAlgorithm algorithm); template - size_t Compress(void *dest, const std::vector &src) { - return Compress((char *) dest, (char *) src.data(), src.size(), sizeof(T)); + int64_t Compress(void *dest, size_t dest_size, const std::vector &src) { + return Compress(dest, dest_size, src.data(), src.size(), sizeof(T)); }; template @@ -34,7 +34,7 @@ public: return Compress(src.data(), src.size(), sizeof(T)); } std::vector Compress(const void* source, size_t nelements, size_t elem_size); - size_t Compress(void *dest, const void* source, size_t nelements, size_t elem_size); + int64_t Compress(void *dest, size_t dest_size, const void* source, size_t nelements, size_t elem_size); private: char scratch[DefaultBlockSize * sizeof(uint64_t)]; }; diff --git a/receiver/FrameTransformation.cpp b/receiver/FrameTransformation.cpp index aa8f7cfb..82ec2d8c 100644 --- a/receiver/FrameTransformation.cpp +++ b/receiver/FrameTransformation.cpp @@ -54,8 +54,8 @@ FrameTransformation::FrameTransformation(const DiffractionExperiment &in_experim } } -size_t FrameTransformation::CompressImage(void *output) { - return compressor.Compress(output, precompression_buffer.data(), experiment.GetPixelsNum(), +int64_t FrameTransformation::CompressImage(void *output, size_t output_size) { + return compressor.Compress(output, output_size, precompression_buffer.data(), experiment.GetPixelsNum(), experiment.GetByteDepthImage()); } @@ -63,14 +63,14 @@ CompressedImage FrameTransformation::GetCompressedImage() { const uint8_t *data; size_t size; if (experiment.GetCompressionAlgorithm() == CompressionAlgorithm::NO_COMPRESSION) { - data = (uint8_t *) precompression_buffer.data(); + data = reinterpret_cast(precompression_buffer.data()); size = experiment.GetPixelsNum() * experiment.GetByteDepthImage(); } else { compressed_buffer.resize(MaxCompressedSize(experiment.GetCompressionAlgorithm(), experiment.GetPixelsNum(), experiment.GetByteDepthImage())); - data = (uint8_t *) compressed_buffer.data(); - size = CompressImage(compressed_buffer.data()); + data = reinterpret_cast(compressed_buffer.data()); + size = static_cast(CompressImage(compressed_buffer.data(), compressed_buffer.size())); } return CompressedImage(data, size,experiment.GetXPixelsNum(), diff --git a/receiver/FrameTransformation.h b/receiver/FrameTransformation.h index e155f013..870f52e5 100644 --- a/receiver/FrameTransformation.h +++ b/receiver/FrameTransformation.h @@ -24,6 +24,6 @@ public: void ProcessModule(const void *input, uint16_t module_number, int data_stream); void FillNotCollectedModule(uint16_t module_number, int data_stream); CompressedImage GetCompressedImage(); - size_t CompressImage(void *output); + int64_t CompressImage(void *output, size_t output_size); // < 0 - error in Compression const void *GetImage() const; }; diff --git a/receiver/JFJochReceiverFPGA.cpp b/receiver/JFJochReceiverFPGA.cpp index e4f3749f..6fe2a27e 100644 --- a/receiver/JFJochReceiverFPGA.cpp +++ b/receiver/JFJochReceiverFPGA.cpp @@ -475,21 +475,23 @@ void JFJochReceiverFPGA::FrameTransformationThread(uint32_t threadid) { serializer.SerializeImage(message); const size_t metadata_size = serializer.GetImageAppendOffset(); - const size_t max_image_size = experiment.GetMaxCompressedSize(); const size_t slot_size = experiment.GetImageBufferLocationSize(); - if (slot_size - metadata_size < max_image_size) { + + const auto compression_start_time = std::chrono::steady_clock::now(); + int64_t image_size = transformation.CompressImage(writer_buffer + serializer.GetImageAppendOffset(), + slot_size - (metadata_size + 32)); // keep 32 byte extra for close, etc. + if (image_size < 0) { // The per-image CBOR metadata (spot/reflection lists, ...) leaves // too little room for the compressed image in the buffer slot. Drop // just this frame instead of aborting the whole data collection. logger.Error("Dropping image {} - CBOR metadata too large to store image: " - "metadata {} B + image {} B do not fit in buffer slot {} B", - message.number, metadata_size, max_image_size, slot_size); + "metadata {} do not fit in buffer slot {} B", + message.number, metadata_size, slot_size); loc->release(); } else { - const auto compression_start_time = std::chrono::steady_clock::now(); - size_t image_size = transformation.CompressImage(writer_buffer + serializer.GetImageAppendOffset()); const auto compression_end_time = std::chrono::steady_clock::now(); - message.compression_time_s = std::chrono::duration(compression_end_time - compression_start_time).count(); + message.compression_time_s = std::chrono::duration( + compression_end_time - compression_start_time).count(); serializer.AppendImage(image_size); compressed_size += image_size; diff --git a/tests/ZSTDCompressorTest.cpp b/tests/ZSTDCompressorTest.cpp index 699f7ce5..0d0d4262 100644 --- a/tests/ZSTDCompressorTest.cpp +++ b/tests/ZSTDCompressorTest.cpp @@ -240,7 +240,7 @@ TEST_CASE("JFJochCompressor_JFJochDecompressor_ZSTD","[ZSTD]") { std::vector tmp(x.GetPixelsNum() * sizeof(int32_t) * 4 + 12); - auto tmp_size = compressor.Compress(tmp.data(), image); + auto tmp_size = compressor.Compress(tmp.data(), tmp.size(), image); tmp.resize(tmp_size); std::vector output; @@ -262,7 +262,7 @@ TEST_CASE("JFJochCompressor_JFJochDecompressor_LZ4","[ZSTD]") { std::vector tmp(x.GetPixelsNum() * sizeof(int32_t) * 4 + 12); - auto tmp_size = compressor.Compress(tmp.data(), image); + auto tmp_size = compressor.Compress(tmp.data(), tmp.size(), image); tmp.resize(tmp_size); std::vector output;