xmtp_id/key_package/
mls_ext_wrapper_encryption.rs

1use openmls::prelude::UnknownExtension;
2use openmls::prelude::{Ciphersuite, Extension, KeyPackage};
3use prost::{EncodeError, Message};
4use xmtp_configuration::WELCOME_WRAPPER_ENCRYPTION_EXTENSION_ID;
5use xmtp_cryptography::configuration::{CIPHERSUITE, POST_QUANTUM_CIPHERSUITE};
6use xmtp_proto::ConversionError;
7use xmtp_proto::xmtp::mls::message_contents::{
8    WelcomePointerWrapperAlgorithm as WelcomePointerWrapperAlgorithmProto,
9    WelcomeWrapperAlgorithm as WrapperAlgorithmProto,
10    WelcomeWrapperEncryption as WelcomeWrapperEncryptionProto,
11};
12
13#[derive(Debug, PartialEq, Clone, Copy)]
14pub enum WrapperAlgorithm {
15    Curve25519,
16    XWingMLKEM768Draft6,
17}
18
19impl WrapperAlgorithm {
20    pub fn to_mls_ciphersuite(self) -> Ciphersuite {
21        match self {
22            WrapperAlgorithm::Curve25519 => CIPHERSUITE,
23            WrapperAlgorithm::XWingMLKEM768Draft6 => POST_QUANTUM_CIPHERSUITE,
24        }
25    }
26    // hardcoded because the functions to do the translations are private
27    // and placed here so that any changes to the this algorithm will have to be handled
28    pub fn to_hpke_config(self) -> hpke_rs::Hpke<hpke_rs::libcrux::HpkeLibcrux> {
29        // Pin XWING to the obsolete 0x004D codepoint so the HPKE
30        // suite_id labels (and thus AEAD keys) match v1.9 / v1.10
31        // clients which shipped on hpke-rs 0.4 where 0x004D was the
32        // only XWING variant. Upstream hpke-rs-libcrux 0.6.1 lacks
33        // dispatch for the Obsolete variant; we use a fork via
34        // [patch.crates-io] in this Cargo.toml. See #3661 for the
35        // d14n cutover plan to the canonical 0x647a codepoint.
36        #[allow(deprecated)]
37        let kem = match self {
38            Self::Curve25519 => hpke_rs::hpke_types::KemAlgorithm::DhKem25519,
39            Self::XWingMLKEM768Draft6 => hpke_rs::hpke_types::KemAlgorithm::XWingDraft06Obsolete,
40        };
41        hpke_rs::Hpke::<hpke_rs::libcrux::HpkeLibcrux>::new(
42            hpke_rs::Mode::Base,
43            kem,
44            hpke_rs::hpke_types::KdfAlgorithm::HkdfSha256,
45            hpke_rs::hpke_types::AeadAlgorithm::ChaCha20Poly1305,
46        )
47    }
48}
49
50impl From<WrapperAlgorithm> for WrapperAlgorithmProto {
51    fn from(value: WrapperAlgorithm) -> Self {
52        match value {
53            WrapperAlgorithm::Curve25519 => WrapperAlgorithmProto::Curve25519,
54            WrapperAlgorithm::XWingMLKEM768Draft6 => WrapperAlgorithmProto::XwingMlkem768Draft6,
55        }
56    }
57}
58
59impl TryFrom<WrapperAlgorithmProto> for WrapperAlgorithm {
60    type Error = xmtp_proto::ConversionError;
61    fn try_from(value: WrapperAlgorithmProto) -> Result<Self, Self::Error> {
62        match value {
63            WrapperAlgorithmProto::Curve25519 | WrapperAlgorithmProto::Unspecified => {
64                Ok(WrapperAlgorithm::Curve25519)
65            }
66            WrapperAlgorithmProto::XwingMlkem768Draft6 => Ok(WrapperAlgorithm::XWingMLKEM768Draft6),
67            WrapperAlgorithmProto::SymmetricKey => Err(xmtp_proto::ConversionError::InvalidValue {
68                item: "WrapperAlgorithm",
69                expected: "Curve25519 or XwingMlkem768Draft6",
70                got: format!("{value:?}"),
71            }),
72        }
73    }
74}
75
76impl From<WrapperAlgorithm> for i32 {
77    fn from(value: WrapperAlgorithm) -> Self {
78        let proto_val: WrapperAlgorithmProto = value.into();
79        proto_val as i32
80    }
81}
82
83impl TryFrom<i32> for WrapperAlgorithm {
84    type Error = xmtp_proto::ConversionError;
85    fn try_from(value: i32) -> Result<Self, Self::Error> {
86        let algorithm = match value {
87            1 => WrapperAlgorithm::Curve25519, // WrapperAlgorithmProto::Curve25519
88            2 => WrapperAlgorithm::XWingMLKEM768Draft6, // WrapperAlgorithmProto::XwingMlkem512
89            3 => {
90                return Err(xmtp_proto::ConversionError::InvalidValue {
91                    item: "WrapperAlgorithm",
92                    expected: "1 or 2",
93                    got: value.to_string(),
94                });
95            }
96            _ => WrapperAlgorithm::Curve25519, // Everything else including unknown
97        };
98        Ok(algorithm)
99    }
100}
101
102impl TryFrom<WelcomePointerWrapperAlgorithmProto> for WrapperAlgorithm {
103    type Error = xmtp_proto::ConversionError;
104    fn try_from(value: WelcomePointerWrapperAlgorithmProto) -> Result<Self, Self::Error> {
105        match value {
106            WelcomePointerWrapperAlgorithmProto::XwingMlkem768Draft6 => {
107                Ok(WrapperAlgorithm::XWingMLKEM768Draft6)
108            }
109            _ => Err(xmtp_proto::ConversionError::InvalidValue {
110                item: "WrapperAlgorithm",
111                expected: "XwingMlkem768Draft6",
112                got: format!("{value:?}"),
113            }),
114        }
115    }
116}
117
118impl TryFrom<WrapperAlgorithm> for WelcomePointerWrapperAlgorithmProto {
119    type Error = xmtp_proto::ConversionError;
120    fn try_from(value: WrapperAlgorithm) -> Result<Self, Self::Error> {
121        match value {
122            WrapperAlgorithm::XWingMLKEM768Draft6 => {
123                Ok(WelcomePointerWrapperAlgorithmProto::XwingMlkem768Draft6)
124            }
125            _ => Err(xmtp_proto::ConversionError::InvalidValue {
126                item: "WrapperAlgorithm",
127                expected: "XwingMlkem768Draft6",
128                got: format!("{value:?}"),
129            }),
130        }
131    }
132}
133
134#[derive(Debug)]
135pub struct WrapperEncryptionExtension {
136    pub algorithm: WrapperAlgorithm,
137    pub pub_key_bytes: Vec<u8>,
138}
139
140impl WrapperEncryptionExtension {
141    pub fn new(algorithm: WrapperAlgorithm, pub_key_bytes: Vec<u8>) -> Self {
142        Self {
143            algorithm,
144            pub_key_bytes,
145        }
146    }
147
148    /// The wrapper a key package advertises: its `WelcomeWrapperEncryption`
149    /// extension, or Curve25519 under the package's init key when it has none.
150    pub fn advertised_by(key_package: &KeyPackage) -> Result<Self, ConversionError> {
151        key_package
152            .extensions()
153            .unknown(WELCOME_WRAPPER_ENCRYPTION_EXTENSION_ID)
154            .map_or_else(
155                || {
156                    Ok(Self::new(
157                        WrapperAlgorithm::Curve25519,
158                        key_package.hpke_init_key().as_slice().to_vec(),
159                    ))
160                },
161                Self::try_from,
162            )
163    }
164}
165
166impl TryFrom<WrapperEncryptionExtension> for Extension {
167    type Error = EncodeError;
168
169    fn try_from(value: WrapperEncryptionExtension) -> Result<Self, Self::Error> {
170        let proto_val = WelcomeWrapperEncryptionProto {
171            pub_key: value.pub_key_bytes,
172            algorithm: value.algorithm.into(),
173        };
174        let mut buf = Vec::new();
175        proto_val.encode(&mut buf)?;
176
177        Ok(Extension::Unknown(
178            WELCOME_WRAPPER_ENCRYPTION_EXTENSION_ID,
179            UnknownExtension(buf),
180        ))
181    }
182}
183
184impl TryFrom<&UnknownExtension> for WrapperEncryptionExtension {
185    type Error = ConversionError;
186
187    fn try_from(value: &UnknownExtension) -> Result<Self, Self::Error> {
188        value.0.as_slice().try_into()
189    }
190}
191
192impl TryFrom<&[u8]> for WrapperEncryptionExtension {
193    type Error = ConversionError;
194
195    fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
196        let proto = WelcomeWrapperEncryptionProto::decode(value)?;
197        let algorithm: WrapperAlgorithm = proto.algorithm.try_into()?;
198        Ok(WrapperEncryptionExtension {
199            algorithm,
200            pub_key_bytes: proto.pub_key,
201        })
202    }
203}
204
205#[cfg(test)]
206mod tests {
207    use super::*;
208
209    #[xmtp_common::test]
210    fn test_serialization() {
211        let algorithm = WrapperAlgorithm::XWingMLKEM768Draft6;
212        let pub_key_bytes = xmtp_common::rand_vec::<32>();
213
214        let extension = WrapperEncryptionExtension::new(algorithm, pub_key_bytes.clone());
215
216        let mls_extension: Extension = extension.try_into().unwrap();
217
218        let Extension::Unknown(id, unknown_extension) = mls_extension else {
219            panic!("Expected unknown extension");
220        };
221
222        assert_eq!(id, WELCOME_WRAPPER_ENCRYPTION_EXTENSION_ID);
223
224        let deserialized: WrapperEncryptionExtension = (&unknown_extension).try_into().unwrap();
225
226        assert_eq!(deserialized.algorithm, algorithm);
227        assert_eq!(deserialized.pub_key_bytes, pub_key_bytes);
228    }
229}
xmtp_id/key_package/mls_ext_wrapper_encryption.rs

xmtp_id/key_package/
mls_ext_wrapper_encryption.rs
