1use 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
18pub struct SessionManager {
20 user_id: String,
22
23 group_manager: GroupManager,
25}
26
27impl SessionManager {
28 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 pub fn group_manager(&self) -> &GroupManager {
39 &self.group_manager
40 }
41
42 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 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 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 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 pub fn join_session(&self, welcome_msg: &WelcomeMessage) -> Result<GroupInfo> {
146 self.verify_welcome_slot(welcome_msg)?;
149
150 let mls_msg = MlsMessageIn::tls_deserialize_exact(&welcome_msg.welcome_data)
152 .map_err(|e| MlsError::Deserialization(e.to_string()))?;
153
154 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 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 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 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 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 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 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 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 assert_eq!(alice_to_bob, bob_to_alice);
344
345 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 assert!(!manager.has_session("bob").unwrap());
355
356 let result = manager.delete_session("bob");
359 assert!(result.is_ok());
361 }
362}