use std::{collections::HashMap, sync::Arc};
use dashmap::DashMap;
use matrix_sdk_common::locks::Mutex;
use ruma::{DeviceId, OwnedDeviceId, OwnedRoomId, OwnedUserId, RoomId, UserId};
use crate::{
identities::ReadOnlyDevice,
olm::{InboundGroupSession, Session},
};
#[derive(Debug, Default, Clone)]
pub struct SessionStore {
entries: Arc<DashMap<String, Arc<Mutex<Vec<Session>>>>>,
}
impl SessionStore {
pub fn new() -> Self {
Self::default()
}
pub async fn add(&self, session: Session) -> bool {
let sessions_lock = self.entries.entry(session.sender_key.to_base64()).or_default();
let mut sessions = sessions_lock.lock().await;
if !sessions.contains(&session) {
sessions.push(session);
true
} else {
false
}
}
pub fn get(&self, sender_key: &str) -> Option<Arc<Mutex<Vec<Session>>>> {
self.entries.get(sender_key).map(|s| s.clone())
}
pub fn set_for_sender(&self, sender_key: &str, sessions: Vec<Session>) {
self.entries.insert(sender_key.to_owned(), Arc::new(Mutex::new(sessions)));
}
}
#[derive(Debug, Default, Clone)]
pub struct GroupSessionStore {
#[allow(clippy::type_complexity)]
entries: Arc<DashMap<OwnedRoomId, HashMap<String, HashMap<String, InboundGroupSession>>>>,
}
impl GroupSessionStore {
pub fn new() -> Self {
Self::default()
}
pub fn add(&self, session: InboundGroupSession) -> bool {
self.entries
.entry((*session.room_id).to_owned())
.or_default()
.entry(session.sender_key.to_base64())
.or_default()
.insert(session.session_id().to_owned(), session)
.is_none()
}
pub fn get_all(&self) -> Vec<InboundGroupSession> {
self.entries
.iter()
.flat_map(|d| {
d.value()
.values()
.flat_map(|t| t.values().cloned().collect::<Vec<InboundGroupSession>>())
.collect::<Vec<InboundGroupSession>>()
})
.collect()
}
pub fn count(&self) -> usize {
self.entries.iter().map(|d| d.value().values().map(|t| t.len()).sum::<usize>()).sum()
}
pub fn get(
&self,
room_id: &RoomId,
sender_key: &str,
session_id: &str,
) -> Option<InboundGroupSession> {
self.entries
.get(room_id)
.and_then(|m| m.get(sender_key).and_then(|m| m.get(session_id).cloned()))
}
}
#[derive(Clone, Debug, Default)]
pub struct DeviceStore {
entries: Arc<DashMap<OwnedUserId, DashMap<OwnedDeviceId, ReadOnlyDevice>>>,
}
impl DeviceStore {
pub fn new() -> Self {
Self::default()
}
pub fn add(&self, device: ReadOnlyDevice) -> bool {
let user_id = device.user_id();
self.entries
.entry(user_id.to_owned())
.or_default()
.insert(device.device_id().into(), device)
.is_none()
}
pub fn get(&self, user_id: &UserId, device_id: &DeviceId) -> Option<ReadOnlyDevice> {
self.entries.get(user_id).and_then(|m| m.get(device_id).map(|d| d.value().clone()))
}
pub fn remove(&self, user_id: &UserId, device_id: &DeviceId) -> Option<ReadOnlyDevice> {
self.entries.get(user_id).and_then(|m| m.remove(device_id)).map(|(_, d)| d)
}
pub fn user_devices(&self, user_id: &UserId) -> HashMap<OwnedDeviceId, ReadOnlyDevice> {
self.entries
.entry(user_id.to_owned())
.or_default()
.iter()
.map(|i| (i.key().to_owned(), i.value().clone()))
.collect()
}
}
#[cfg(test)]
mod tests {
use matrix_sdk_test::async_test;
use ruma::room_id;
use vodozemac::{Curve25519PublicKey, Ed25519PublicKey};
use crate::{
identities::device::testing::get_device,
olm::{tests::get_account_and_session, InboundGroupSession},
store::caches::{DeviceStore, GroupSessionStore, SessionStore},
};
#[async_test]
async fn test_session_store() {
let (_, session) = get_account_and_session().await;
let store = SessionStore::new();
assert!(store.add(session.clone()).await);
assert!(!store.add(session.clone()).await);
let sessions = store.get(&session.sender_key.to_base64()).unwrap();
let sessions = sessions.lock().await;
let loaded_session = &sessions[0];
assert_eq!(&session, loaded_session);
}
#[async_test]
async fn test_session_store_bulk_storing() {
let (_, session) = get_account_and_session().await;
let store = SessionStore::new();
store.set_for_sender(&session.sender_key.to_base64(), vec![session.clone()]);
let sessions = store.get(&session.sender_key.to_base64()).unwrap();
let sessions = sessions.lock().await;
let loaded_session = &sessions[0];
assert_eq!(&session, loaded_session);
}
#[async_test]
async fn test_group_session_store() {
let (account, _) = get_account_and_session().await;
let room_id = room_id!("!test:localhost");
let curve_key = "Nn0L2hkcCMFKqynTjyGsJbth7QrVmX3lbrksMkrGOAw";
let (outbound, _) = account.create_group_session_pair_with_defaults(room_id).await;
assert_eq!(0, outbound.message_index().await);
assert!(!outbound.shared());
outbound.mark_as_shared();
assert!(outbound.shared());
let inbound = InboundGroupSession::new(
Curve25519PublicKey::from_base64(curve_key).unwrap(),
Ed25519PublicKey::from_base64("ee3Ek+J2LkkPmjGPGLhMxiKnhiX//xcqaVL4RP6EypE").unwrap(),
room_id,
&outbound.session_key().await,
outbound.settings().algorithm.to_owned(),
None,
)
.unwrap();
let store = GroupSessionStore::new();
store.add(inbound.clone());
let loaded_session = store.get(room_id, curve_key, outbound.session_id()).unwrap();
assert_eq!(inbound, loaded_session);
}
#[async_test]
async fn test_device_store() {
let device = get_device();
let store = DeviceStore::new();
assert!(store.add(device.clone()));
assert!(!store.add(device.clone()));
let loaded_device = store.get(device.user_id(), device.device_id()).unwrap();
assert_eq!(device, loaded_device);
let user_devices = store.user_devices(device.user_id());
assert_eq!(&**user_devices.keys().next().unwrap(), device.device_id());
assert_eq!(user_devices.values().next().unwrap(), &device);
let loaded_device = user_devices.get(device.device_id()).unwrap();
assert_eq!(&device, loaded_device);
store.remove(device.user_id(), device.device_id());
let loaded_device = store.get(device.user_id(), device.device_id());
assert!(loaded_device.is_none());
}
}