xmtp_id/key_package/
mls_ext_wrapper_encryption.rs

1use openmls::prelude::UnknownExtension;
2use openmls::prelude::{Ciphersuite, Extension};
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
149impl TryFrom<WrapperEncryptionExtension> for Extension {
150    type Error = EncodeError;
151
152    fn try_from(value: WrapperEncryptionExtension) -> Result<Self, Self::Error> {
153        let proto_val = WelcomeWrapperEncryptionProto {
154            pub_key: value.pub_key_bytes,
155            algorithm: value.algorithm.into(),
156        };
157        let mut buf = Vec::new();
158        proto_val.encode(&mut buf)?;
159
160        Ok(Extension::Unknown(
161            WELCOME_WRAPPER_ENCRYPTION_EXTENSION_ID,
162            UnknownExtension(buf),
163        ))
164    }
165}
166
167impl TryFrom<&UnknownExtension> for WrapperEncryptionExtension {
168    type Error = ConversionError;
169
170    fn try_from(value: &UnknownExtension) -> Result<Self, Self::Error> {
171        value.0.as_slice().try_into()
172    }
173}
174
175impl TryFrom<&[u8]> for WrapperEncryptionExtension {
176    type Error = ConversionError;
177
178    fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
179        let proto = WelcomeWrapperEncryptionProto::decode(value)?;
180        let algorithm: WrapperAlgorithm = proto.algorithm.try_into()?;
181        Ok(WrapperEncryptionExtension {
182            algorithm,
183            pub_key_bytes: proto.pub_key,
184        })
185    }
186}
187
188#[cfg(test)]
189mod tests {
190    use super::*;
191
192    #[xmtp_common::test]
193    fn test_serialization() {
194        let algorithm = WrapperAlgorithm::XWingMLKEM768Draft6;
195        let pub_key_bytes = xmtp_common::rand_vec::<32>();
196
197        let extension = WrapperEncryptionExtension::new(algorithm, pub_key_bytes.clone());
198
199        let mls_extension: Extension = extension.try_into().unwrap();
200
201        let Extension::Unknown(id, unknown_extension) = mls_extension else {
202            panic!("Expected unknown extension");
203        };
204
205        assert_eq!(id, WELCOME_WRAPPER_ENCRYPTION_EXTENSION_ID);
206
207        let deserialized: WrapperEncryptionExtension = (&unknown_extension).try_into().unwrap();
208
209        assert_eq!(deserialized.algorithm, algorithm);
210        assert_eq!(deserialized.pub_key_bytes, pub_key_bytes);
211    }
212}
xmtp_id/key_package/mls_ext_wrapper_encryption.rs

xmtp_id/key_package/
mls_ext_wrapper_encryption.rs
