Skip to main content

openmls/treesync/node/
leaf_node.rs

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