1 /*
2  * Copyright (c) 2024-2024 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 
39 private:
40     static const int FSVERITY_HASH_PAGE_SIZE;
41     static const int64_t INPUTSTREAM_MAX_SIZE;
42     static const int CHUNK_SIZE;
43     static const long MAX_READ_SIZE;
44     static const int MAX_PROCESSORS;
45     static const int BLOCKINGQUEUE;
46     static ByteBuffer* Slice(ByteBuffer* buffer, int begin, int end);
47     static std::vector<int64_t> GetOffsetArrays(long dataSize, int digestSize);
48     static std::vector<long> GetLevelSize(long dataSize, int digestSize);
49     static long GetChunkCount(long dataSize, long divisor);
50     static long GetFullChunkSize(long dataSize, long divisor, long multiplier);
51     void SetAlgorithm(const std::string& algorithm);
52     void TransInputStreamToHashData(std::istream& inputStream, long size,
53                                     ByteBuffer* outputBuffer, int bufStartIdx);
54     std::vector<std::vector<int8_t>> GetDataHashes(std::istream& inputStream, long size);
55 
56     void RunHashTask(std::vector<std::vector<int8_t>>& hashes, ByteBuffer* buffer,
57                      int readChunkIndex, int bufStartIdx);
58     void TransInputDataToHashData(ByteBuffer* inputBuffer, ByteBuffer* outputBuffer,
59                                   int64_t bufStartIdx, int64_t outputStartIdx);
60     void GenerateHashDataByInputData(std::istream& inputStream, long size, ByteBuffer* outputBuffer,
61                                      std::vector<int64_t>& offsetArrays, int digestSize);
62     void GenerateHashDataByHashData(ByteBuffer* buffer, std::vector<int64_t>& offsetArrays, int digestSize);
63     MerkleTree* GetMerkleTree(ByteBuffer* dataBuffer, long inputDataSize,
64                               FsVerityHashAlgorithm fsVerityHashAlgorithm);
65     void DataRoundupChunkSize(ByteBuffer* data, long originalDataSize, int digestSize);
66     bool CheckCalculateHashResult;
67     const int POOL_SIZE = std::min(MAX_PROCESSORS, static_cast<int>(std::thread::hardware_concurrency()));
68     std::string mAlgorithm = "SHA-256";
69     std::shared_ptr<Uscript::ThreadPool> mPools;
70 };
71 } // namespace SignatureTools
72 } // namespace OHOS
73 #endif // SIGNATURETOOLS_MERKLE_TREE_BUILDER_H