To fix TreeSlice declared leaf count drives unterminating loop before authentication

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;

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);

test/credential.cpp

@@ -170,7 +170,7 @@ TEST_CASE("X509 Credential EC certificates")
   const std::vector<bytes> 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<X509Credential>();

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, {} }, {});
 

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<OptionalNode>(dirpath.size());
+  slice.copath_hashes = std::vector<bytes>(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);
+}