@@ -36,7 +36,7 @@ use ruma::{
3636use serde:: { Deserialize , Serialize } ;
3737use serde_json:: { value:: RawValue as RawJsonValue , Value } ;
3838use sha2:: { Digest , Sha256 } ;
39- use tracing:: { debug, info, trace, warn} ;
39+ use tracing:: { debug, info, instrument , trace, warn, Span } ;
4040use vodozemac:: {
4141 olm:: {
4242 Account as InnerAccount , AccountPickle , IdentityKeys , OlmMessage , PreKeyMessage ,
@@ -286,66 +286,80 @@ impl Account {
286286 }
287287
288288 /// Decrypt an Olm message, creating a new Olm session if possible.
289+ #[ instrument( skip( self , message) , fields( session_id) ) ]
289290 async fn decrypt_olm_message (
290291 & self ,
291292 sender : & UserId ,
292293 sender_key : Curve25519PublicKey ,
293294 message : & OlmMessage ,
294295 ) -> OlmResult < ( SessionType , DecryptionResult ) > {
295296 // First try to decrypt using an existing session.
296- let ( session, plaintext) =
297- if let Some ( d ) = self . decrypt_with_existing_sessions ( sender_key, message) . await ? {
298- // Decryption succeeded, de-structure the session/plaintext out of
299- // the Option.
300- ( SessionType :: Existing ( d . 0 ) , d . 1 )
301- } else {
302- // Decryption failed with every known session, let's try to create a
303- // new session.
304- match message {
305- // A new session can only be created using a pre-key message,
306- // return with an error if it isn't one.
307- OlmMessage :: Normal ( _ ) => {
308- warn ! (
309- ?sender_key ,
310- "Failed to decrypt a non-pre-key message with all \
311- available sessions" ,
312- ) ;
313-
314- return Err ( OlmError :: SessionWedged ( sender . to_owned ( ) , sender_key ) ) ;
315- }
297+ let ( session, plaintext) = if let Some ( d ) =
298+ self . decrypt_with_existing_sessions ( sender_key, message) . await ?
299+ {
300+ // Decryption succeeded, de-structure the session/plaintext out of
301+ // the Option.
302+ ( SessionType :: Existing ( d . 0 ) , d . 1 )
303+ } else {
304+ // Decryption failed with every known session, let's try to create a
305+ // new session.
306+ match message {
307+ // A new session can only be created using a pre-key message,
308+ // return with an error if it isn't one.
309+ OlmMessage :: Normal ( _ ) => {
310+ let session_ids = if let Some ( sessions ) =
311+ self . store . get_sessions ( & sender_key . to_base64 ( ) ) . await ?
312+ {
313+ sessions . lock ( ) . await . iter ( ) . map ( |s| s . session_id ( ) . to_owned ( ) ) . collect ( )
314+ } else {
315+ vec ! [ ]
316+ } ;
316317
317- OlmMessage :: PreKey ( m) => {
318- // Create the new session.
319- let result = match self . inner . create_inbound_session ( sender_key, m) . await {
320- Ok ( r) => r,
321- Err ( e) => {
322- warn ! (
323- ?sender_key,
324- session_keys = ?m. session_keys( ) ,
325- "Failed to create a new Olm session from a \
326- pre-key message: {e:?}",
327- ) ;
328- return Err ( OlmError :: SessionWedged ( sender. to_owned ( ) , sender_key) ) ;
329- }
330- } ;
318+ warn ! (
319+ ?session_ids,
320+ "Failed to decrypt a non-pre-key message with all available sessions" ,
321+ ) ;
331322
332- // We need to add the new session to the session cache, otherwise
333- // we might try to create the same session again.
334- // TODO separate the session cache from the storage so we only add
335- // it to the cache but don't store it.
336- let changes = Changes {
337- account : Some ( self . inner . clone ( ) ) ,
338- sessions : vec ! [ result. session. clone( ) ] ,
339- ..Default :: default ( )
340- } ;
341- self . store . save_changes ( changes) . await ?;
323+ return Err ( OlmError :: SessionWedged ( sender. to_owned ( ) , sender_key) ) ;
324+ }
342325
343- ( SessionType :: New ( result. session ) , result. plaintext )
344- }
326+ OlmMessage :: PreKey ( m) => {
327+ // Create the new session.
328+ let result = match self . inner . create_inbound_session ( sender_key, m) . await {
329+ Ok ( r) => r,
330+ Err ( e) => {
331+ warn ! (
332+ session_keys = ?m. session_keys( ) ,
333+ "Failed to create a new Olm session from a pre-key message: {e:?}" ,
334+ ) ;
335+
336+ return Err ( OlmError :: SessionWedged ( sender. to_owned ( ) , sender_key) ) ;
337+ }
338+ } ;
339+
340+ // We need to add the new session to the session cache, otherwise
341+ // we might try to create the same session again.
342+ // TODO: separate the session cache from the storage so we only add
343+ // it to the cache but don't store it.
344+ let changes = Changes {
345+ account : Some ( self . inner . clone ( ) ) ,
346+ sessions : vec ! [ result. session. clone( ) ] ,
347+ ..Default :: default ( )
348+ } ;
349+ self . store . save_changes ( changes) . await ?;
350+
351+ ( SessionType :: New ( result. session ) , result. plaintext )
345352 }
346- } ;
353+ }
354+ } ;
347355
348- trace ! ( ?sender_key, "Successfully decrypted an Olm message" ) ;
356+ let session_id = match & session {
357+ SessionType :: New ( s) => s. session_id ( ) ,
358+ SessionType :: Existing ( s) => s. session_id ( ) ,
359+ } ;
360+
361+ Span :: current ( ) . record ( "session_id" , session_id) ;
362+ trace ! ( "Successfully decrypted an Olm message" ) ;
349363
350364 match self . parse_decrypted_to_device_event ( sender, sender_key, plaintext) . await {
351365 Ok ( result) => Ok ( ( session, result) ) ,
@@ -368,7 +382,6 @@ impl Account {
368382 }
369383
370384 warn ! (
371- sender_key = sender_key. to_base64( ) ,
372385 error = ?e,
373386 "A to-device message was successfully decrypted but \
374387 parsing and checking the event fields failed"
@@ -1045,22 +1058,28 @@ impl ReadOnlyAccount {
10451058 ///
10461059 /// * `message` - A pre-key Olm message that was sent to us by the other
10471060 /// account.
1061+ #[ instrument(
1062+ skip_all,
1063+ fields(
1064+ sender_key = ?their_identity_key,
1065+ session_id = message. session_id( ) ,
1066+ session_keys = ?message. session_keys( ) ,
1067+ )
1068+ ) ]
10481069 pub async fn create_inbound_session (
10491070 & self ,
10501071 their_identity_key : Curve25519PublicKey ,
10511072 message : & PreKeyMessage ,
10521073 ) -> Result < InboundCreationResult , SessionCreationError > {
1053- debug ! (
1054- sender_key = ?their_identity_key,
1055- session_keys = ?message. session_keys( ) ,
1056- "Creating a new Olm session from a pre-key message"
1057- ) ;
1074+ debug ! ( "Creating a new Olm session from a pre-key message" ) ;
10581075
10591076 let result = self . inner . lock ( ) . await . create_inbound_session ( their_identity_key, message) ?;
10601077
10611078 let now = SecondsSinceUnixEpoch :: now ( ) ;
10621079 let session_id = result. session . session_id ( ) ;
10631080
1081+ trace ! ( ?session_id, "Olm session created successfully" ) ;
1082+
10641083 let session = Session {
10651084 user_id : self . user_id . clone ( ) ,
10661085 device_id : self . device_id . clone ( ) ,
0 commit comments