Fix failure on light client rejoining

lib/mls_ds/src/tree_follower.cpp

@@ -145,8 +145,7 @@ TreeFollower::update(const MLSMessage& commit_message,
   const auto commit_auth_content =
     commit_public_message.authenticated_content();
   const auto group_content = commit_auth_content.content;
-  const auto& commit =
-    var::get<Commit>(commit_auth_content.content.content);
+  const auto& commit = var::get<Commit>(commit_auth_content.content.content);
 
   // Apply proposals
   apply(_tree, _suite, group_content.sender, commit.proposals, extra_proposals);
@@ -155,8 +154,7 @@ TreeFollower::update(const MLSMessage& commit_message,
 
   // Merge the update path
   if (commit.path) {
-    const auto sender =
-      var::get<MemberSender>(group_content.sender.sender);
+    const auto sender = var::get<MemberSender>(group_content.sender.sender);
     const auto from = LeafIndex(sender.sender);
     const auto& path = opt::get(commit.path);
     _tree.merge(from, path);

src/treekem.cpp

@@ -100,7 +100,7 @@ TreeKEMPrivateKey::joiner(const TreeKEMPublicKey& pub,
   auto priv = TreeKEMPrivateKey{ pub.suite, index, {}, {}, {} };
   priv.private_key_cache.insert({ NodeIndex(index), std::move(leaf_priv) });
   if (path_secret) {
-    priv.implant(pub, intersect, opt::get(path_secret));
+    priv.implant_matching(pub, intersect, opt::get(path_secret));
   }
   return priv;
 }

test/state.cpp

@@ -511,6 +511,62 @@ TEST_CASE_METHOD(StateTest, "Light client can participate")
   REQUIRE(first4 == third4);
 }
 
+TEST_CASE_METHOD(StateTest, "Light client can rejoin")
+{
+  // Initialize the creator's state
+  auto first0 = State{ group_id,
+                       suite,
+                       leaf_privs[0],
+                       identity_privs[0],
+                       key_packages[0].leaf_node,
+                       {} };
+
+  // Add the second and third participants
+  auto add1a = first0.add_proposal(key_packages[1]);
+  auto add1b = first0.add_proposal(key_packages[2]);
+  auto [commit1, welcome1, first1_] = first0.commit(
+    fresh_secret(), CommitOpts{ { add1a, add1b }, true, false, {} }, {});
+  silence_unused(commit1);
+  auto first1 = first1_;
+
+  auto third1 = State{ init_privs[2],
+                       leaf_privs[2],
+                       identity_privs[2],
+                       key_packages[2],
+                       welcome1,
+                       std::nullopt,
+                       {} };
+
+  REQUIRE(first1 == third1);
+
+  // Remove the second participant and re-add them in the same commit
+  auto remove2 = first1.remove_proposal(LeafIndex{ 1 });
+  auto add2 = first1.add_proposal(key_packages[1]);
+  auto [commit2, welcome2, first2_] = first1.commit(
+    fresh_secret(), CommitOpts{ { remove2, add2 }, false, false, {} }, {});
+  silence_unused(welcome2);
+  auto first2 = first2_;
+
+  auto third2 = opt::get(third1.handle(commit2));
+
+  REQUIRE(first2 == third2);
+
+  // Second participant (re-)joins as a light client
+  const auto annotated_welcome_2 = AnnotatedWelcome::from(
+    welcome2, first2.tree(), LeafIndex{ 0 }, LeafIndex{ 1 });
+
+  auto second2 = State{ init_privs[1],
+                        leaf_privs[1],
+                        identity_privs[1],
+                        key_packages[1],
+                        annotated_welcome_2.welcome,
+                        annotated_welcome_2.tree(),
+                        {} };
+
+  REQUIRE(first2 == second2);
+  REQUIRE_FALSE(second2.is_full_client());
+}
+
 TEST_CASE_METHOD(StateTest, "Light client can upgrade after several commits")
 {
   // Initialize the first two users

test/treekem.cpp

@@ -61,8 +61,8 @@ TEST_CASE_METHOD(TreeKEMTest, "Node public key")
 TEST_CASE_METHOD(TreeKEMTest, "TreeKEM Private Key")
 {
   const auto size = LeafCount{ 5 };
-  const auto index = LeafIndex{ 2 };
-  const auto intersect = NodeIndex{ 3 };
+  const auto adder_index = LeafIndex{ 1 };
+  const auto joiner_index = LeafIndex{ 2 };
   const auto random = random_bytes(32);
   const auto priv = HPKEPrivateKey::derive(suite, random);
   const auto priv2 = HPKEPrivateKey::generate(suite);
@@ -75,22 +75,27 @@ TEST_CASE_METHOD(TreeKEMTest, "TreeKEM Private Key")
     pub.add_leaf(leaf_node);
   }
 
-  // create() populates the direct path
-  auto priv_create = TreeKEMPrivateKey::create(pub, index, random);
-  REQUIRE(priv_create.path_secrets.find(NodeIndex(4)) !=
+  // create() populates the direct path in the private key
+  auto priv_create = TreeKEMPrivateKey::create(pub, adder_index, random);
+  REQUIRE(priv_create.path_secrets.find(NodeIndex(2)) !=
           priv_create.path_secrets.end());
-  REQUIRE(priv_create.path_secrets.find(NodeIndex(5)) !=
+  REQUIRE(priv_create.path_secrets.find(NodeIndex(1)) !=
           priv_create.path_secrets.end());
   REQUIRE(priv_create.path_secrets.find(NodeIndex(3)) !=
           priv_create.path_secrets.end());
   REQUIRE(priv_create.path_secrets.find(NodeIndex(7)) !=
           priv_create.path_secrets.end());
   REQUIRE(priv_create.update_secret.size() == hash_size);
 
+  // Populate the direct path from the  into the public key manually
+  pub.node_at(NodeIndex{ 1 }) = OptionalNode{ Node{ ParentNode{} } };
+  pub.node_at(NodeIndex{ 3 }) = OptionalNode{ Node{ ParentNode{} } };
+  pub.node_at(NodeIndex{ 7 }) = OptionalNode{ Node{ ParentNode{} } };
+
   // joiner() populates the leaf and the path above the ancestor,
   // but not the direct path in the middle
-  auto priv_joiner =
-    TreeKEMPrivateKey::joiner(pub, index, priv, intersect, random);
+  auto priv_joiner = TreeKEMPrivateKey::joiner(
+    pub, joiner_index, priv, joiner_index.ancestor(adder_index), random);
   REQUIRE(priv_joiner.private_key(NodeIndex(4)));
   REQUIRE(priv_joiner.path_secrets.find(NodeIndex(3)) !=
           priv_joiner.path_secrets.end());
@@ -103,7 +108,7 @@ TEST_CASE_METHOD(TreeKEMTest, "TreeKEM Private Key")
 
   // set_leaf_priv() properly sets the leaf secret
   priv_joiner.set_leaf_priv(priv2);
-  REQUIRE(priv_joiner.private_key(NodeIndex(index)) == priv2);
+  REQUIRE(priv_joiner.private_key(NodeIndex(joiner_index)) == priv2);
   REQUIRE(priv_joiner.update_secret.size() == hash_size);
   REQUIRE(priv_joiner.update_secret == last_update_secret);