Skip to main content

offline_protocol_mls/
session.rs

1//! 1:1 session management using MLS groups.
2//!
3//! This module provides a higher-level abstraction for 1:1 encrypted messaging
4//! by using 2-person MLS groups under the hood.
5
6use crate::error::{MlsError, Result};
7use crate::group::GroupManager;
8use crate::provider::MlsProvider;
9use crate::storage::MlsStorage;
10use crate::types::{EncryptedMessage, GroupId, GroupInfo, MlsMessageType, WelcomeMessage};
11
12use openmls::prelude::tls_codec::{Deserialize as TlsDeserialize, Serialize as TlsSerialize};
13use openmls::prelude::*;
14use openmls_traits::signatures::Signer;
15use std::sync::Arc;
16use tracing::{debug, info};
17
18/// Manages 1:1 encrypted sessions.
19pub struct SessionManager {
20    /// The local user's ID.
21    user_id: String,
22
23    /// Group manager for underlying MLS operations.
24    group_manager: GroupManager,
25}
26
27impl SessionManager {
28    /// Creates a new session manager.
29    pub fn new(user_id: String, storage: Arc<dyn MlsStorage>, provider: MlsProvider) -> Self {
30        let group_manager = GroupManager::new(storage, provider);
31        Self {
32            user_id,
33            group_manager,
34        }
35    }
36
37    /// Returns a reference to the underlying group manager.
38    pub fn group_manager(&self) -> &GroupManager {
39        &self.group_manager
40    }
41
42    /// Gets the session ID for a 1:1 conversation with another user.
43    ///
44    /// Fails with [`crate::MlsError::InvalidGroupId`] if `other_user_id`
45    /// contains storage-hostile characters (the session id is a raw storage
46    /// key).
47    pub fn get_session_id(&self, other_user_id: &str) -> Result<GroupId> {
48        GroupId::for_session(&self.user_id, other_user_id)
49    }
50
51    /// Checks if a session exists with another user.
52    pub fn has_session(&self, other_user_id: &str) -> Result<bool> {
53        let session_id = self.get_session_id(other_user_id)?;
54        let group = self.group_manager.load_group(&session_id)?;
55        Ok(group.is_some())
56    }
57
58    /// Creates a new session with another user using their key package.
59    pub fn create_session(
60        &self,
61        other_user_id: &str,
62        their_key_package: KeyPackage,
63        credential_with_key: &CredentialWithKey,
64        signer: &impl Signer,
65    ) -> Result<WelcomeMessage> {
66        let session_id = self.get_session_id(other_user_id)?;
67
68        if self.group_manager.load_group(&session_id)?.is_some() {
69            return Err(MlsError::GroupCreation(format!(
70                "Session already exists with user: {}",
71                other_user_id
72            )));
73        }
74
75        let mut group =
76            self.group_manager
77                .create_group(&session_id, credential_with_key, signer)?;
78
79        let (_commit, welcome) =
80            self.group_manager
81                .add_member(&mut group, their_key_package, signer)?;
82
83        self.group_manager.save_group(&session_id, &group)?;
84
85        let welcome_bytes = welcome
86            .tls_serialize_detached()
87            .map_err(|e| MlsError::Serialization(e.to_string()))?;
88
89        info!(
90            session_id = %session_id,
91            other_user = %other_user_id,
92            "Created new 1:1 session"
93        );
94
95        Ok(WelcomeMessage {
96            group_id: session_id,
97            welcome_data: welcome_bytes,
98            inviter_id: self.user_id.clone(),
99            group_name: None,
100            timestamp_ms: chrono::Utc::now().timestamp_millis() as u64,
101        })
102    }
103
104    /// Verifies a Welcome's `group_id` names the `(self, inviter)` session slot
105    /// — the only 1:1 slot an authenticated inviter may install (SEC-M6) — and
106    /// returns the mismatch error **without touching any state**, so callers
107    /// can reject a forged/tampered slot before performing any mutation.
108    ///
109    /// `welcome_msg.group_id` is an attacker-controllable wire field that
110    /// becomes the raw storage key for the adopted group, and the inviter is
111    /// authenticated one layer up (`handle_welcome_message` drops any Welcome
112    /// whose `inviter_id` != the transport-authenticated sender). The only 1:1
113    /// slot this Welcome may legitimately install is therefore
114    /// `session:<self>:<inviter>`. Without this check an authenticated peer
115    /// could name a *third* party's slot (e.g. "session:alice:bob") and, via
116    /// the stage-then-replace adopt in `join_group_replacing`, overwrite or
117    /// seed the victim's session with that party — hijacking it so the victim's
118    /// `encrypt_for_user` output goes to the attacker's group. This mirrors the
119    /// invite-side (SEC-M5) and decrypt-side (SEC-M1) identity bindings on the
120    /// previously-unguarded join side. The sibling group-Welcome path is
121    /// guarded separately by the reserved-namespace check in
122    /// [`crate::MlsManager::join_group`], which refuses any `session:`-prefixed
123    /// `group_id`. Checked before deserialization so a tampered slot is
124    /// rejected regardless of the Welcome blob.
125    pub(crate) fn verify_welcome_slot(&self, welcome_msg: &WelcomeMessage) -> Result<()> {
126        let expected_id = GroupId::for_session(&self.user_id, &welcome_msg.inviter_id)?;
127        if welcome_msg.group_id != expected_id {
128            return Err(MlsError::WelcomeIdentityMismatch {
129                expected: expected_id.to_string(),
130                found: welcome_msg.group_id.to_string(),
131            });
132        }
133        Ok(())
134    }
135
136    /// Joins a session using a Welcome message.
137    ///
138    /// Implements a **non-destructive** "welcome-wins" strategy: if we already
139    /// created a session for this peer (both-create race), the incoming Welcome
140    /// is staged first and only replaces our group once staging succeeds (see
141    /// [`crate::group::GroupManager::join_group_replacing`]). A retransmitted
142    /// Welcome we already adopted, or a first-contact key-package race, fails to
143    /// stage and leaves the existing converged group intact — so adoption is
144    /// idempotent and never bricks a working session.
145    pub fn join_session(&self, welcome_msg: &WelcomeMessage) -> Result<GroupInfo> {
146        // SEC-M6: bind the Welcome's session slot to the (self, inviter) pair
147        // before touching any group state (see `verify_welcome_slot`).
148        self.verify_welcome_slot(welcome_msg)?;
149
150        // Deserialize the Welcome from the MlsMessageOut bytes
151        let mls_msg = MlsMessageIn::tls_deserialize_exact(&welcome_msg.welcome_data)
152            .map_err(|e| MlsError::Deserialization(e.to_string()))?;
153
154        // Extract the Welcome from the message
155        let welcome = match mls_msg.extract() {
156            MlsMessageBodyIn::Welcome(w) => w,
157            _ => {
158                return Err(MlsError::WelcomeProcessing(
159                    "Not a Welcome message".to_string(),
160                ))
161            }
162        };
163
164        // Stage-then-swap: replaces any existing group only after the incoming
165        // Welcome stages successfully, so a failed stage is a safe no-op.
166        let group = self
167            .group_manager
168            .join_group_replacing(welcome, &welcome_msg.group_id)?;
169        let info = self
170            .group_manager
171            .get_group_info(&group, &welcome_msg.group_id);
172
173        info!(
174            session_id = %welcome_msg.group_id,
175            inviter = %welcome_msg.inviter_id,
176            "Joined/adopted 1:1 session (non-destructive)"
177        );
178
179        Ok(info)
180    }
181
182    /// Encrypts a message for the session partner.
183    pub fn encrypt_message(
184        &self,
185        other_user_id: &str,
186        plaintext: &[u8],
187        signer: &impl Signer,
188    ) -> Result<EncryptedMessage> {
189        let session_id = self.get_session_id(other_user_id)?;
190
191        let mut group = self
192            .group_manager
193            .load_group(&session_id)?
194            .ok_or_else(|| MlsError::SessionNotFound(other_user_id.to_string()))?;
195
196        let mls_message = self
197            .group_manager
198            .encrypt_message(&mut group, plaintext, signer)?;
199
200        self.group_manager.save_group(&session_id, &group)?;
201
202        let ciphertext = mls_message
203            .tls_serialize_detached()
204            .map_err(|e| MlsError::Serialization(e.to_string()))?;
205
206        Ok(EncryptedMessage {
207            group_id: session_id,
208            message_type: MlsMessageType::Application,
209            epoch: group.epoch().as_u64(),
210            ciphertext,
211            sender_id: self.user_id.clone(),
212            timestamp_ms: chrono::Utc::now().timestamp_millis() as u64,
213        })
214    }
215
216    /// Decrypts a message from a session.
217    ///
218    /// `claimed_sender` is the transport-level sender this message will be
219    /// attributed to; it must match the MLS-authenticated credential.
220    ///
221    /// The envelope's `group_id` is bound to `claimed_sender` before the group
222    /// is loaded, because every failure OpenMLS raises below happens *before*
223    /// it authenticates anything — see [`MlsError::SessionIdentityMismatch`].
224    /// This is the single chokepoint for the binding: `decrypt_from_user` is
225    /// reachable both from the receive path and from the public
226    /// `manual_mls_decrypt_from_user` API, so checking it in either caller
227    /// alone leaves the other as a bypass.
228    pub fn decrypt_message(
229        &self,
230        encrypted: &EncryptedMessage,
231        claimed_sender: &str,
232    ) -> Result<Option<Vec<u8>>> {
233        let expected_id = GroupId::for_session(&self.user_id, claimed_sender)?;
234        if encrypted.group_id != expected_id {
235            return Err(MlsError::SessionIdentityMismatch {
236                expected: expected_id.to_string(),
237                found: encrypted.group_id.to_string(),
238            });
239        }
240
241        let mut group = self
242            .group_manager
243            .load_group(&encrypted.group_id)?
244            .ok_or_else(|| MlsError::GroupNotFound(encrypted.group_id.to_string()))?;
245
246        let mls_message = MlsMessageIn::tls_deserialize_exact(&encrypted.ciphertext)
247            .map_err(|e| MlsError::Deserialization(e.to_string()))?;
248
249        let result = self.group_manager.decrypt_message(
250            &mut group,
251            &encrypted.group_id,
252            mls_message,
253            claimed_sender,
254        )?;
255
256        self.group_manager.save_group(&encrypted.group_id, &group)?;
257
258        Ok(result)
259    }
260
261    /// Lists the peers this device has a 1:1 session with.
262    ///
263    /// A session id is `session:<a>:<b>` where exactly one of the two halves
264    /// is us — ids cannot contain `:` (`GroupId::validate` runs the storage-key
265    /// charset check per segment), so the split is unambiguous. Anything that
266    /// does not have that exact shape is skipped rather than guessed at: a slot
267    /// naming two other parties has no "other peer" to report, and returning
268    /// one half of it anyway would surface a peer this device has no session
269    /// with.
270    pub fn list_sessions(&self) -> Result<Vec<String>> {
271        let sessions: Vec<String> = self
272            .group_manager
273            .list_groups()?
274            .into_iter()
275            .filter_map(|g| {
276                let parts: Vec<&str> = g.as_str().split(':').collect();
277                match parts.as_slice() {
278                    ["session", a, b] if *a == self.user_id => Some(b.to_string()),
279                    ["session", a, b] if *b == self.user_id => Some(a.to_string()),
280                    _ => None,
281                }
282            })
283            .collect();
284
285        Ok(sessions)
286    }
287
288    /// Deletes a session.
289    pub fn delete_session(&self, other_user_id: &str) -> Result<()> {
290        let session_id = self.get_session_id(other_user_id)?;
291        self.group_manager.delete_group(&session_id)?;
292        debug!(session_id = %session_id, "Deleted session");
293        Ok(())
294    }
295
296    /// Gets information about a session.
297    pub fn get_session_info(&self, other_user_id: &str) -> Result<Option<GroupInfo>> {
298        let session_id = self.get_session_id(other_user_id)?;
299        let group = self.group_manager.load_group(&session_id)?;
300        Ok(group.map(|g| self.group_manager.get_group_info(&g, &session_id)))
301    }
302}
303
304#[cfg(test)]
305mod tests {
306    use super::*;
307    use crate::storage::InMemoryStorage;
308    use crate::storage_adapter::MlsStorageAdapter;
309
310    fn create_test_session_manager(user_id: &str) -> SessionManager {
311        let storage = Arc::new(InMemoryStorage::new());
312        let adapter = MlsStorageAdapter::new(storage.clone());
313        let provider = MlsProvider::new(adapter);
314        SessionManager::new(user_id.to_string(), storage, provider)
315    }
316
317    #[test]
318    fn test_session_id_generation() {
319        let manager = create_test_session_manager("alice");
320        let id1 = manager.get_session_id("bob").unwrap();
321        let id2 = GroupId::for_session("bob", "alice").unwrap();
322        assert_eq!(id1, id2);
323        assert!(id1.as_str().starts_with("session:"));
324    }
325
326    #[test]
327    fn test_no_session_initially() {
328        let manager = create_test_session_manager("alice");
329        assert!(!manager.has_session("bob").unwrap());
330        assert!(manager.list_sessions().unwrap().is_empty());
331    }
332
333    #[test]
334    fn test_session_id_is_deterministic() {
335        // Session IDs should be the same regardless of which user creates them
336        let alice_manager = create_test_session_manager("alice");
337        let bob_manager = create_test_session_manager("bob");
338
339        let alice_to_bob = alice_manager.get_session_id("bob").unwrap();
340        let bob_to_alice = bob_manager.get_session_id("alice").unwrap();
341
342        // Both should produce the same session ID
343        assert_eq!(alice_to_bob, bob_to_alice);
344
345        // The session ID should have a consistent format with sorted user IDs
346        assert_eq!(alice_to_bob.as_str(), "session:alice:bob");
347    }
348
349    #[test]
350    fn test_delete_session() {
351        let manager = create_test_session_manager("alice");
352
353        // Initially no session
354        assert!(!manager.has_session("bob").unwrap());
355
356        // Delete non-existent session should not panic (delete_group handles this)
357        // This is important for welcome-wins where we delete before knowing if exists
358        let result = manager.delete_session("bob");
359        // Should succeed even if nothing to delete
360        assert!(result.is_ok());
361    }
362}