Implement LightCommit more or less along the lines in the draft

include/mls/messages.h

@@ -279,6 +279,15 @@ struct Welcome
     const std::vector<PSKWithSecret>& psks);
 };
 
+struct LightCommit
+{
+  GroupContext group_context;
+  bytes confirmation_tag;
+  TreeSlice sender_membership_proof;
+  std::optional<HPKECiphertext> encrypted_path_secret;
+  std::optional<NodeIndex> decryption_node_index;
+};
+
 ///
 /// Proposals & Commit
 ///
@@ -645,6 +654,10 @@ struct PublicMessage
   bytes membership_mac(CipherSuite suite,
                        const bytes& membership_key,
                        const std::optional<GroupContext>& context) const;
+
+  // XXX(RLB) This is a hack to avoid unwrapping across epochs.  We should do
+  // something more elegant, like unchecked_content()
+  friend class State;
 };
 
 struct PrivateMessage

include/mls/state.h

@@ -138,6 +138,12 @@ class State
   std::optional<State> handle(const ValidatedContent& content_auth,
                               std::optional<State> cached_state);
 
+  ///
+  /// Light MLS
+  ///
+  LightCommit lighten_for(LeafIndex leaf, const MLSMessage& commit) const;
+  State handle(const LightCommit& light_commit) const;
+
   ///
   /// PSK management
   ///

include/mls/treekem.h

@@ -120,6 +120,9 @@ struct TreeKEMPrivateKey
   void implant(const TreeKEMPublicKey& pub,
                NodeIndex start,
                const bytes& path_secret);
+  void implant_matching(const TreeKEMPublicKey& pub,
+                        NodeIndex start,
+                        const bytes& path_secret);
 };
 
 struct TreeKEMPublicKey
@@ -167,6 +170,9 @@ struct TreeKEMPublicKey
 
   TreeSlice extract_slice(LeafIndex leaf) const;
   void implant_slice(const TreeSlice& slice);
+  std::tuple<HPKECiphertext, NodeIndex> slice_path(UpdatePath path,
+                                                   LeafIndex from,
+                                                   LeafIndex to) const;
 
   template<typename UnaryPredicate>
   bool all_leaves(const UnaryPredicate& pred) const

src/messages.cpp

