| // Copyright 2026 Google LLC |
| // |
| // Licensed under the Apache License, Version 2.0 (the "License"); |
| // you may not use this file except in compliance with the License. |
| // You may obtain a copy of the License at |
| // |
| // https://www.apache.org/licenses/LICENSE-2.0 |
| // |
| // Unless required by applicable law or agreed to in writing, software |
| // distributed under the License is distributed on an "AS IS" BASIS, |
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| // See the License for the specific language governing permissions and |
| // limitations under the License. |
| // |
| |
| #include "bindings/merkle.h" |
| |
| #include <cstdint> |
| #include <memory> |
| #include <optional> |
| #include <string> |
| #include <utility> |
| #include <vector> |
| |
| #include "absl/status/status.h" |
| #include "absl/status/status_macros.h" |
| #include "absl/status/statusor.h" |
| #include "absl/strings/str_format.h" |
| #include "absl/strings/string_view.h" |
| #include "crypto/hash.h" |
| #include "google/protobuf/repeated_field.h" |
| #include "proto/bmff_based_hash_assertion.pb.h" |
| #include "riegeli/bytes/reader.h" |
| |
| namespace credentio { |
| |
| namespace { |
| |
| const int64_t kMaxPowerOfTwo = 0x8000'0000; |
| const int64_t kMaxBlockCount = 1048576; |
| |
| absl::StatusOr<int64_t> NextPowerOfTwo(int64_t value) { |
| if ((value & (value - 1)) == 0) { |
| // value is already a power of two |
| return value; |
| } |
| if (value < 0 || value >= kMaxPowerOfTwo) { |
| return absl::InvalidArgumentError( |
| absl::StrFormat("Value is negative or too large: %d", value)); |
| } |
| int64_t next_power_of_two = 1; |
| while (next_power_of_two < value) { |
| next_power_of_two <<= 1; |
| } |
| return next_power_of_two; |
| } |
| |
| absl::Status PadToPowerOfTwo(std::vector<std::string>& entries) { |
| ABSL_ASSIGN_OR_RETURN(auto next_power_of_two, NextPowerOfTwo(entries.size())); |
| if (next_power_of_two > entries.size()) { |
| entries.resize(next_power_of_two, ""); |
| } |
| return absl::OkStatus(); |
| } |
| |
| absl::StatusOr<std::vector<std::string>> ReduceLayer( |
| HasherFactory* factory, std::vector<std::string> layer) { |
| std::vector<std::string> reduced_layer; |
| for (int64_t i = 0; i + 1 < layer.size(); i += 2) { |
| ABSL_ASSIGN_OR_RETURN(auto hasher, factory->Create()); |
| hasher->Update(layer[i]); |
| hasher->Update(layer[i + 1]); |
| reduced_layer.push_back(hasher->Digest()); |
| } |
| return reduced_layer; |
| } |
| |
| absl::StatusOr<MerkleTree> PopulateLayers( |
| HasherFactory* factory, std::vector<std::string> leaf_hashes) { |
| MerkleTree tree = {leaf_hashes}; |
| while (tree[tree.size() - 1].size() > 1) { |
| ABSL_ASSIGN_OR_RETURN(auto reduced_layer, |
| ReduceLayer(factory, tree.back())); |
| tree.push_back(std::move(reduced_layer)); |
| } |
| return tree; |
| } |
| |
| // Returns the number of rows between leaf_row_size and hashes_row_size. |
| // For example, if leaf_row_size is 16 and hashes_row_size is 4, the function |
| // will return 4. |
| absl::StatusOr<int64_t> RowsAway(int64_t leaf_row_size, |
| int64_t hashes_row_size) { |
| if (leaf_row_size <= 0 || hashes_row_size <= 0) { |
| return absl::InvalidArgumentError( |
| "leaf_row_size and hashes_row_size must be positive"); |
| } |
| if (hashes_row_size > leaf_row_size) { |
| return absl::InvalidArgumentError( |
| "hashes row size is larger than the leaf row size"); |
| } |
| ABSL_ASSIGN_OR_RETURN(auto leaf_pow2, NextPowerOfTwo(leaf_row_size)); |
| ABSL_ASSIGN_OR_RETURN(auto hashes_pow2, NextPowerOfTwo(hashes_row_size)); |
| int64_t rows_away = 0; |
| while (leaf_pow2 != hashes_pow2) { |
| leaf_pow2 >>= 1; |
| rows_away++; |
| } |
| return rows_away; |
| } |
| |
| absl::StatusOr<std::vector<int64_t>> GetFixedBlockSizes( |
| uint64_t mdat_box_size, int64_t fixed_block_size) { |
| if (fixed_block_size <= 0) { |
| return absl::InvalidArgumentError( |
| "merkle map has an invalid fixed block size"); |
| } |
| |
| uint64_t num_blocks = mdat_box_size / fixed_block_size; |
| if (mdat_box_size % fixed_block_size != 0) { |
| num_blocks++; |
| } |
| |
| if (num_blocks > kMaxBlockCount) { |
| return absl::InvalidArgumentError(absl::StrFormat( |
| "merkle map block count exceeds maximum allowable limit: %d > %d", |
| num_blocks, kMaxBlockCount)); |
| } |
| |
| std::vector<int64_t> block_sizes; |
| block_sizes.reserve(num_blocks); |
| block_sizes.insert(block_sizes.end(), mdat_box_size / fixed_block_size, |
| fixed_block_size); |
| if (mdat_box_size % fixed_block_size != 0) { |
| block_sizes.push_back(mdat_box_size % fixed_block_size); |
| } |
| return block_sizes; |
| } |
| |
| absl::StatusOr<std::vector<int64_t>> GetVariableBlockSizes( |
| uint64_t mdat_box_size, |
| const google::protobuf::RepeatedField<int64_t>& variable_block_sizes) { |
| if (variable_block_sizes.size() > kMaxBlockCount) { |
| return absl::InvalidArgumentError(absl::StrFormat( |
| "merkle map block count exceeds maximum allowable limit: %d > %d", |
| variable_block_sizes.size(), kMaxBlockCount)); |
| } |
| |
| int64_t total_block_size = 0; |
| std::vector<int64_t> block_sizes; |
| block_sizes.reserve(variable_block_sizes.size()); |
| for (int64_t block_size : variable_block_sizes) { |
| if (block_size <= 0) { |
| return absl::InvalidArgumentError( |
| "merkle map has an invalid variable block size"); |
| } |
| if (total_block_size > INT64_MAX - block_size) { |
| return absl::InvalidArgumentError( |
| "merkle map's variable block sizes overflow int64_t"); |
| } |
| total_block_size += block_size; |
| block_sizes.push_back(block_size); |
| } |
| if (total_block_size != mdat_box_size) { |
| return absl::InvalidArgumentError( |
| "merkle map's variable block sizes do not sum to the mdat box size"); |
| } |
| return block_sizes; |
| } |
| |
| absl::StatusOr<std::vector<int64_t>> GetSingleBoxSizes(uint64_t mdat_box_size) { |
| if (mdat_box_size > INT64_MAX) { |
| return absl::InvalidArgumentError( |
| "mdat atom's box size is too large to fit in an int64_t"); |
| } |
| return std::vector<int64_t>{static_cast<int64_t>(mdat_box_size)}; |
| } |
| |
| } // namespace |
| |
| absl::StatusOr<std::vector<int64_t>> DeriveMerkleBlockSizes( |
| const BmffMerkle& merkle, uint64_t mdat_box_size) { |
| if (merkle.has_fixed_block_size() && merkle.variable_block_sizes_size() > 0) { |
| return absl::InvalidArgumentError( |
| "merkle map has both fixed and variable block sizes"); |
| } |
| |
| std::vector<int64_t> block_sizes; |
| if (merkle.has_fixed_block_size()) { |
| int64_t fixed_block_size = merkle.fixed_block_size(); |
| if (fixed_block_size <= 0) { |
| return absl::InvalidArgumentError( |
| "merkle map has an invalid fixed block size"); |
| } |
| uint64_t num_blocks = mdat_box_size / fixed_block_size; |
| if (mdat_box_size % fixed_block_size != 0) { |
| num_blocks++; |
| } |
| if (num_blocks != merkle.count()) { |
| return absl::InvalidArgumentError( |
| "merkle map's block sizes do not match the expected leaf count"); |
| } |
| |
| ABSL_ASSIGN_OR_RETURN( |
| block_sizes, |
| GetFixedBlockSizes(mdat_box_size, merkle.fixed_block_size())); |
| } else if (merkle.variable_block_sizes_size() > 0) { |
| ABSL_ASSIGN_OR_RETURN( |
| block_sizes, |
| GetVariableBlockSizes(mdat_box_size, merkle.variable_block_sizes())); |
| } else { |
| ABSL_ASSIGN_OR_RETURN(block_sizes, GetSingleBoxSizes(mdat_box_size)); |
| } |
| |
| if (block_sizes.size() != merkle.count()) { |
| return absl::InvalidArgumentError( |
| "merkle map's block sizes do not match the expected leaf count"); |
| } |
| |
| return block_sizes; |
| } |
| |
| absl::StatusOr<DerivedTreeData> DeriveMerkleTreeData(int64_t leaf_count, |
| int64_t hashes_count, |
| int64_t auxiliary_count) { |
| ABSL_ASSIGN_OR_RETURN(auto leaf_pow2, NextPowerOfTwo(leaf_count)); |
| ABSL_ASSIGN_OR_RETURN(auto hashes_pow2, NextPowerOfTwo(hashes_count)); |
| ABSL_ASSIGN_OR_RETURN(auto delta_rows, RowsAway(leaf_pow2, hashes_pow2)); |
| int64_t expected_hashes_count = leaf_count; |
| for (int64_t r = 0; r < delta_rows; ++r) { |
| expected_hashes_count = (expected_hashes_count + 1) >> 1; |
| } |
| |
| if (hashes_count != expected_hashes_count) { |
| return absl::InvalidArgumentError(absl::StrFormat( |
| "hashes count is not equal to the expected hashes count: " |
| "Hashes Count: %d, Leaf Count: %d, Expected Hashes Count: %d", |
| hashes_count, leaf_count, expected_hashes_count)); |
| } |
| |
| if ((delta_rows > 0 || auxiliary_count > 0) && |
| auxiliary_count != leaf_count) { |
| return absl::InvalidArgumentError( |
| absl::StrFormat("incorrect count of auxiliary data boxes: " |
| "Auxiliary Data Count: %d, Leaf Count: %d", |
| auxiliary_count, leaf_count)); |
| } |
| |
| return DerivedTreeData{ |
| .full_leaf_count = leaf_pow2, |
| .full_hashes_count = hashes_pow2, |
| .delta_rows = delta_rows, |
| }; |
| } |
| |
| absl::StatusOr<MerkleTree> CreateMerkleTree(HasherFactory* factory, |
| riegeli::Reader& input, |
| int64_t starting_offset, |
| int64_t length, |
| CreateMerkleTreeOptions options) { |
| if (options.fixed_block_size.has_value() && |
| !options.variable_block_sizes.empty()) { |
| return absl::InvalidArgumentError( |
| "Must specify a fixed block size or variable block sizes, not both."); |
| } |
| |
| if (!options.fixed_block_size.has_value() && |
| options.variable_block_sizes.empty()) { |
| // The entire input is a single leaf node. |
| return CreateMerkleTree(factory, input, starting_offset, length, |
| std::vector<int64_t>{length}); |
| } |
| |
| if (options.fixed_block_size.has_value()) { |
| if (options.fixed_block_size.value() <= 0) { |
| return absl::InvalidArgumentError( |
| absl::StrFormat("Fixed block size must be greater than 0. Found: %d", |
| options.fixed_block_size.value())); |
| } |
| std::vector<int64_t> block_sizes(length / options.fixed_block_size.value(), |
| options.fixed_block_size.value()); |
| if (length % options.fixed_block_size.value() != 0) { |
| block_sizes.push_back(length % options.fixed_block_size.value()); |
| } |
| return CreateMerkleTree(factory, input, starting_offset, length, |
| block_sizes); |
| } |
| |
| return CreateMerkleTree(factory, input, starting_offset, length, |
| options.variable_block_sizes); |
| } |
| |
| absl::StatusOr<MerkleTree> CreateMerkleTree(HasherFactory* factory, |
| riegeli::Reader& input, |
| int64_t starting_offset, |
| int64_t length, |
| std::vector<int64_t> block_sizes) { |
| int64_t total_blocks_length = 0; |
| std::vector<std::string> leaf_hashes; |
| leaf_hashes.reserve(block_sizes.size()); |
| if (!input.Seek(starting_offset) || input.pos() != starting_offset) { |
| return input.StatusOrAnnotate( |
| absl::DataLossError("Failed to seek to starting offset")); |
| } |
| for (int64_t block_size : block_sizes) { |
| if (block_size <= 0) { |
| return absl::InvalidArgumentError(absl::StrFormat( |
| "Block size must be greater than 0. Found: %d", block_size)); |
| } |
| total_blocks_length += block_size; |
| if (total_blocks_length > length) { |
| return absl::InvalidArgumentError(absl::StrFormat( |
| "Total block size is too large. Total Block size: %d, Length: %d", |
| total_blocks_length, length)); |
| } |
| std::string leaf_data; |
| if (!input.Read(block_size, leaf_data)) { |
| return input.StatusOrAnnotate( |
| absl::DataLossError("Failed to read leaf data")); |
| } |
| ABSL_ASSIGN_OR_RETURN(auto hasher, factory->Create()); |
| hasher->Update(leaf_data); |
| leaf_hashes.push_back(hasher->Digest()); |
| } |
| ABSL_RETURN_IF_ERROR(PadToPowerOfTwo(leaf_hashes)); |
| |
| return PopulateLayers(factory, leaf_hashes); |
| } |
| |
| } // namespace credentio |
