From b71f6867c58383b6eb026602589eaa49b91b6ad5 Mon Sep 17 00:00:00 2001 From: yyc12345 Date: Wed, 11 Dec 2024 20:19:48 +0800 Subject: [PATCH] fix: update NlpCodec --- NlpCodec/NlpCodec.cpp | 79 +++++++++++++++++++++++-------------------- 1 file changed, 43 insertions(+), 36 deletions(-) diff --git a/NlpCodec/NlpCodec.cpp b/NlpCodec/NlpCodec.cpp index 4dfbf6f..f9effcd 100644 --- a/NlpCodec/NlpCodec.cpp +++ b/NlpCodec/NlpCodec.cpp @@ -13,6 +13,8 @@ namespace NlpCodec { +#pragma region Help Structs and Functions + /// @brief NlpCodec universal exception. /// @details Once this exception was thrown, it means that somethings went wrong. /// and main function should catch it, output error message and exit program immediately. @@ -32,12 +34,29 @@ namespace NlpCodec { static constexpr _TyTo SafeCast(_TyFrom value) { if (!std::in_range<_TyTo>(value)) throw NlpException( - "Fail to cast integral number because given value is greater than container. " + "Fail to cast integral number because given value is greater than the type can hold. " "This is usually caused by your input or output file is too long."); return static_cast<_TyTo>(value); } -#pragma region "Encryption Stuff" { + /// @brief The safe version of `std::ifstream::read`. + /// Throw exception if fail to read. + static void SafeRead(std::ifstream& fin, char* s, std::streamsize count) { + fin.read(s, count); + if (!fin.good() || fin.gcount() != count) + throw NlpException("Fail to read data from file."); + } + /// @brief The safe version of `std::ofstream::write`. + /// Throw exception if fail to write. + static void SafeWrite(std::ofstream& fout, const char* s, std::streamsize count) { + fout.write(s, count); + if (!fout.good()) + throw NlpException("Fail to write data into file."); + } + +#pragma endregion + +#pragma region Encryption Stuff /* @@ -155,7 +174,7 @@ namespace NlpCodec { } } -#pragma endregion } +#pragma endregion /// @brief Get the length of given file stream. static uint32_t GetFileLength(std::ifstream& fin) { @@ -195,9 +214,7 @@ namespace NlpCodec { throw NlpException("Fail to allocate memory."); // Read data from file to input buffer - fin.read(inbuf.get(), raw_size); - if (!fin.good() || fin.gcount() != raw_size) - throw NlpException("Fail to read file data into buffer."); + SafeRead(fin, inbuf.get(), raw_size); // Do XOR operation CircularXorOperation(inbuf.get(), raw_size); @@ -209,9 +226,9 @@ namespace NlpCodec { reinterpret_cast(inbuf.get()), static_cast(raw_size), Z_BEST_COMPRESSION ); - // Check ZLib result. + // Check zlib result. if (ret != Z_OK) - throw NlpException("Zlib compress() failed."); + throw NlpException("zlib compress() failed."); // Fetch final compressed size. uint32_t compressed_size = SafeCast(dest_len); @@ -219,21 +236,15 @@ namespace NlpCodec { uint32_t checksum = static_cast(adler32(0u, reinterpret_cast(outbuf.get()), SafeCast(compressed_size))); // Write compressed data into file - fout.write(outbuf.get(), compressed_size); - if (!fout.good()) - throw NlpException("Fail to write data into file."); + SafeWrite(fout, outbuf.get(), compressed_size); // Raw size and checksum need some extra encryption before writting raw_size = EncryptRawFileLength(raw_size); checksum = EncryptChecksum(checksum); // Write raw size and checksum - fout.write(reinterpret_cast(&raw_size), sizeof(uint32_t)); - if (!fout.good()) - throw NlpException("Fail to write raw size into file."); - fout.write(reinterpret_cast(&checksum), sizeof(uint32_t)); - if (!fout.good()) - throw NlpException("Fail to write checksum into file."); + SafeWrite(fout, reinterpret_cast(&raw_size), sizeof(uint32_t)); + SafeWrite(fout, reinterpret_cast(&checksum), sizeof(uint32_t)); } @@ -247,8 +258,8 @@ namespace NlpCodec { compressed_size -= TAIL_SIZE; fin.seekg(compressed_size, std::ios_base::beg); uint32_t expected_raw_size = 0u, expected_checksum = 0u; - fin.read(reinterpret_cast(&expected_raw_size), sizeof(uint32_t)); - fin.read(reinterpret_cast(&expected_checksum), sizeof(uint32_t)); + SafeRead(fin, reinterpret_cast(&expected_raw_size), sizeof(uint32_t)); + SafeRead(fin, reinterpret_cast(&expected_checksum), sizeof(uint32_t)); fin.seekg(0, std::ios_base::beg); // Raw size and checksum data need to do some extra decryption. @@ -262,9 +273,7 @@ namespace NlpCodec { throw NlpException("Fail to allocate memory."); // Read file into buffer - fin.read(inbuf.get(), compressed_size); - if (!fin.good() || fin.gcount() != compressed_size) - throw NlpException("Fail to read data into buffer.\n"); + SafeRead(fin, inbuf.get(), compressed_size); // Test checksum uint32_t checksum = static_cast(adler32(0u, reinterpret_cast(inbuf.get()), SafeCast(compressed_size))); @@ -281,15 +290,13 @@ namespace NlpCodec { ); // Check zlib result if (ret != Z_OK) - throw NlpException("Zlib uncompress() failed."); + throw NlpException("zlib uncompress() failed."); - // do xor operation + // Do XOR operation CircularXorOperation(outbuf.get(), expected_raw_size); // Write result into file - fout.write(outbuf.get(), expected_raw_size); - if (!fout.good()) - throw NlpException("Fail to write data into file."); + SafeWrite(fout, outbuf.get(), expected_raw_size); } @@ -316,16 +323,16 @@ namespace NlpCodec::Runtime { << "NlpCodec [encode | decode | version | help] " << std::endl << std::endl << "version - print version info about this program." << std::endl - << "help - print this page." << std::endl + << "help - print this page." << std::endl << std::endl - << "encode - encode text file into NLP file." << std::endl - << "decode - decode NLP file into text file." << std::endl - << " - the source file." << std::endl - << " the path to text file in encode mode." << std::endl - << " the path to NLP file in decode mode." << std::endl - << " - the destination file." << std::endl - << " the path to NLP file in encode mode." << std::endl - << " the path to text file in decode mode." << std::endl + << "encode - encode text file into NLP file." << std::endl + << "decode - decode NLP file into text file." << std::endl + << " - the source file." << std::endl + << " encode mode: the path to text file." << std::endl + << " decode mode: the path to NLP file." << std::endl + << " - the destination file." << std::endl + << " encode mode: the path to NLP file." << std::endl + << " decode mode: the path to text file." << std::endl << "" << std::endl; }