rust/src/lib.rs

1// sophax_mls — MLS (RFC 9420) wrapper for SophaxChat
2//
3// Pure-Rust RustCrypto backend — no OpenSSL, iOS-compatible.
4// Stateless design: every function takes serialized group state, performs one
5// operation, returns new serialized state. Swift owns persistence.
6
7use std::collections::HashSet;
8use std::sync::Arc;
9
10use mls_rs::{
11    identity::{
12        basic::{BasicCredential, BasicIdentityProvider},
13        SigningIdentity,
14    },
15    storage_provider::in_memory::InMemoryGroupStateStorage,
16    CipherSuite, CipherSuiteProvider, Client, CryptoProvider, ExtensionList, MlsMessage,
17};
18use mls_rs_core::{
19    crypto::{SignaturePublicKey, SignatureSecretKey},
20    group::{GroupState, GroupStateStorage},
21};
22use mls_rs_crypto_rustcrypto::RustCryptoProvider;
23
24uniffi::setup_scaffolding!();
25
26const CIPHER_SUITE: CipherSuite = CipherSuite::CURVE25519_AES128;
27
28// ───────────────────────────────────────────────────────────────────── Errors
29
30#[derive(Debug, thiserror::Error, uniffi::Error)]
31pub enum MlsError {
32    #[error("{msg}")]
33    Protocol { msg: String },
34    #[error("Decryption failed or message type is not application data")]
35    DecryptionFailed,
36    #[error("{msg}")]
37    InvalidState { msg: String },
38    #[error("{msg}")]
39    Serialization { msg: String },
40    #[error("Member not found in roster")]
41    UnknownMember,
42    #[error("Expected a Commit message")]
43    InvalidCommit,
44}
45
46impl From<mls_rs::error::MlsError> for MlsError {
47    fn from(e: mls_rs::error::MlsError) -> Self {
48        MlsError::Protocol { msg: e.to_string() }
49    }
50}
51
52// ──────────────────────────────────────────────────────────────────── Records
53
54#[derive(Debug, Clone, uniffi::Record)]
55pub struct MemberWelcome {
56    pub peer_id: String,
57    pub welcome_bytes: Vec<u8>,
58    /// Ratchet tree bytes — send alongside welcome for clients without tree extension
59    pub ratchet_tree_bytes: Vec<u8>,
60}
61
62#[derive(Debug, Clone, uniffi::Record)]
63pub struct CreateGroupOutput {
64    /// Coordinator's group state — persist encrypted immediately
65    pub group_state: Vec<u8>,
66    /// One Welcome per invited member; send via DR unicast
67    pub welcomes: Vec<MemberWelcome>,
68}
69
70#[derive(Debug, Clone, uniffi::Record)]
71pub struct CommitOutput {
72    /// Broadcast to all group members
73    pub commit_bytes: Vec<u8>,
74    /// Send to new member (present only for AddMember commits)
75    pub welcome_bytes: Option<Vec<u8>>,
76    /// Updated state — replace persisted state after distributing commit
77    pub new_group_state: Vec<u8>,
78    pub new_epoch: u64,
79}
80
81#[derive(Debug, Clone, uniffi::Record)]
82pub struct ProcessedCommit {
83    pub added_peer_ids: Vec<String>,
84    pub removed_peer_ids: Vec<String>,
85    pub new_group_state: Vec<u8>,
86    pub new_epoch: u64,
87}
88
89#[derive(Debug, Clone, uniffi::Record)]
90pub struct EncryptOutput {
91    pub ciphertext: Vec<u8>,
92    pub new_group_state: Vec<u8>,
93}
94
95#[derive(Debug, Clone, uniffi::Record)]
96pub struct DecryptOutput {
97    pub plaintext: Vec<u8>,
98    pub new_group_state: Vec<u8>,
99}
100
101// ─────────────────────────────────────────────────────────── MlsClientHandle
102
103/// Local MLS identity derived from the peer's existing Ed25519 key.
104#[derive(uniffi::Object)]
105pub struct MlsClientHandle {
106    peer_id: String,
107    signing_secret: Vec<u8>,
108    crypto: RustCryptoProvider,
109}
110
111#[uniffi::export]
112impl MlsClientHandle {
113    #[uniffi::constructor]
114    pub fn new(peer_id: String, signing_key_bytes: Vec<u8>) -> Result<Arc<Self>, MlsError> {
115        if signing_key_bytes.len() != 32 {
116            return Err(MlsError::InvalidState {
117                msg: format!("signing key must be 32 bytes, got {}", signing_key_bytes.len()),
118            });
119        }
120        Ok(Arc::new(MlsClientHandle {
121            peer_id,
122            signing_secret: signing_key_bytes,
123            crypto: RustCryptoProvider::default(),
124        }))
125    }
126
127    /// Generate a fresh single-use KeyPackage for this identity.
128    pub fn generate_key_package(&self) -> Result<Vec<u8>, MlsError> {
129        let client = self.make_client(InMemoryGroupStateStorage::new())?;
130        let kp = client
131            .generate_key_package_message(Default::default(), Default::default(), None)
132            .map_err(MlsError::from)?;
133        kp.to_bytes()
134            .map_err(|e| MlsError::Serialization { msg: e.to_string() })
135    }
136}
137
138impl MlsClientHandle {
139    /// Build a ready-to-use MLS client, seeded with the provided group state storage.
140    /// `Client::builder()` starts from `BaseConfig` which already includes in-memory
141    /// key-package, PSK, and group-state storage + DefaultMlsRules. We override only
142    /// the group-state storage (to use our pre-seeded instance) and supply the identity.
143    fn make_client(&self, group_storage: InMemoryGroupStateStorage) -> Result<Client<impl mls_rs::client_builder::MlsConfig>, MlsError> {
144        let cs = self.crypto.cipher_suite_provider(CIPHER_SUITE).ok_or_else(|| {
145            MlsError::InvalidState { msg: "cipher suite unavailable".to_string() }
146        })?;
147        let secret = SignatureSecretKey::from(self.signing_secret.clone());
148        let public: SignaturePublicKey = cs
149            .signature_key_derive_public(&secret)
150            .map_err(|e| MlsError::InvalidState { msg: e.to_string() })?;
151        let credential = BasicCredential::new(self.peer_id.as_bytes().to_vec());
152        let signing_identity = SigningIdentity::new(credential.into_credential(), public);
153        Ok(Client::builder()
154            .identity_provider(BasicIdentityProvider)
155            .crypto_provider(self.crypto.clone())
156            .group_state_storage(group_storage)
157            .signing_identity(signing_identity, secret, CIPHER_SUITE)
158            .build())
159    }
160}
161
162// ──────────────────────────────────────────────────────── Free functions
163
164/// Create a new MLS group (coordinator only).
165#[uniffi::export]
166pub fn mls_create_group(
167    client: Arc<MlsClientHandle>,
168    group_id: Vec<u8>,
169    member_key_packages: Vec<Vec<u8>>,
170) -> Result<CreateGroupOutput, MlsError> {
171    let gs = InMemoryGroupStateStorage::new();
172    let c = client.make_client(gs.clone())?;
173    let mut group = c
174        .create_group_with_id(group_id.clone(), ExtensionList::default(), Default::default(), None)
175        .map_err(MlsError::from)?;
176
177    let mut welcomes = Vec::new();
178    for kp_bytes in member_key_packages {
179        let kp = MlsMessage::from_bytes(&kp_bytes)
180            .map_err(|e| MlsError::Serialization { msg: e.to_string() })?;
181        let peer_id = peer_id_from_key_package(&kp)?;
182        let commit = group.commit_builder().add_member(kp).map_err(MlsError::from)?.build().map_err(MlsError::from)?;
183        group.apply_pending_commit().map_err(MlsError::from)?;
184        let welcome_bytes = commit.welcome_messages.into_iter().next()
185            .ok_or_else(|| MlsError::InvalidState { msg: "no welcome message".to_string() })?
186            .to_bytes().map_err(|e| MlsError::Serialization { msg: e.to_string() })?;
187        let ratchet_tree_bytes = group.export_tree().to_bytes()
188            .map_err(|e| MlsError::Serialization { msg: e.to_string() })?;
189        welcomes.push(MemberWelcome { peer_id, welcome_bytes, ratchet_tree_bytes });
190    }
191
192    group.write_to_storage().map_err(MlsError::from)?;
193    let group_state = extract_state(&gs, &group_id)?;
194    Ok(CreateGroupOutput { group_state, welcomes })
195}
196
197/// Join a group by processing a Welcome message.
198#[uniffi::export]
199pub fn mls_process_welcome(
200    client: Arc<MlsClientHandle>,
201    welcome_bytes: Vec<u8>,
202    ratchet_tree_bytes: Option<Vec<u8>>,
203) -> Result<Vec<u8>, MlsError> {
204    let gs = InMemoryGroupStateStorage::new();
205    let c = client.make_client(gs.clone())?;
206    let welcome = MlsMessage::from_bytes(&welcome_bytes)
207        .map_err(|e| MlsError::Serialization { msg: e.to_string() })?;
208    let ratchet_tree = ratchet_tree_bytes
209        .as_deref()
210        .map(|bytes| mls_rs::group::ExportedTree::from_bytes(bytes))
211        .transpose()
212        .map_err(|e| MlsError::Serialization { msg: e.to_string() })?;
213    let (mut group, _) = c.join_group(ratchet_tree, &welcome, None).map_err(MlsError::from)?;
214    group.write_to_storage().map_err(MlsError::from)?;
215    let group_id = group.group_id().to_vec();
216    extract_state(&gs, &group_id)
217}
218
219/// Encrypt an application message.
220#[uniffi::export]
221pub fn mls_encrypt(
222    client: Arc<MlsClientHandle>,
223    group_id: Vec<u8>,
224    group_state: Vec<u8>,
225    plaintext: Vec<u8>,
226) -> Result<EncryptOutput, MlsError> {
227    let mut gs = InMemoryGroupStateStorage::new();
228    seed_storage(&mut gs, &group_id, group_state)?;
229    let c = client.make_client(gs.clone())?;
230    let mut group = c.load_group(&group_id).map_err(MlsError::from)?;
231    let msg = group.encrypt_application_message(&plaintext, Default::default()).map_err(MlsError::from)?;
232    let ciphertext = msg.to_bytes().map_err(|e| MlsError::Serialization { msg: e.to_string() })?;
233    group.write_to_storage().map_err(MlsError::from)?;
234    let new_group_state = extract_state(&gs, &group_id)?;
235    Ok(EncryptOutput { ciphertext, new_group_state })
236}
237
238/// Decrypt an MLS application message.
239#[uniffi::export]
240pub fn mls_decrypt(
241    client: Arc<MlsClientHandle>,
242    group_id: Vec<u8>,
243    group_state: Vec<u8>,
244    ciphertext: Vec<u8>,
245) -> Result<DecryptOutput, MlsError> {
246    let mut gs = InMemoryGroupStateStorage::new();
247    seed_storage(&mut gs, &group_id, group_state)?;
248    let c = client.make_client(gs.clone())?;
249    let mut group = c.load_group(&group_id).map_err(MlsError::from)?;
250    let msg = MlsMessage::from_bytes(&ciphertext)
251        .map_err(|e| MlsError::Serialization { msg: e.to_string() })?;
252    let received = group.process_incoming_message(msg).map_err(MlsError::from)?;
253    let plaintext = match received {
254        mls_rs::group::ReceivedMessage::ApplicationMessage(app_msg) => app_msg.data().to_vec(),
255        _ => return Err(MlsError::DecryptionFailed),
256    };
257    group.write_to_storage().map_err(MlsError::from)?;
258    let new_group_state = extract_state(&gs, &group_id)?;
259    Ok(DecryptOutput { plaintext, new_group_state })
260}
261
262/// Process a Commit broadcast (coordinator's output). Validates coordinator in Swift.
263#[uniffi::export]
264pub fn mls_process_commit(
265    client: Arc<MlsClientHandle>,
266    group_id: Vec<u8>,
267    group_state: Vec<u8>,
268    commit_bytes: Vec<u8>,
269) -> Result<ProcessedCommit, MlsError> {
270    let mut gs = InMemoryGroupStateStorage::new();
271    seed_storage(&mut gs, &group_id, group_state)?;
272    let c = client.make_client(gs.clone())?;
273    let mut group = c.load_group(&group_id).map_err(MlsError::from)?;
274
275    // Snapshot current roster before commit
276    let before: HashSet<String> = roster_peer_ids(&group);
277
278    let msg = MlsMessage::from_bytes(&commit_bytes)
279        .map_err(|e| MlsError::Serialization { msg: e.to_string() })?;
280    let received = group.process_incoming_message(msg).map_err(MlsError::from)?;
281    match received {
282        mls_rs::group::ReceivedMessage::Commit(_) => {}
283        _ => return Err(MlsError::InvalidCommit),
284    };
285
286    let after: HashSet<String> = roster_peer_ids(&group);
287    let added_peer_ids: Vec<String> = after.difference(&before).cloned().collect();
288    let removed_peer_ids: Vec<String> = before.difference(&after).cloned().collect();
289    let new_epoch = group.current_epoch();
290
291    group.write_to_storage().map_err(MlsError::from)?;
292    let new_group_state = extract_state(&gs, &group_id)?;
293    Ok(ProcessedCommit { added_peer_ids, removed_peer_ids, new_group_state, new_epoch })
294}
295
296/// Add a member (coordinator only). Returns commit to broadcast + welcome to send.
297#[uniffi::export]
298pub fn mls_add_member(
299    client: Arc<MlsClientHandle>,
300    group_id: Vec<u8>,
301    group_state: Vec<u8>,
302    key_package_bytes: Vec<u8>,
303) -> Result<CommitOutput, MlsError> {
304    let mut gs = InMemoryGroupStateStorage::new();
305    seed_storage(&mut gs, &group_id, group_state)?;
306    let c = client.make_client(gs.clone())?;
307    let mut group = c.load_group(&group_id).map_err(MlsError::from)?;
308    let kp = MlsMessage::from_bytes(&key_package_bytes)
309        .map_err(|e| MlsError::Serialization { msg: e.to_string() })?;
310    let commit = group.commit_builder().add_member(kp).map_err(MlsError::from)?.build().map_err(MlsError::from)?;
311    group.apply_pending_commit().map_err(MlsError::from)?;
312    let commit_bytes = commit.commit_message.to_bytes()
313        .map_err(|e| MlsError::Serialization { msg: e.to_string() })?;
314    let welcome_bytes = commit.welcome_messages.into_iter().next()
315        .map(|w| w.to_bytes())
316        .transpose()
317        .map_err(|e| MlsError::Serialization { msg: e.to_string() })?;
318    let new_epoch = group.current_epoch();
319    group.write_to_storage().map_err(MlsError::from)?;
320    let new_group_state = extract_state(&gs, &group_id)?;
321    Ok(CommitOutput { commit_bytes, welcome_bytes, new_group_state, new_epoch })
322}
323
324/// Remove a member (coordinator only).
325#[uniffi::export]
326pub fn mls_remove_member(
327    client: Arc<MlsClientHandle>,
328    group_id: Vec<u8>,
329    group_state: Vec<u8>,
330    peer_id: String,
331) -> Result<CommitOutput, MlsError> {
332    let mut gs = InMemoryGroupStateStorage::new();
333    seed_storage(&mut gs, &group_id, group_state)?;
334    let c = client.make_client(gs.clone())?;
335    let mut group = c.load_group(&group_id).map_err(MlsError::from)?;
336    let peer_id_bytes = peer_id.as_bytes().to_vec();
337    let leaf_index = group.roster().members_iter()
338        .find(|m| m.signing_identity().credential.as_basic()
339            .map(|c| c.identifier == peer_id_bytes).unwrap_or(false))
340        .map(|m| m.index())
341        .ok_or(MlsError::UnknownMember)?;
342    let commit = group.commit_builder().remove_member(leaf_index).map_err(MlsError::from)?.build().map_err(MlsError::from)?;
343    group.apply_pending_commit().map_err(MlsError::from)?;
344    let commit_bytes = commit.commit_message.to_bytes()
345        .map_err(|e| MlsError::Serialization { msg: e.to_string() })?;
346    let new_epoch = group.current_epoch();
347    group.write_to_storage().map_err(MlsError::from)?;
348    let new_group_state = extract_state(&gs, &group_id)?;
349    Ok(CommitOutput { commit_bytes, welcome_bytes: None, new_group_state, new_epoch })
350}
351
352// ─────────────────────────────────────────────────────────────────── Helpers
353
354fn peer_id_from_key_package(kp_msg: &MlsMessage) -> Result<String, MlsError> {
355    let kp = kp_msg.as_key_package()
356        .ok_or_else(|| MlsError::Serialization { msg: "not a key package".to_string() })?;
357    let basic = kp.signing_identity().credential.as_basic()
358        .ok_or_else(|| MlsError::Serialization { msg: "key package has no BasicCredential".to_string() })?;
359    String::from_utf8(basic.identifier.clone())
360        .map_err(|e| MlsError::Serialization { msg: e.to_string() })
361}
362
363fn seed_storage(
364    storage: &mut InMemoryGroupStateStorage,
365    group_id: &[u8],
366    state_bytes: Vec<u8>,
367) -> Result<(), MlsError> {
368    let gs = GroupState { id: group_id.to_vec(), data: state_bytes.into() };
369    storage.write(gs, vec![], vec![])
370        .map_err(|e| MlsError::Serialization { msg: format!("{e}") })
371}
372
373fn extract_state(
374    storage: &InMemoryGroupStateStorage,
375    group_id: &[u8],
376) -> Result<Vec<u8>, MlsError> {
377    storage.state(group_id)
378        .map_err(|e| MlsError::Serialization { msg: format!("{e}") })?
379        .map(|z| z.to_vec())
380        .ok_or(MlsError::InvalidState { msg: "group state missing after write".to_string() })
381}
382
383fn roster_peer_ids<C: mls_rs::client_builder::MlsConfig>(group: &mls_rs::Group<C>) -> HashSet<String> {
384    group.roster().members_iter()
385        .filter_map(|m| {
386            m.signing_identity().credential.as_basic().and_then(|c| {
387                String::from_utf8(c.identifier.clone()).ok()
388            })
389        })
390        .collect()
391}