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
31namespace OHOS {
32namespace SignatureTools {
33class MerkleTreeBuilder {
34public:
35    MerkleTreeBuilder();
36    MerkleTree* GenerateMerkleTree(std::istream& inputStream, long size,
37                                   const FsVerityHashAlgorithm& fsVerityHashAlgorithm);
38
39private:
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