diff --git a/src/tree_math.cpp b/src/tree_math.cpp index 869f5b36..0c4ada97 100644 --- a/src/tree_math.cpp +++ b/src/tree_math.cpp @@ -14,7 +14,7 @@ log2(uint32_t x) } uint32_t k = 0; - while ((x >> k) > 0) { + while (k < 32 && (x >> k) > 0) { k += 1; } return k - 1; diff --git a/src/treekem.cpp b/src/treekem.cpp index 35899b7e..30f787e8 100644 --- a/src/treekem.cpp +++ b/src/treekem.cpp @@ -729,6 +729,47 @@ TreeKEMPublicKey::slice_path(UpdatePath path, void TreeKEMPublicKey::implant_slice_unchecked(const TreeSlice& slice) { + // Validate TreeSlice parameters before any tree math operations to prevent + // malformed inputs from triggering undefined behavior in log2() or other + // tree operations. A genuine slice must have consistent n_leaves with the + // node vectors it carries. + + // Check basic bounds + if (slice.n_leaves.val == 0 || slice.n_leaves.val > 100000) { + throw InvalidParameterError("Malformed tree slice (invalid n_leaves)"); + } + + if (slice.leaf_index.val >= slice.n_leaves.val) { + throw InvalidParameterError( + "Malformed tree slice (leaf index out of bounds)"); + } + + // Limit path sizes to prevent resource exhaustion + const auto max_depth = size_t(20); + if (slice.direct_path_nodes.size() > max_depth + 1 || + slice.copath_hashes.size() > max_depth) { + throw InvalidParameterError("Malformed tree slice (oversized paths)"); + } + + // The direct path of a leaf in a tree of n leaves has log2(NodeCount(n)) + // entries. Require the declared size to match the supplied vectors so a + // forged n_leaves cannot be larger than the evidence supplied for it. + { + const auto w = uint64_t{ 2 } * (uint64_t{ slice.n_leaves.val } - 1) + 1; + auto depth = size_t(0); + while ((w >> depth) > 1) { + depth += 1; + } + // dirpath includes the leaf itself after the insert below, so the + // serialized vector must have depth + 1 entries (or 1 for a + // single-leaf tree whose dirpath is empty). + const auto expected_dirpath = (slice.n_leaves.val == 1) ? 1 : depth + 1; + if (slice.direct_path_nodes.size() != expected_dirpath) { + throw InvalidParameterError( + "Malformed tree slice (direct path does not match declared size)"); + } + } + const auto n = NodeIndex(slice.leaf_index); auto dirpath = n.dirpath(size); dirpath.insert(dirpath.begin(), n); diff --git a/test/credential.cpp b/test/credential.cpp index bac4b2c5..cb067ade 100644 --- a/test/credential.cpp +++ b/test/credential.cpp @@ -170,7 +170,7 @@ TEST_CASE("X509 Credential EC certificates") const std::vector der_in{ cert }; auto key = SignaturePrivateKey::parse_der( - mls::CipherSuite::ID::P256_AES128GCM_SHA256_P256, keydata); + CipherSuite::ID::P256_AES128GCM_SHA256_P256, keydata); auto cred = Credential::x509(der_in); auto x509 = cred.get(); diff --git a/test/state.cpp b/test/state.cpp index de938edd..b9990d20 100644 --- a/test/state.cpp +++ b/test/state.cpp @@ -1592,8 +1592,7 @@ TEST_CASE_METHOD(StateTest, kp0.sign(id0); auto [init1, leaf1, id1, kp1] = make_client(); - // SAME app_id - kp1.leaf_node.extensions.add(ApplicationIDExtension{ app_id }); + kp1.leaf_node.extensions.add(ApplicationIDExtension{ app_id }); // SAME app_id kp1.leaf_node.sign(suite, id1, std::nullopt); kp1.sign(id1); @@ -1700,9 +1699,9 @@ TEST_CASE_METHOD( StateTest, "Allow Commit with Remove and Re-Add Same ApplicationIDExtension") { - // This test verifies that removing a member and adding a new member with - // the same ApplicationIDExtension in the same commit is allowed (e.g., - // device replacement) + // This test verifies that removing a member and adding a new member with the + // same ApplicationIDExtension in the same commit is allowed (e.g., device + // replacement) const auto app_id = from_ascii("reused-app-id"); @@ -1735,10 +1734,11 @@ TEST_CASE_METHOD( REQUIRE(first1.roster().size() == 2); REQUIRE(second1.roster().size() == 2); - // Create a replacement key package with the SAME ApplicationIDExtension - // as kp1 + // Create a replacement key package with the SAME ApplicationIDExtension as + // kp1 auto [init2, leaf2, id2, kp2] = make_client(); - kp2.leaf_node.extensions.add(ApplicationIDExtension{ different_app_id }); + kp2.leaf_node.extensions.add( + ApplicationIDExtension{ different_app_id }); // SAME as kp1 kp2.leaf_node.capabilities.extensions.push_back(ApplicationIDExtension::type); kp2.leaf_node.sign(suite, id2, std::nullopt); kp2.sign(id2); @@ -1747,8 +1747,8 @@ TEST_CASE_METHOD( auto remove = first1.remove_proposal(LeafIndex{ 1 }); auto add2 = first1.add_proposal(kp2); - // This should succeed because we're removing the old member before adding - // the new one + // This should succeed because we're removing the old member before adding the + // new one auto [commit2, welcome2, first2] = first1.commit( fresh_secret(), CommitOpts{ { remove, add2 }, true, false, {} }, {}); diff --git a/test/treekem.cpp b/test/treekem.cpp index 6bf18ea5..9e468953 100644 --- a/test/treekem.cpp +++ b/test/treekem.cpp @@ -306,3 +306,133 @@ TEST_CASE("TreeKEM Interop", "[.][all]") } } } + +TEST_CASE_METHOD(TreeKEMTest, + "Security: Reject malformed TreeSlice with zero n_leaves") +{ + TreeSlice slice; + slice.leaf_index = LeafIndex{ 0 }; + slice.n_leaves = LeafCount{ 0 }; + + REQUIRE_THROWS_AS(TreeKEMPublicKey(suite, slice), InvalidParameterError); +} + +TEST_CASE_METHOD(TreeKEMTest, + "Security: Reject malformed TreeSlice with excessive n_leaves") +{ + TreeSlice slice; + slice.leaf_index = LeafIndex{ 0 }; + slice.n_leaves = LeafCount{ 0x40000001u }; + + REQUIRE_THROWS_AS(TreeKEMPublicKey(suite, slice), InvalidParameterError); +} + +TEST_CASE_METHOD( + TreeKEMTest, + "Security: Reject malformed TreeSlice with n_leaves exceeding limit") +{ + TreeSlice slice; + slice.leaf_index = LeafIndex{ 0 }; + slice.n_leaves = LeafCount{ 100001 }; + + REQUIRE_THROWS_AS(TreeKEMPublicKey(suite, slice), InvalidParameterError); +} + +TEST_CASE_METHOD( + TreeKEMTest, + "Security: Accept well-formed TreeSlice at n_leaves limit boundary") +{ + const auto n_leaves = LeafCount{ 100000 }; + const auto leaf_index = LeafIndex{ 0 }; + const auto n = NodeIndex(leaf_index); + + auto dirpath = n.dirpath(n_leaves); + dirpath.insert(dirpath.begin(), n); + const auto copath = n.copath(n_leaves); + + TreeSlice slice; + slice.leaf_index = leaf_index; + slice.n_leaves = n_leaves; + slice.direct_path_nodes = std::vector(dirpath.size()); + slice.copath_hashes = std::vector(copath.size()); + + REQUIRE_NOTHROW(TreeKEMPublicKey(suite, slice)); +} + +TEST_CASE_METHOD( + TreeKEMTest, + "Security: Reject malformed TreeSlice with leaf_index out of bounds") +{ + TreeSlice slice; + slice.leaf_index = LeafIndex{ 10 }; + slice.n_leaves = LeafCount{ 5 }; + + REQUIRE_THROWS_AS(TreeKEMPublicKey(suite, slice), InvalidParameterError); +} + +TEST_CASE_METHOD( + TreeKEMTest, + "Security: Reject malformed TreeSlice with oversized direct_path") +{ + TreeSlice slice; + slice.leaf_index = LeafIndex{ 0 }; + slice.n_leaves = LeafCount{ 10 }; + + for (size_t i = 0; i < 22; ++i) { + slice.direct_path_nodes.push_back(OptionalNode{}); + } + + REQUIRE_THROWS_AS(TreeKEMPublicKey(suite, slice), InvalidParameterError); +} + +TEST_CASE_METHOD(TreeKEMTest, + "Security: Reject malformed TreeSlice with oversized copath") +{ + TreeSlice slice; + slice.leaf_index = LeafIndex{ 0 }; + slice.n_leaves = LeafCount{ 10 }; + + for (size_t i = 0; i < 21; ++i) { + slice.copath_hashes.push_back(bytes{}); + } + + REQUIRE_THROWS_AS(TreeKEMPublicKey(suite, slice), InvalidParameterError); +} + +TEST_CASE_METHOD(TreeKEMTest, + "Security: Reject TreeSlice with mismatched direct_path size") +{ + TreeSlice slice; + slice.leaf_index = LeafIndex{ 0 }; + slice.n_leaves = LeafCount{ 100 }; + + for (size_t i = 0; i < 5; ++i) { + slice.direct_path_nodes.push_back(OptionalNode{}); + } + + REQUIRE_THROWS_AS(TreeKEMPublicKey(suite, slice), InvalidParameterError); +} + +TEST_CASE("Security: log2 does not infinite loop on large values") +{ + auto root_index = NodeIndex::root(LeafCount{ 1000 }); + REQUIRE(root_index.val > 0); + + auto root_index2 = NodeIndex::root(LeafCount{ 99999 }); + REQUIRE(root_index2.val > 0); +} + +TEST_CASE("Security: log2 does not infinite loop when the shift count " + "reaches the type width") +{ + // NodeCount(LeafCount{0x40000001}) is 0x80000001, which has bit 31 set. + // Before the fix, log2() kept shifting past bit width 32, and on + // platforms where shifting a uint32_t by 32 wraps around (undefined + // behavior treated as shift-by-0), this caused an infinite loop. + auto root_index = NodeIndex::root(LeafCount{ 0x40000001u }); + REQUIRE(root_index.val > 0); + + // The theoretical maximum, to exercise the boundary case directly. + auto root_index2 = NodeIndex::root(LeafCount{ 0xFFFFFFFFu }); + REQUIRE(root_index2.val > 0); +}