Skip to main content

openmls/treesync/node/
leaf_node.rs

1//! This module contains the [`LeafNode`] struct and its implementation.
2use std::collections::HashSet;
3
4use openmls_traits::{
5    crypto::OpenMlsCrypto, random::OpenMlsRand, signatures::Signer, types::Ciphersuite,
6};
7use serde::{Deserialize, Serialize};
8use thiserror::Error;
9use tls_codec::{
10    Serialize as TlsSerializeTrait, TlsDeserialize, TlsDeserializeBytes, TlsSerialize, TlsSize,
11    VLBytes,
12};
13
14use super::encryption_keys::{EncryptionKey, EncryptionKeyPair};
15use crate::{
16    binary_tree::array_representation::LeafNodeIndex,
17    ciphersuite::{
18        signable::{Signable, SignedStruct, Verifiable, VerifiedStruct},
19        Signature, SignaturePublicKey,
20    },
21    credentials::{Credential, CredentialType, CredentialWithKey},
22    error::LibraryError,
23    extensions::{ExtensionType, Extensions},
24    group::GroupId,
25    key_packages::{KeyPackage, Lifetime},
26    prelude::KeyPackageBundle,
27    storage::OpenMlsProvider,
28};
29
30use crate::treesync::errors::LeafNodeValidationError;
31
32mod capabilities;
33mod codec;
34
35pub use capabilities::*;
36
37pub(crate) struct NewLeafNodeParams {
38    pub(crate) ciphersuite: Ciphersuite,
39    pub(crate) credential_with_key: CredentialWithKey,
40    pub(crate) leaf_node_source: LeafNodeSource,
41    pub(crate) capabilities: Capabilities,
42    pub(crate) extensions: Extensions<LeafNode>,
43    pub(crate) tree_info_tbs: TreeInfoTbs,
44}
45
46/// Set of LeafNode parameters that are used when regenerating a LeafNodes
47/// during an update operation.
48#[derive(Debug, PartialEq, Clone)]
49pub(crate) struct UpdateLeafNodeParams {
50    pub(crate) credential_with_key: CredentialWithKey,
51    pub(crate) capabilities: Capabilities,
52    pub(crate) extensions: Extensions<LeafNode>,
53}
54
55impl UpdateLeafNodeParams {
56    #[cfg(test)]
57    pub(crate) fn derive(leaf_node: &LeafNode) -> Self {
58        Self {
59            credential_with_key: CredentialWithKey {
60                credential: leaf_node.payload.credential.clone(),
61                signature_key: leaf_node.payload.signature_key.clone(),
62            },
63            capabilities: leaf_node.payload.capabilities.clone(),
64            extensions: leaf_node.payload.extensions.clone(),
65        }
66    }
67}
68
69/// Parameters for a leaf node that can be chosen by the application.
70#[derive(Debug, PartialEq, Clone, Default)]
71pub struct LeafNodeParameters {
72    credential_with_key: Option<CredentialWithKey>,
73    capabilities: Option<Capabilities>,
74    extensions: Option<Extensions<LeafNode>>,
75}
76
77impl LeafNodeParameters {
78    /// Create a new [`LeafNodeParametersBuilder`].
79    pub fn builder() -> LeafNodeParametersBuilder {
80        LeafNodeParametersBuilder::default()
81    }
82
83    /// Returns the credential with key.
84    pub fn credential_with_key(&self) -> Option<&CredentialWithKey> {
85        self.credential_with_key.as_ref()
86    }
87
88    /// Returns the capabilities.
89    pub fn capabilities(&self) -> Option<&Capabilities> {
90        self.capabilities.as_ref()
91    }
92
93    /// Returns the extensions.
94    pub fn extensions(&self) -> Option<&Extensions<LeafNode>> {
95        self.extensions.as_ref()
96    }
97
98    pub(crate) fn is_empty(&self) -> bool {
99        self.credential_with_key.is_none()
100            && self.capabilities.is_none()
101            && self.extensions.is_none()
102    }
103
104    pub(crate) fn set_credential_with_key(&mut self, credential_with_key: CredentialWithKey) {
105        self.credential_with_key = Some(credential_with_key);
106    }
107
108    #[cfg(feature = "virtual-clients-draft")]
109    pub(crate) fn set_extensions(&mut self, extensions: Extensions<LeafNode>) {
110        self.extensions = Some(extensions);
111    }
112}
113
114/// Builder for [`LeafNodeParameters`].
115#[derive(Debug, Default)]
116pub struct LeafNodeParametersBuilder {
117    credential_with_key: Option<CredentialWithKey>,
118    capabilities: Option<Capabilities>,
119    extensions: Option<Extensions<LeafNode>>,
120}
121
122impl LeafNodeParametersBuilder {
123    /// Set the credential with key.
124    pub fn with_credential_with_key(mut self, credential_with_key: CredentialWithKey) -> Self {
125        self.credential_with_key = Some(credential_with_key);
126        self
127    }
128
129    /// Set the capabilities.
130    pub fn with_capabilities(mut self, capabilities: Capabilities) -> Self {
131        self.capabilities = Some(capabilities);
132        self
133    }
134
135    /// Set the extensions.
136    ///
137    /// Returns an error if one or more of the extensions is invalid in leaf nodes.
138    pub fn with_extensions(mut self, extensions: Extensions<LeafNode>) -> Self {
139        self.extensions = Some(extensions);
140        self
141    }
142
143    /// Build the [`LeafNodeParameters`].
144    pub fn build(self) -> LeafNodeParameters {
145        LeafNodeParameters {
146            credential_with_key: self.credential_with_key,
147            capabilities: self.capabilities,
148            extensions: self.extensions,
149        }
150    }
151}
152
153/// This struct implements the MLS leaf node.
154///
155/// ```c
156/// // draft-ietf-mls-protocol-17
157/// struct {
158///     HPKEPublicKey encryption_key;
159///     SignaturePublicKey signature_key;
160///     Credential credential;
161///     Capabilities capabilities;
162///
163///     LeafNodeSource leaf_node_source;
164///     select (LeafNode.leaf_node_source) {
165///         case key_package:
166///             Lifetime lifetime;
167///
168///         case update:
169///             struct{};
170///
171///         case commit:
172///             opaque parent_hash<V>;
173///     };
174///
175///     Extension extensions<V>;
176///     /* SignWithLabel(., "LeafNodeTBS", LeafNodeTBS) */
177///     opaque signature<V>;
178/// } LeafNode;
179/// ```
180#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, TlsSerialize, TlsSize)]
181pub struct LeafNode {
182    payload: LeafNodePayload,
183    signature: Signature,
184}
185
186impl LeafNode {
187    /// Create a new [`LeafNode`].
188    /// This first creates a `LeadNodeTbs` and returns the result of signing
189    /// it.
190    ///
191    /// This function generates a fresh HPKE key pair for the leaf node and
192    /// returns the HPKE key pair along with the new leaf node.
193    /// The caller is responsible for storing the private key.
194    pub(crate) fn new(
195        provider: &impl OpenMlsProvider,
196        signer: &impl Signer,
197        new_leaf_node_params: NewLeafNodeParams,
198    ) -> Result<(Self, EncryptionKeyPair), LibraryError> {
199        let NewLeafNodeParams {
200            ciphersuite,
201            credential_with_key,
202            leaf_node_source,
203            capabilities,
204            extensions,
205            tree_info_tbs,
206        } = new_leaf_node_params;
207
208        // Create a new encryption key pair.
209        let encryption_key_pair =
210            EncryptionKeyPair::random(provider.rand(), provider.crypto(), ciphersuite)?;
211
212        let leaf_node = Self::new_with_key(
213            encryption_key_pair.public_key().clone(),
214            credential_with_key,
215            leaf_node_source,
216            capabilities,
217            extensions,
218            tree_info_tbs,
219            signer,
220        )?;
221
222        Ok((leaf_node, encryption_key_pair))
223    }
224
225    /// Create a new [`LeafNode`] from a caller-provided encryption key pair.
226    ///
227    /// Mirrors [`LeafNode::new`] but uses `encryption_key_pair` instead of
228    /// generating a fresh one. This is the virtual-clients KeyPackage build
229    /// hook: the encryption key is derived from the per-operation secret so a
230    /// sibling can reproduce it.
231    #[cfg(feature = "virtual-clients-draft")]
232    pub(crate) fn new_with_encryption_key_pair(
233        signer: &impl Signer,
234        new_leaf_node_params: NewLeafNodeParams,
235        encryption_key_pair: EncryptionKeyPair,
236    ) -> Result<(Self, EncryptionKeyPair), LibraryError> {
237        let NewLeafNodeParams {
238            ciphersuite: _,
239            credential_with_key,
240            leaf_node_source,
241            capabilities,
242            extensions,
243            tree_info_tbs,
244        } = new_leaf_node_params;
245
246        let leaf_node = Self::new_with_key(
247            encryption_key_pair.public_key().clone(),
248            credential_with_key,
249            leaf_node_source,
250            capabilities,
251            extensions,
252            tree_info_tbs,
253            signer,
254        )?;
255
256        Ok((leaf_node, encryption_key_pair))
257    }
258
259    /// Creates a new placeholder [`LeafNode`] that is used to build external
260    /// commits.
261    ///
262    /// Note: This is not a valid leaf node and it must be rekeyed and signed
263    /// before it can be used.
264    pub(crate) fn new_placeholder() -> Self {
265        let payload = LeafNodePayload {
266            encryption_key: EncryptionKey::from(Vec::new()),
267            signature_key: Vec::new().into(),
268            credential: Credential::new(CredentialType::Basic, Vec::new()),
269            capabilities: Capabilities::default(),
270            leaf_node_source: LeafNodeSource::Update,
271            extensions: Extensions::default(),
272        };
273
274        Self {
275            payload,
276            signature: Vec::new().into(),
277        }
278    }
279
280    /// Create a new leaf node with a given HPKE encryption key pair.
281    /// The key pair must be stored in the key store by the caller.
282    fn new_with_key(
283        encryption_key: EncryptionKey,
284        credential_with_key: CredentialWithKey,
285        leaf_node_source: LeafNodeSource,
286        capabilities: Capabilities,
287        extensions: Extensions<LeafNode>,
288        tree_info_tbs: TreeInfoTbs,
289        signer: &impl Signer,
290    ) -> Result<Self, LibraryError> {
291        let leaf_node_tbs = LeafNodeTbs::new(
292            encryption_key,
293            credential_with_key,
294            capabilities,
295            leaf_node_source,
296            extensions,
297            tree_info_tbs,
298        );
299
300        leaf_node_tbs
301            .sign(signer)
302            .map_err(|_| LibraryError::custom("Signing failed"))
303    }
304
305    /// New [`LeafNode`] with a parent hash.
306    ///
307    /// With the `virtual-clients-draft` feature, an
308    /// `encryption_key_pair_override` may be supplied. If `Some`, it is used
309    /// as the leaf's encryption keypair instead of generating a fresh one.
310    /// This is the hook for the virtual-clients-draft sender.
311    #[allow(clippy::too_many_arguments)]
312    pub(in crate::treesync) fn new_with_parent_hash(
313        rand: &impl OpenMlsRand,
314        crypto: &impl OpenMlsCrypto,
315        ciphersuite: Ciphersuite,
316        parent_hash: &[u8],
317        leaf_node_params: UpdateLeafNodeParams,
318        group_id: GroupId,
319        leaf_index: LeafNodeIndex,
320        signer: &impl Signer,
321        #[cfg(feature = "virtual-clients-draft")] encryption_key_pair_override: Option<
322            EncryptionKeyPair,
323        >,
324    ) -> Result<(Self, EncryptionKeyPair), LibraryError> {
325        #[cfg(feature = "virtual-clients-draft")]
326        let encryption_key_pair = match encryption_key_pair_override {
327            Some(kp) => kp,
328            None => EncryptionKeyPair::random(rand, crypto, ciphersuite)?,
329        };
330        #[cfg(not(feature = "virtual-clients-draft"))]
331        let encryption_key_pair = EncryptionKeyPair::random(rand, crypto, ciphersuite)?;
332
333        let leaf_node_tbs = LeafNodeTbs::new(
334            encryption_key_pair.public_key().clone(),
335            leaf_node_params.credential_with_key,
336            leaf_node_params.capabilities,
337            LeafNodeSource::Commit(parent_hash.into()),
338            leaf_node_params.extensions,
339            TreeInfoTbs::Commit(TreePosition {
340                group_id,
341                leaf_index,
342            }),
343        );
344
345        // Sign the leaf node
346        let leaf_node = leaf_node_tbs
347            .sign(signer)
348            .map_err(|_| LibraryError::custom("Signing failed"))?;
349
350        Ok((leaf_node, encryption_key_pair))
351    }
352
353    /// Generate a fresh leaf node.
354    ///
355    /// This includes generating a new encryption key pair that is stored in the
356    /// key store.
357    ///
358    /// This function can be used when generating an update. In most other cases
359    /// a leaf node should be generated as part of a new [`KeyPackage`].
360    #[cfg(all(test, feature = "generate-kats"))]
361    pub(crate) fn generate_update<Provider: OpenMlsProvider>(
362        ciphersuite: Ciphersuite,
363        credential_with_key: CredentialWithKey,
364        capabilities: Capabilities,
365        extensions: Extensions<LeafNode>,
366        tree_info_tbs: TreeInfoTbs,
367        provider: &Provider,
368        signer: &impl Signer,
369    ) -> Result<Self, LeafNodeGenerationError<Provider::StorageError>> {
370        // Note that this function is supposed to be used in the public API only
371        // because it is interacting with the key store.
372
373        let new_leaf_node_params = NewLeafNodeParams {
374            ciphersuite,
375            credential_with_key,
376            leaf_node_source: LeafNodeSource::Update,
377            capabilities,
378            extensions,
379            tree_info_tbs,
380        };
381
382        let (leaf_node, encryption_key_pair) = Self::new(provider, signer, new_leaf_node_params)?;
383
384        // Store the encryption key pair in the key store.
385        encryption_key_pair
386            .write(provider.storage())
387            .map_err(LeafNodeGenerationError::StorageError)?;
388
389        Ok(leaf_node)
390    }
391
392    /// Update a leaf node.
393    ///
394    /// This function generates a new encryption key pair that is stored in the
395    /// key store and also returned.
396    ///
397    /// This function can be used when generating an update. In most other cases
398    /// a leaf node should be generated as part of a new [`KeyPackage`].
399    pub(crate) fn update<Provider: OpenMlsProvider>(
400        &mut self,
401        ciphersuite: Ciphersuite,
402        provider: &Provider,
403        signer: &impl Signer,
404        group_id: GroupId,
405        leaf_index: LeafNodeIndex,
406        leaf_node_parmeters: LeafNodeParameters,
407    ) -> Result<EncryptionKeyPair, LeafNodeUpdateError<Provider::StorageError>> {
408        let tree_info = TreeInfoTbs::Update(TreePosition::new(group_id, leaf_index));
409        let mut leaf_node_tbs = LeafNodeTbs::from(self.clone(), tree_info);
410
411        // Update credential
412        if let Some(credential_with_key) = leaf_node_parmeters.credential_with_key {
413            leaf_node_tbs.payload.credential = credential_with_key.credential;
414            leaf_node_tbs.payload.signature_key = credential_with_key.signature_key;
415        }
416
417        // Update extensions
418        if let Some(extensions) = leaf_node_parmeters.extensions {
419            leaf_node_tbs.payload.extensions = extensions;
420        }
421
422        // Update capabilities
423        if let Some(capabilities) = leaf_node_parmeters.capabilities {
424            leaf_node_tbs.payload.capabilities = capabilities;
425        }
426
427        // Create a new encryption key pair
428        let encryption_key_pair =
429            EncryptionKeyPair::random(provider.rand(), provider.crypto(), ciphersuite)?;
430        leaf_node_tbs.payload.encryption_key = encryption_key_pair.public_key().clone();
431
432        // Store the encryption key pair in the key store.
433        encryption_key_pair
434            .write(provider.storage())
435            .map_err(LeafNodeUpdateError::Storage)?;
436
437        // Set the leaf node source to update
438        leaf_node_tbs.payload.leaf_node_source = LeafNodeSource::Update;
439
440        // Sign the leaf node
441        let leaf_node = leaf_node_tbs.sign(signer)?;
442        self.payload = leaf_node.payload;
443        self.signature = leaf_node.signature;
444
445        Ok(encryption_key_pair)
446    }
447
448    /// Returns the `encryption_key`.
449    pub fn encryption_key(&self) -> &EncryptionKey {
450        &self.payload.encryption_key
451    }
452
453    /// Returns the `signature_key` as byte slice.
454    pub fn signature_key(&self) -> &SignaturePublicKey {
455        &self.payload.signature_key
456    }
457
458    /// Returns the `credential`.
459    pub fn credential(&self) -> &Credential {
460        &self.payload.credential
461    }
462
463    /// Returns the `parent_hash` as byte slice or `None`.
464    pub fn parent_hash(&self) -> Option<&[u8]> {
465        match &self.payload.leaf_node_source {
466            LeafNodeSource::Commit(ph) => Some(ph.as_slice()),
467            _ => None,
468        }
469    }
470
471    /// Returns the [`Lifetime`] if present.
472    /// `None` otherwise.
473    pub(crate) fn life_time(&self) -> Option<&Lifetime> {
474        if let LeafNodeSource::KeyPackage(life_time) = &self.payload.leaf_node_source {
475            Some(life_time)
476        } else {
477            None
478        }
479    }
480
481    /// Returns a reference to the [`Signature`] of this leaf.
482    pub fn signature(&self) -> &Signature {
483        &self.signature
484    }
485
486    /// Return a reference to [`Capabilities`].
487    pub fn capabilities(&self) -> &Capabilities {
488        &self.payload.capabilities
489    }
490
491    /// Return a reference to the leaf node source.
492    pub fn leaf_node_source(&self) -> &LeafNodeSource {
493        &self.payload.leaf_node_source
494    }
495
496    /// Return a reference to the leaf node extensions.
497    pub fn extensions(&self) -> &Extensions<LeafNode> {
498        &self.payload.extensions
499    }
500
501    /// The virtual-client derivation info this leaf might contain.
502    #[cfg(feature = "virtual-clients-draft")]
503    pub(crate) fn vc_derivation_info(
504        &self,
505    ) -> Result<
506        Option<crate::components::vc_derivation_info::DerivationInfo>,
507        crate::components::vc_derivation_info::VirtualClientsError,
508    > {
509        use tls_codec::DeserializeBytes as _;
510
511        use crate::components::vc_derivation_info::{
512            DerivationInfo, VirtualClientsError, VC_COMPONENT_ID,
513        };
514
515        let Some(bytes) = self
516            .extensions()
517            .app_data_dictionary()
518            .and_then(|dict| dict.dictionary().get(&VC_COMPONENT_ID))
519        else {
520            return Ok(None);
521        };
522        DerivationInfo::tls_deserialize_exact_bytes(bytes)
523            .map(Some)
524            .map_err(|e| {
525                log::error!("vc: leaf derivation info deserialize failed: {e:?}");
526                VirtualClientsError::DerivationInfoMalformed
527            })
528    }
529
530    /// Returns `true` if the [`ExtensionType`] is supported by this leaf node.
531    pub(crate) fn supports_extension(&self, extension_type: &ExtensionType) -> bool {
532        extension_type.is_default()
533            || self
534                .payload
535                .capabilities
536                .extensions
537                .contains(extension_type)
538    }
539
540    /// Check whether the this leaf node supports all the required extensions
541    /// in the provided list.
542    pub(crate) fn check_extension_support(
543        &self,
544        extensions: &[ExtensionType],
545    ) -> Result<(), LeafNodeValidationError> {
546        let mut required = extensions.iter().filter(|e| !e.is_default()).peekable();
547
548        // Skip building the lookup if there are no non-default extensions.
549        if required.peek().is_none() {
550            return Ok(());
551        }
552
553        let supported: HashSet<ExtensionType> = self
554            .payload
555            .capabilities
556            .extensions
557            .iter()
558            .copied()
559            .collect();
560
561        if let Some(unsupported) = required.find(|e| !supported.contains(e)) {
562            log::error!(
563                "Leaf node does not support required extension {:?}\n
564                    Supported extensions: {:?}",
565                unsupported,
566                self.payload.capabilities.extensions
567            );
568            return Err(LeafNodeValidationError::UnsupportedExtensions);
569        }
570
571        Ok(())
572    }
573
574    /// Perform all checks that can be done without further context:
575    /// - the used extensions are not known to be invalid in leaf nodes
576    /// - the types of the used extensions are covered by the capabilities
577    /// - the type of the credential is covered by the capabilities
578    pub(crate) fn validate_locally(&self) -> Result<(), LeafNodeValidationError> {
579        // Check that no extension is invalid when used in leaf nodes.
580        // https://validation.openmls.tech/#valn1601
581        // NOTE: This check is conducted manually for now, instead of using the method
582        // Extensions::validate_extension_types_for_leaf_node(),
583        // in order to collect the invalid extension types for the log message below.
584        // However, it could be better to instead return the list of invalid extension types
585        // as part of Extensions::validate_extension_types_for_leaf_node(),
586        // as part of the error message.
587        let invalid_extension_types = self
588            .extensions()
589            .iter()
590            .filter(|ext| !ext.extension_type().is_valid_in_leaf_node())
591            .collect::<Vec<_>>();
592        if !invalid_extension_types.is_empty() {
593            log::error!("Invalid extension used in leaf node: {invalid_extension_types:?}");
594            return Err(LeafNodeValidationError::UnsupportedExtensions);
595        }
596
597        // Check that all extensions are contained in the capabilities.
598        if !self.capabilities().contains_extensions(self.extensions()) {
599            log::error!(
600                "Leaf node does not support all extensions it uses\n
601                Supported extensions: {:?}\n
602                Used extensions: {:?}",
603                self.payload.capabilities.extensions,
604                self.extensions()
605            );
606            return Err(LeafNodeValidationError::UnsupportedExtensions);
607        }
608
609        // Check that the capabilities contain the leaf node's credential type.
610        // (https://validation.openmls.tech/#valn0113)
611        if !self
612            .capabilities()
613            .contains_credential(self.credential().credential_type())
614        {
615            return Err(LeafNodeValidationError::UnsupportedCredentials);
616        }
617
618        Ok(())
619    }
620}
621
622/// The payload of a [`LeafNode`]
623///
624/// ```text
625/// struct {
626///     HPKEPublicKey encryption_key;
627///     SignaturePublicKey signature_key;
628///     Credential credential;
629///     Capabilities capabilities;
630///
631///     LeafNodeSource leaf_node_source;
632///     select (LeafNode.leaf_node_source) {
633///         case key_package:
634///             Lifetime lifetime;
635///
636///         case update:
637///             struct{};
638///
639///         case commit:
640///             opaque parent_hash<V>;
641///     };
642///
643///     Extension extensions<V>;
644///     ...
645/// } LeafNode;
646/// ```
647#[derive(
648    Debug,
649    Clone,
650    PartialEq,
651    Eq,
652    Serialize,
653    Deserialize,
654    TlsSerialize,
655    TlsDeserialize,
656    TlsDeserializeBytes,
657    TlsSize,
658)]
659struct LeafNodePayload {
660    encryption_key: EncryptionKey,
661    signature_key: SignaturePublicKey,
662    credential: Credential,
663    capabilities: Capabilities,
664    leaf_node_source: LeafNodeSource,
665    extensions: Extensions<LeafNode>,
666}
667
668/// The source of the `LeafNode`.
669#[derive(
670    Debug,
671    Clone,
672    PartialEq,
673    Eq,
674    Serialize,
675    Deserialize,
676    TlsSerialize,
677    TlsDeserialize,
678    TlsDeserializeBytes,
679    TlsSize,
680)]
681#[repr(u8)]
682pub enum LeafNodeSource {
683    /// The leaf node was added to the group as part of a key package.
684    #[tls_codec(discriminant = 1)]
685    KeyPackage(Lifetime),
686    /// The leaf node was added through an Update proposal.
687    Update,
688    /// The leaf node was added via a Commit.
689    Commit(ParentHash),
690}
691
692pub type ParentHash = VLBytes;
693
694/// To-be-signed leaf node.
695///
696/// ```c
697/// // draft-ietf-mls-protocol-17
698/// struct {
699///     HPKEPublicKey encryption_key;
700///     SignaturePublicKey signature_key;
701///     Credential credential;
702///     Capabilities capabilities;
703///
704///     LeafNodeSource leaf_node_source;
705///     select (LeafNodeTBS.leaf_node_source) {
706///         case key_package:
707///             Lifetime lifetime;
708///
709///         case update:
710///             struct{};
711///
712///         case commit:
713///             opaque parent_hash<V>;
714///     };
715///
716///     Extension extensions<V>;
717///
718///     // ... continued in [`TreeInfo`] ...
719/// } LeafNodeTBS;
720/// ```
721#[derive(Debug, TlsSerialize, TlsSize)]
722pub struct LeafNodeTbs {
723    payload: LeafNodePayload,
724    tree_info_tbs: TreeInfoTbs,
725}
726
727impl LeafNodeTbs {
728    /// Build a [`LeafNodeTbs`] from a [`LeafNode`] and a [`TreeInfo`]
729    /// to update a leaf node.
730    pub(crate) fn from(leaf_node: LeafNode, tree_info_tbs: TreeInfoTbs) -> Self {
731        Self {
732            payload: leaf_node.payload,
733            tree_info_tbs,
734        }
735    }
736
737    /// Build a new [`LeafNodeTbs`] from a [`KeyPackage`] and [`Credential`].
738    /// To get the [`LeafNode`] call [`LeafNode::sign`].
739    pub(crate) fn new(
740        encryption_key: EncryptionKey,
741        credential_with_key: CredentialWithKey,
742        capabilities: Capabilities,
743        leaf_node_source: LeafNodeSource,
744        extensions: Extensions<LeafNode>,
745        tree_info_tbs: TreeInfoTbs,
746    ) -> Self {
747        let payload = LeafNodePayload {
748            encryption_key,
749            signature_key: credential_with_key.signature_key,
750            credential: credential_with_key.credential,
751            capabilities,
752            leaf_node_source,
753            extensions,
754        };
755
756        LeafNodeTbs {
757            payload,
758            tree_info_tbs,
759        }
760    }
761}
762
763/// Helper struct that holds additional information required to sign a leaf node.
764///
765/// ```c
766/// // draft-ietf-mls-protocol-17
767/// struct {
768///     // ... continued from [`LeafNodeTbs`] ...
769///
770///     select (LeafNodeTBS.leaf_node_source) {
771///         case key_package:
772///             struct{};
773///
774///         case update:
775///             opaque group_id<V>;
776///             uint32 leaf_index;
777///
778///         case commit:
779///             opaque group_id<V>;
780///             uint32 leaf_index;
781///     };
782/// } LeafNodeTBS;
783/// ```
784#[derive(Debug)]
785pub(crate) enum TreeInfoTbs {
786    KeyPackage,
787    Update(TreePosition),
788    Commit(TreePosition),
789}
790
791#[derive(Debug, Clone, PartialEq, Eq, TlsSerialize, TlsSize)]
792pub(crate) struct TreePosition {
793    group_id: GroupId,
794    leaf_index: LeafNodeIndex,
795}
796
797impl TreePosition {
798    pub(crate) fn new(group_id: GroupId, leaf_index: LeafNodeIndex) -> Self {
799        Self {
800            group_id,
801            leaf_index,
802        }
803    }
804
805    #[cfg(feature = "test-utils")]
806    pub(crate) fn into_parts(self) -> (GroupId, LeafNodeIndex) {
807        (self.group_id, self.leaf_index)
808    }
809}
810
811const LEAF_NODE_SIGNATURE_LABEL: &str = "LeafNodeTBS";
812
813#[derive(
814    Debug,
815    Clone,
816    PartialEq,
817    Eq,
818    Serialize,
819    Deserialize,
820    TlsSerialize,
821    TlsDeserialize,
822    TlsDeserializeBytes,
823    TlsSize,
824)]
825pub struct LeafNodeIn {
826    payload: LeafNodePayload,
827    signature: Signature,
828}
829
830impl LeafNodeIn {
831    pub(crate) fn into_verifiable_leaf_node(self) -> VerifiableLeafNode {
832        match self.payload.leaf_node_source {
833            LeafNodeSource::KeyPackage(_) => {
834                let verifiable = VerifiableKeyPackageLeafNode {
835                    payload: self.payload,
836                    signature: self.signature,
837                };
838                VerifiableLeafNode::KeyPackage(verifiable)
839            }
840            LeafNodeSource::Update => {
841                let verifiable = VerifiableUpdateLeafNode {
842                    payload: self.payload,
843                    signature: self.signature,
844                    tree_position: None,
845                };
846                VerifiableLeafNode::Update(verifiable)
847            }
848            LeafNodeSource::Commit(_) => {
849                let verifiable = VerifiableCommitLeafNode {
850                    payload: self.payload,
851                    signature: self.signature,
852                    tree_position: None,
853                };
854                VerifiableLeafNode::Commit(verifiable)
855            }
856        }
857    }
858
859    /// Returns the `encryption_key` as byte slice.
860    pub fn encryption_key(&self) -> &EncryptionKey {
861        &self.payload.encryption_key
862    }
863
864    /// Returns the `signature_key` as byte slice.
865    pub fn signature_key(&self) -> &SignaturePublicKey {
866        &self.payload.signature_key
867    }
868
869    /// Returns the `signature_key` as byte slice.
870    pub fn credential(&self) -> &Credential {
871        &self.payload.credential
872    }
873
874    /// Assume that signature is valid and return the corresponding [`LeafNode`].
875    ///
876    /// # Safety
877    ///
878    /// The caller must guarantee that the leaf node is verified.
879    #[cfg(feature = "unchecked-conversions")]
880    pub fn into_unchecked(self) -> LeafNode {
881        LeafNode {
882            payload: self.payload,
883            signature: self.signature,
884        }
885    }
886}
887
888impl From<LeafNode> for LeafNodeIn {
889    fn from(leaf_node: LeafNode) -> Self {
890        Self {
891            payload: leaf_node.payload,
892            signature: leaf_node.signature,
893        }
894    }
895}
896
897#[cfg(any(feature = "test-utils", test))]
898impl From<LeafNodeIn> for LeafNode {
899    fn from(deserialized: LeafNodeIn) -> Self {
900        Self {
901            payload: deserialized.payload,
902            signature: deserialized.signature,
903        }
904    }
905}
906
907impl From<KeyPackage> for LeafNode {
908    fn from(key_package: KeyPackage) -> Self {
909        key_package.leaf_node().clone()
910    }
911}
912
913impl From<KeyPackageBundle> for LeafNode {
914    fn from(key_package: KeyPackageBundle) -> Self {
915        key_package.key_package().leaf_node().clone()
916    }
917}
918
919#[derive(Debug, Clone, PartialEq, Eq)]
920pub(crate) enum VerifiableLeafNode {
921    KeyPackage(VerifiableKeyPackageLeafNode),
922    Update(VerifiableUpdateLeafNode),
923    Commit(VerifiableCommitLeafNode),
924}
925
926impl VerifiableLeafNode {
927    pub(crate) fn signature_key(&self) -> &SignaturePublicKey {
928        match self {
929            VerifiableLeafNode::KeyPackage(v) => v.signature_key(),
930            VerifiableLeafNode::Update(v) => v.signature_key(),
931            VerifiableLeafNode::Commit(v) => v.signature_key(),
932        }
933    }
934}
935
936#[derive(Debug, Clone, PartialEq, Eq)]
937pub(crate) struct VerifiableKeyPackageLeafNode {
938    payload: LeafNodePayload,
939    signature: Signature,
940}
941
942impl VerifiableKeyPackageLeafNode {
943    pub(crate) fn signature_key(&self) -> &SignaturePublicKey {
944        &self.payload.signature_key
945    }
946}
947
948// https://validation.openmls.tech/#valn0102
949impl Verifiable for VerifiableKeyPackageLeafNode {
950    type VerifiedStruct = LeafNode;
951
952    fn unsigned_payload(&self) -> Result<Vec<u8>, tls_codec::Error> {
953        self.payload.tls_serialize_detached()
954    }
955
956    fn signature(&self) -> &Signature {
957        &self.signature
958    }
959
960    fn label(&self) -> &str {
961        LEAF_NODE_SIGNATURE_LABEL
962    }
963
964    fn verify(
965        self,
966        crypto: &impl openmls_traits::crypto::OpenMlsCrypto,
967        pk: &crate::ciphersuite::OpenMlsSignaturePublicKey,
968    ) -> Result<Self::VerifiedStruct, crate::ciphersuite::signable::SignatureError> {
969        self.verify_no_out(crypto, pk)?;
970        Ok(LeafNode {
971            payload: self.payload,
972            signature: self.signature,
973        })
974    }
975}
976
977impl VerifiedStruct for LeafNode {}
978
979#[derive(Debug, Clone, PartialEq, Eq)]
980pub(crate) struct VerifiableUpdateLeafNode {
981    payload: LeafNodePayload,
982    signature: Signature,
983    tree_position: Option<TreePosition>,
984}
985
986impl VerifiableUpdateLeafNode {
987    pub(crate) fn add_tree_position(&mut self, tree_info: TreePosition) {
988        self.tree_position = Some(tree_info);
989    }
990
991    pub(crate) fn signature_key(&self) -> &SignaturePublicKey {
992        &self.payload.signature_key
993    }
994}
995
996impl Verifiable for VerifiableUpdateLeafNode {
997    type VerifiedStruct = LeafNode;
998
999    fn unsigned_payload(&self) -> Result<Vec<u8>, tls_codec::Error> {
1000        let tree_info_tbs = match &self.tree_position {
1001            Some(tree_position) => TreeInfoTbs::Commit(tree_position.clone()),
1002            None => return Err(tls_codec::Error::InvalidInput),
1003        };
1004        let leaf_node_tbs = LeafNodeTbs {
1005            payload: self.payload.clone(),
1006            tree_info_tbs,
1007        };
1008        leaf_node_tbs.tls_serialize_detached()
1009    }
1010
1011    fn signature(&self) -> &Signature {
1012        &self.signature
1013    }
1014
1015    fn label(&self) -> &str {
1016        LEAF_NODE_SIGNATURE_LABEL
1017    }
1018
1019    fn verify(
1020        self,
1021        crypto: &impl openmls_traits::crypto::OpenMlsCrypto,
1022        pk: &crate::ciphersuite::OpenMlsSignaturePublicKey,
1023    ) -> Result<Self::VerifiedStruct, crate::ciphersuite::signable::SignatureError> {
1024        self.verify_no_out(crypto, pk)?;
1025        Ok(LeafNode {
1026            payload: self.payload,
1027            signature: self.signature,
1028        })
1029    }
1030}
1031
1032#[derive(Debug, Clone, PartialEq, Eq)]
1033pub(crate) struct VerifiableCommitLeafNode {
1034    payload: LeafNodePayload,
1035    signature: Signature,
1036    tree_position: Option<TreePosition>,
1037}
1038
1039impl VerifiableCommitLeafNode {
1040    pub(crate) fn add_tree_position(&mut self, tree_info: TreePosition) {
1041        self.tree_position = Some(tree_info);
1042    }
1043
1044    pub(crate) fn signature_key(&self) -> &SignaturePublicKey {
1045        &self.payload.signature_key
1046    }
1047}
1048
1049impl Verifiable for VerifiableCommitLeafNode {
1050    type VerifiedStruct = LeafNode;
1051
1052    fn unsigned_payload(&self) -> Result<Vec<u8>, tls_codec::Error> {
1053        let tree_info_tbs = match &self.tree_position {
1054            Some(tree_position) => TreeInfoTbs::Commit(tree_position.clone()),
1055            None => return Err(tls_codec::Error::InvalidInput),
1056        };
1057        let leaf_node_tbs = LeafNodeTbs {
1058            payload: self.payload.clone(),
1059            tree_info_tbs,
1060        };
1061
1062        leaf_node_tbs.tls_serialize_detached()
1063    }
1064
1065    fn signature(&self) -> &Signature {
1066        &self.signature
1067    }
1068
1069    fn label(&self) -> &str {
1070        LEAF_NODE_SIGNATURE_LABEL
1071    }
1072
1073    fn verify(
1074        self,
1075        crypto: &impl openmls_traits::crypto::OpenMlsCrypto,
1076        pk: &crate::ciphersuite::OpenMlsSignaturePublicKey,
1077    ) -> Result<Self::VerifiedStruct, crate::ciphersuite::signable::SignatureError> {
1078        self.verify_no_out(crypto, pk)?;
1079        Ok(LeafNode {
1080            payload: self.payload,
1081            signature: self.signature,
1082        })
1083    }
1084}
1085
1086impl Signable for LeafNodeTbs {
1087    type SignedOutput = LeafNode;
1088
1089    fn unsigned_payload(&self) -> Result<Vec<u8>, tls_codec::Error> {
1090        self.tls_serialize_detached()
1091    }
1092
1093    fn label(&self) -> &str {
1094        LEAF_NODE_SIGNATURE_LABEL
1095    }
1096}
1097
1098impl SignedStruct<LeafNodeTbs> for LeafNode {
1099    fn from_payload(tbs: LeafNodeTbs, signature: Signature, _serialized_payload: Vec<u8>) -> Self {
1100        Self {
1101            payload: tbs.payload,
1102            signature,
1103        }
1104    }
1105}
1106
1107#[cfg(all(test, feature = "generate-kats"))]
1108#[derive(Error, Debug, PartialEq, Clone)]
1109pub enum LeafNodeGenerationError<StorageError> {
1110    /// See [`LibraryError`] for more details.
1111    #[error(transparent)]
1112    LibraryError(#[from] LibraryError),
1113
1114    /// Error storing leaf private key in storage.
1115    #[error("Error storing leaf private key.")]
1116    StorageError(StorageError),
1117}
1118
1119/// Leaf Node Update Error
1120#[derive(Error, Debug, PartialEq, Clone)]
1121pub enum LeafNodeUpdateError<StorageError> {
1122    /// See [`LibraryError`] for more details.
1123    #[error(transparent)]
1124    LibraryError(#[from] LibraryError),
1125
1126    /// Error storing leaf private key in storage.
1127    #[error("Error storing leaf private key.")]
1128    Storage(StorageError),
1129
1130    /// Signature error.
1131    #[error(transparent)]
1132    Signature(#[from] crate::ciphersuite::signable::SignatureError),
1133}