fix signal issues not found

rust/src/api/messages/incoming/client2client/messages.rs

@@ -11,6 +11,7 @@ use crate::bridge::callbacks::get_callbacks;
 use crate::context::Context;
 use crate::database::app::tables::{Contact, MediaFile, NewReceipt, Receipt};
 use crate::error::{twonly_error, Result, TwonlyError};
+use crate::services::contacts::ContactService;
 use crate::utils::new_uuid_v4;
 use prost::Message as ProstMessage;
 use proto::encrypted_content::error_messages::Type;
@@ -252,6 +253,37 @@ pub(crate) fn spawn_receipt_delivery(ctx: &Arc<Context>, receipt_id: String) {
     });
 }
 
+async fn encrypt_v2_with_session_recovery(
+    ctx: &Arc<Context>,
+    contact_id: i64,
+    plaintext: Vec<u8>,
+) -> Result<Vec<u8>> {
+    let encrypt = |plaintext| async move {
+        let engine = ctx.get_signal_engine().lock().await;
+        engine
+            .as_ref()
+            .ok_or(TwonlyError::SignalIdentityNotFound)?
+            .encrypt_message(contact_id.to_string(), 1, plaintext)
+            .await
+    };
+
+    match encrypt(plaintext.clone()).await {
+        Err(TwonlyError::Signal(message))
+            if message.contains("session with") && message.contains("not found") =>
+        {
+            tracing::warn!(
+                contact_id,
+                "Signal session missing; rebuilding it from the server prekey bundle"
+            );
+            ContactService::new(ctx)
+                .establish_signal_session(contact_id)
+                .await?;
+            encrypt(plaintext).await
+        }
+        result => result,
+    }
+}
+
 pub(crate) async fn send_queued_receipt(ctx: &Arc<Context>, receipt_id: &str) -> Result<()> {
     #[cfg(not(debug_assertions))]
     {
@@ -314,14 +346,8 @@ pub(crate) async fn send_queued_receipt(ctx: &Arc<Context>, receipt_id: &str) ->
                 .encrypted_content
                 .take()
                 .ok_or_else(|| TwonlyError::Generic("queued V2 message has no plaintext".into()))?;
-            let engine = ctx.get_signal_engine().lock().await;
-            message.encrypted_content = Some(
-                engine
-                    .as_ref()
-                    .ok_or(TwonlyError::SignalIdentityNotFound)?
-                    .encrypt_message(row.contact_id.to_string(), 1, plaintext)
-                    .await?,
-            );
+            message.encrypted_content =
+                Some(encrypt_v2_with_session_recovery(ctx, row.contact_id, plaintext).await?);
         }
         _ => {}
     }
@@ -375,7 +401,7 @@ pub(crate) async fn send_queued_receipt(ctx: &Arc<Context>, receipt_id: &str) ->
 }
 
 pub(crate) async fn prepare_queued_receipt(
-    ctx: &Context,
+    ctx: &Arc<Context>,
     receipt_id: &str,
 ) -> Result<Option<(Vec<u8>, Option<Vec<u8>>)>> {
     let database = ctx.get_app_database().await;
@@ -413,15 +439,8 @@ pub(crate) async fn prepare_queued_receipt(
                 .take()
                 .ok_or_else(|| TwonlyError::Generic("queued message has no content".into()))?;
 
-            let engine = ctx.get_signal_engine().lock().await;
-
-            message.encrypted_content = Some(
-                engine
-                    .as_ref()
-                    .ok_or(TwonlyError::SignalIdentityNotFound)?
-                    .encrypt_message(row.contact_id.to_string(), 1, plaintext)
-                    .await?,
-            );
+            message.encrypted_content =
+                Some(encrypt_v2_with_session_recovery(ctx, row.contact_id, plaintext).await?);
         }
         _ => {}
     }

rust/src/services/contacts.rs

@@ -34,6 +34,47 @@ impl ContactService {
             }
         };
 
+        self.process_user_prekey_bundle(&user).await?;
+
+        let database = self.ctx.get_app_database().await;
+        let mut transaction = database.pool.begin().await?;
+        UpdateContact::builder()
+            .user_id(user.user_id)
+            .username(username)
+            .signal_version("v2".to_owned())
+            .requested(false)
+            .blocked(false)
+            .deleted_by_user(false)
+            .build()
+            .insert_on_conflict_update(&mut transaction)
+            .await?;
+        transaction.commit().await?;
+        database.notify_committed(["contacts"]);
+
+        self.send_contact_request(
+            user.user_id,
+            encrypted_content::contact_request::Type::Request,
+            blocking,
+        )
+        .await
+    }
+
+    pub(crate) async fn establish_signal_session(&self, user_id: i64) -> Result<()> {
+        let user = match Server::get_user_by_id(&self.ctx, user_id).await? {
+            ServerResult::Ok(user) => user,
+            ServerResult::ErrorCode(code) => {
+                return Err(TwonlyError::Generic(format!(
+                    "Could not load prekey bundle for user {user_id}: server error {code}"
+                )));
+            }
+        };
+        self.process_user_prekey_bundle(&user).await
+    }
+
+    async fn process_user_prekey_bundle(
+        &self,
+        user: &crate::api::proto::server_to_client::response::UserData,
+    ) -> Result<()> {
         let missing = TwonlyError::ApiResponseMissingField;
         let pqc_bundle = user.pqc_bundle.as_ref().ok_or(missing("pqc_bundle"))?;
         let identity_key = user
@@ -86,28 +127,7 @@ impl ContactService {
                 },
             )
             .await?;
-
-        let database = self.ctx.get_app_database().await;
-        let mut transaction = database.pool.begin().await?;
-        UpdateContact::builder()
-            .user_id(user.user_id)
-            .username(username)
-            .signal_version("v2".to_owned())
-            .requested(false)
-            .blocked(false)
-            .deleted_by_user(false)
-            .build()
-            .insert_on_conflict_update(&mut transaction)
-            .await?;
-        transaction.commit().await?;
-        database.notify_committed(["contacts"]);
-
-        self.send_contact_request(
-            user.user_id,
-            encrypted_content::contact_request::Type::Request,
-            blocking,
-        )
-        .await
+        Ok(())
     }
 
     pub async fn accept_request(&self, contact_id: i64, blocking: bool) -> Result<()> {

rust/tests/api/tester.rs

@@ -469,11 +469,6 @@ impl Tester {
 
         // Reload the API configuration so the new login_token/user_id is picked up by the API Client
         ApiRuntime::reload_configuration(&self.context).await?;
-
-        // After the user has register, it should reconnected and reauthenticated
-        ApiRuntime::close(&self.context).await?;
-        ApiRuntime::connect(&self.context).await?;
-
         self.wait_until(ApiConnectionState::Authenticated).await?;
 
         Server::generate_and_upload_pqc_pre_keys(&self.context).await?;