Skip to content

Commit

Permalink
crypto: Helpers for NSE test
Browse files Browse the repository at this point in the history
  • Loading branch information
andybalaam committed May 10, 2024
1 parent bca8452 commit a5517ff
Show file tree
Hide file tree
Showing 4 changed files with 69 additions and 9 deletions.
5 changes: 4 additions & 1 deletion crates/matrix-sdk-crypto/src/machine.rs
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ use tokio::sync::Mutex;
use tracing::{
debug, error,
field::{debug, display},
info, instrument, warn, Span,
info, instrument, trace, warn, Span,
};
use vodozemac::{
megolm::{DecryptionError, SessionOrdering},
Expand Down Expand Up @@ -1918,9 +1918,12 @@ impl OlmMachine {
&self,
generation: &Mutex<Option<u64>>,
) -> StoreResult<()> {
trace!("initialize_crypto_store_generation()");

// Avoid reentrant initialization by taking the lock for the entire's function
// scope.
let mut gen_guard = generation.lock().await;
trace!(" generation={gen_guard:?}");

let prev_generation =
self.inner.store.get_custom_value(Self::CURRENT_GENERATION_STORE_KEY).await?;
Expand Down
2 changes: 1 addition & 1 deletion crates/matrix-sdk-crypto/src/olm/account.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1277,7 +1277,7 @@ impl Account {
// session in lieu of an OTK.)

warn!(
session_id = session.session_id(),
?session,
"Failed to decrypt a pre-key message with the corresponding session"
);

Expand Down
2 changes: 2 additions & 0 deletions crates/matrix-sdk/src/encryption/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1215,6 +1215,8 @@ impl Encryption {
///
/// The provided `lock_value` must be a unique identifier for this process.
pub async fn enable_cross_process_store_lock(&self, lock_value: String) -> Result<(), Error> {
trace!("enable_cross_process_store_lock({lock_value})");

// If the lock has already been created, don't recreate it from scratch.
if let Some(prev_lock) = self.client.locks().cross_process_crypto_store_lock.get() {
let prev_holder = prev_lock.lock_holder();
Expand Down
69 changes: 62 additions & 7 deletions testing/matrix-sdk-integration-testing/src/helpers.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ use std::{
collections::HashMap,
ops::Deref,
option_env,
path::{Path, PathBuf},
sync::{Arc, Mutex as StdMutex},
time::Duration,
};
Expand All @@ -21,9 +22,14 @@ use tokio::sync::Mutex;

static USERS: Lazy<Mutex<HashMap<String, (Client, TempDir)>>> = Lazy::new(Mutex::default);

enum SqlitePath {
Random,
Path(PathBuf),
}

pub struct TestClientBuilder {
username: String,
use_sqlite: bool,
use_sqlite_dir: Option<SqlitePath>,
encryption_settings: EncryptionSettings,
http_proxy: Option<String>,
}
Expand All @@ -32,7 +38,7 @@ impl TestClientBuilder {
pub fn new(username: impl Into<String>) -> Self {
Self {
username: username.into(),
use_sqlite: false,
use_sqlite_dir: None,
encryption_settings: Default::default(),
http_proxy: None,
}
Expand All @@ -45,7 +51,16 @@ impl TestClientBuilder {
}

pub fn use_sqlite(mut self) -> Self {
self.use_sqlite = true;
self.use_sqlite_dir = Some(SqlitePath::Random);
self
}

/// Create or re-use a Sqlite store (with no passphrase) in the supplied
/// directory. Note: this path must remain valid throughout the use of
/// the constructed Client, so if you created a TempDir you must hang on
/// to a reference to it throughout the test.
pub fn use_sqlite_dir(mut self, path: &Path) -> Self {
self.use_sqlite_dir = Some(SqlitePath::Path(path.to_owned()));
self
}

Expand All @@ -59,6 +74,42 @@ impl TestClientBuilder {
self
}

pub async fn duplicate(self, other: &Client) -> Result<Client> {
let homeserver_url =
option_env!("HOMESERVER_URL").unwrap_or("http://localhost:8228").to_owned();
let sliding_sync_proxy_url =
option_env!("SLIDING_SYNC_PROXY_URL").unwrap_or("http://localhost:8338").to_owned();

let mut client_builder = Client::builder()
.user_agent("matrix-sdk-integration-tests")
.homeserver_url(homeserver_url)
.sliding_sync_proxy(sliding_sync_proxy_url)
.with_encryption_settings(self.encryption_settings)
.request_config(RequestConfig::short_retry());

if let Some(proxy) = self.http_proxy {
client_builder = client_builder.proxy(proxy);
}

let client = match self.use_sqlite_dir {
Some(SqlitePath::Path(path_buf)) => {
client_builder.sqlite_store(&path_buf, None).build().await?
}
_ => {
panic!("You must call use_sqlite_dir for a duplicate client!");
}
};

client
.restore_session(
other.session().expect("Session must be logged in before we can duplicate it"),
)
.await
.expect("Failed to restore session");

Ok(client)
}

pub async fn build(self) -> Result<Client> {
let mut users = USERS.lock().await;
if let Some((client, _)) = users.get(&self.username) {
Expand All @@ -83,10 +134,14 @@ impl TestClientBuilder {
client_builder = client_builder.proxy(proxy);
}

let client = if self.use_sqlite {
client_builder.sqlite_store(tmp_dir.path(), None).build().await?
} else {
client_builder.build().await?
let client = match self.use_sqlite_dir {
None => client_builder.build().await?,
Some(SqlitePath::Random) => {
client_builder.sqlite_store(tmp_dir.path(), None).build().await?
}
Some(SqlitePath::Path(path_buf)) => {
client_builder.sqlite_store(&path_buf, None).build().await?
}
};

// safe to assume we have not registered this user yet, but ignore if we did
Expand Down

0 comments on commit a5517ff

Please sign in to comment.