1 /* 2 * Copyright (c) 2025-2025 Huawei Device Co., Ltd. 3 * Licensed under the Apache License, Version 2.0 (the "License"); 4 * you may not use this file except in compliance with the License. 5 * You may obtain a copy of the License at 6 * 7 * http://www.apache.org/licenses/LICENSE-2.0 8 * 9 * Unless required by applicable law or agreed to in writing, software 10 * distributed under the License is distributed on an "AS IS" BASIS, 11 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 * See the License for the specific language governing permissions and 13 * limitations under the License. 14 */ 15 #ifndef SIGNATURETOOLS_MERKLE_TREE_BUILDER_H 16 #define SIGNATURETOOLS_MERKLE_TREE_BUILDER_H 17 18 #include <algorithm> 19 #include <string> 20 #include <vector> 21 #include <iostream> 22 #include <memory> 23 #include <sstream> 24 25 #include "thread_pool.h" 26 #include "byte_buffer.h" 27 #include "merkle_tree.h" 28 #include "fs_digest_utils.h" 29 #include "signature_tools_log.h" 30 31 namespace OHOS { 32 namespace SignatureTools { 33 class MerkleTreeBuilder { 34 public: 35 MerkleTreeBuilder(); 36 MerkleTree* GenerateMerkleTree(std::istream& inputStream, long size, 37 const FsVerityHashAlgorithm& fsVerityHashAlgorithm); 38 SetCsOffset(uint64_t & csOffset)39 void SetCsOffset(uint64_t& csOffset) 40 { 41 this->csOffset = csOffset; 42 } 43 44 private: 45 static const int FSVERITY_HASH_PAGE_SIZE; 46 static const int64_t INPUTSTREAM_MAX_SIZE; 47 static const int CHUNK_SIZE; 48 static const long MAX_READ_SIZE; 49 static const int MAX_PROCESSORS; 50 static const int BLOCKINGQUEUE; 51 static ByteBuffer* Slice(ByteBuffer* buffer, int begin, int end); 52 static std::vector<int64_t> GetOffsetArrays(long dataSize, int digestSize); 53 static std::vector<long> GetLevelSize(long dataSize, int digestSize); 54 static long GetChunkCount(long dataSize, long divisor); 55 static long GetFullChunkSize(long dataSize, long divisor, long multiplier); 56 void SetAlgorithm(const std::string& algorithm); 57 void TransInputStreamToHashData(std::istream& inputStream, long size, 58 ByteBuffer* outputBuffer, int bufStartIdx); 59 std::vector<std::vector<int8_t>> GetDataHashes(std::istream& inputStream, long size); 60 61 void RunHashTask(std::vector<std::vector<int8_t>>& hashes, ByteBuffer* buffer, 62 int readChunkIndex, int bufStartIdx); 63 void TransInputDataToHashData(ByteBuffer* inputBuffer, ByteBuffer* outputBuffer, 64 int64_t bufStartIdx, int64_t outputStartIdx); 65 void GenerateHashDataByInputData(std::istream& inputStream, long size, ByteBuffer* outputBuffer, 66 std::vector<int64_t>& offsetArrays, int digestSize); 67 void GenerateHashDataByHashData(ByteBuffer* buffer, std::vector<int64_t>& offsetArrays, int digestSize); 68 MerkleTree* GetMerkleTree(ByteBuffer* dataBuffer, long inputDataSize, 69 FsVerityHashAlgorithm fsVerityHashAlgorithm); 70 void DataRoundupChunkSize(ByteBuffer* data, long originalDataSize, int digestSize); 71 bool CheckCalculateHashResult; 72 const int POOL_SIZE = std::min(MAX_PROCESSORS, static_cast<int>(std::thread::hardware_concurrency())); 73 std::string mAlgorithm = "SHA-256"; 74 std::shared_ptr<Uscript::ThreadPool> mPools; 75 uint64_t csOffset = 0; 76 }; 77 } // namespace SignatureTools 78 } // namespace OHOS 79 #endif // SIGNATURETOOLS_MERKLE_TREE_BUILDER_H