@@ -12,7 +12,8 @@ namespace MLS_NAMESPACE {
 
 const Extension::Type ExternalPubExtension::type = ExtensionType::external_pub;
 const Extension::Type RatchetTreeExtension::type = ExtensionType::ratchet_tree;
-const Extension::Type MembershipProofExtension::type = ExtensionType::membership_proof;
+const Extension::Type MembershipProofExtension::type =
+  ExtensionType::membership_proof;
 const Extension::Type ExternalSendersExtension::type =
   ExtensionType::external_senders;
 const Extension::Type FlagsExtension::type = ExtensionType::flags;

src/session.cpp

@@ -305,8 +305,10 @@ Session::commit()
 {
   auto commit_secret = inner->fresh_secret();
   auto encrypt = inner->encrypt_handshake;
-  auto [commit, welcome, new_state] = inner->history.front().commit(
-    commit_secret, CommitOpts{ {}, true, encrypt, false, {} }, { encrypt, {}, 0 });
+  auto [commit, welcome, new_state] =
+    inner->history.front().commit(commit_secret,
+                                  CommitOpts{ {}, true, encrypt, false, {} },
+                                  { encrypt, {}, 0 });
 
   auto commit_msg = tls::marshal(commit);
   auto welcome_msg = tls::marshal(welcome);

src/state.cpp

@@ -29,13 +29,13 @@ State::State(bytes group_id,
   }
 
   _index = _tree.add_leaf(leaf_node);
-  _tree.set_hash_all();
   _tree_priv = TreeKEMPrivateKey::solo(suite, _index, std::move(enc_priv));
   if (!_tree_priv.consistent(_tree)) {
     throw InvalidParameterError("LeafNode inconsistent with private key");
   }
 
   // XXX(RLB): Convert KeyScheduleEpoch to take GroupContext?
+  _tree.set_hash_all();
   auto ctx = tls::marshal(group_context());
   _key_schedule =
     KeyScheduleEpoch(_suite, random_bytes(_suite.secret_size()), ctx);
@@ -53,7 +53,8 @@ State::import_tree(const bytes& tree_hash,
 {
   auto tree = TreeKEMPublicKey(_suite);
   auto maybe_tree_extn = extensions.find<RatchetTreeExtension>();
-  auto maybe_membership_proof_extn = extensions.find<MembershipProofExtension>();
+  auto maybe_membership_proof_extn =
+    extensions.find<MembershipProofExtension>();
 
   if (external) {
     tree = opt::get(external);
@@ -1010,6 +1011,97 @@ State::handle(const ValidatedContent& val_content,
   return next;
 }
 
+LightCommit
+State::lighten_for(LeafIndex leaf, const MLSMessage& commit_msg) const
+{
+  // Check that the current epoch is one higher than commit.epoch
+  if (_epoch != commit_msg.epoch() + 1) {
+    throw InvalidParameterError("Invalid epoch for lightening operation");
+  }
+
+  // Pull the GroupContext for the current state
+  const auto ctx = group_context();
+
+  // Make a memberhsip proof
+  const auto& public_msg = var::get<PublicMessage>(commit_msg.message);
+  const auto& sender =
+    var::get<MemberSender>(public_msg.content.sender.sender).sender;
+
+  const auto confirmation_tag = opt::get(public_msg.auth.confirmation_tag);
+  const auto sender_membership_proof = _tree.extract_slice(sender);
+
+  // Extract the correct path secret for the recipient
+  const auto& commit = var::get<Commit>(public_msg.content.content);
+  auto encrypted_path_secret = std::optional<HPKECiphertext>{};
+  auto decryption_node_index = std::optional<NodeIndex>{};
+  if (commit.path) {
+    const auto [secret, index] =
+      _tree.slice_path(opt::get(commit.path), sender, leaf);
+    encrypted_path_secret = secret;
+    decryption_node_index = index;
+  }
+
+  return {
+    ctx,
+    confirmation_tag,
+    sender_membership_proof,
+    encrypted_path_secret,
+    decryption_node_index,
+  };
+}
+
+State
+State::handle(const LightCommit& light_commit) const
+{
+  // Verify the membership proof
+  // TODO(RLB) Also verify the signature (?)
+  // XXX(RLB) Should this use the new or old tree hash?
+  if (light_commit.sender_membership_proof.tree_hash(_suite) !=
+      light_commit.group_context.tree_hash) {
+    throw ProtocolError("Invalid sender membership proof");
+  }
+
+  // Import the GroupContext
+  // TODO(RLB) Verify that version, cipher_suite, group_id, epoch are as
+  // expected
+  auto next = successor();
+  next._epoch += 1;
+  next._tree =
+    TreeKEMPublicKey(next._suite, light_commit.sender_membership_proof);
+  next._extensions = light_commit.group_context.extensions;
+
+  // Decrypt the commit secret
+  auto commit_secret = _suite.zero();
+  if (light_commit.encrypted_path_secret) {
+    const auto& encrypted_path_secret =
+      opt::get(light_commit.encrypted_path_secret);
+    const auto& decryption_node_index =
+      opt::get(light_commit.decryption_node_index);
+
+    const auto priv = opt::get(_tree_priv.private_key(decryption_node_index));
+    const auto context = tls::marshal(next.group_context());
+    const auto path_secret = priv.decrypt(next._suite,
+                                          encrypt_label::update_path_node,
+                                          context,
+                                          encrypted_path_secret);
+    const auto ancestor =
+      next._index.ancestor(light_commit.sender_membership_proof.leaf_index);
+    next._tree_priv.implant_matching(next._tree, ancestor, path_secret);
+
+    commit_secret = next._tree_priv.update_secret;
+  }
+
+  // Update the key schedule
+  // TODO(RLB) Need to accommodate PSKs for light clients
+  next._transcript_hash =
+    TranscriptHash(next._suite,
+                   light_commit.group_context.confirmed_transcript_hash,
+                   light_commit.confirmation_tag);
+  next.update_epoch_secrets(commit_secret, {}, std::nullopt);
+
+  return next;
+}
+
 ///
 /// Subgroup branching
 ///
@@ -2057,6 +2149,7 @@ State::update_epoch_secrets(const bytes& commit_secret,
     _transcript_hash.confirmed,
     _extensions,
   });
+
   _key_schedule =
     _key_schedule.next(commit_secret, psks, force_init_secret, ctx);
   _keys = _key_schedule.encryption_keys(_tree.size);

src/treekem.cpp

@@ -125,6 +125,30 @@ TreeKEMPrivateKey::implant(const TreeKEMPublicKey& pub,
   update_secret = pub.suite.derive_secret(secret, "path");
 }
 
+void
+TreeKEMPrivateKey::implant_matching(const TreeKEMPublicKey& pub,
+                                    NodeIndex start,
+                                    const bytes& path_secret)
+{
+  auto secret = path_secret;
+
+  path_secrets.insert_or_assign(start, secret);
+  private_key_cache.erase(start);
+
+  const auto dp = start.dirpath(pub.size);
+  for (const auto& n : dp) {
+    if (pub.node_at(n).blank()) {
+      continue;
+    }
+
+    secret = pub.suite.derive_secret(secret, "path");
+    path_secrets.insert_or_assign(n, secret);
+    private_key_cache.erase(n);
+  }
+
+  update_secret = pub.suite.derive_secret(secret, "path");
+}
+
 std::optional<HPKEPrivateKey>
 TreeKEMPrivateKey::private_key(NodeIndex n) const
 {
@@ -589,7 +613,7 @@ TreeKEMPublicKey::parent_hash_valid() const
 bool
 TreeKEMPublicKey::is_complete() const
 {
-  return nodes.size() == NodeCount{size}.val;
+  return nodes.size() == NodeCount{ size }.val;
 }
 
 std::vector<NodeIndex>
@@ -655,6 +679,34 @@ TreeKEMPublicKey::implant_slice(const TreeSlice& slice)
   implant_slice_unchecked(slice);
 }
 
+std::tuple<HPKECiphertext, NodeIndex>
+TreeKEMPublicKey::slice_path(UpdatePath path,
+                             LeafIndex from,
+                             LeafIndex to) const
+{
+  const auto toi = NodeIndex(to);
+  const auto fdp = filtered_direct_path(NodeIndex(from));
+
+  for (auto i = size_t(0); i < fdp.size(); i++) {
+    const auto& [dpi, res] = fdp.at(i);
+
+    if (!toi.is_below(dpi)) {
+      continue;
+    }
+
+    for (auto j = size_t(0); j < res.size(); j++) {
+      const auto resi = res.at(j);
+      if (!toi.is_below(resi)) {
+        continue;
+      }
+
+      return { path.nodes.at(i).encrypted_path_secret.at(j), resi };
+    }
+  }
+
+  throw ProtocolError("Decryption node not found");
+}
+
 void
 TreeKEMPublicKey::implant_slice_unchecked(const TreeSlice& slice)
 {