matrix-rust-sdk/src/crypto/memory_stores.rs

162 lines
4.6 KiB
Rust
Raw Normal View History

// Copyright 2020 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use std::collections::HashMap;
use std::sync::Arc;
use dashmap::{DashMap, ReadOnlyView};
use tokio::sync::Mutex;
use super::device::Device;
use super::olm::{InboundGroupSession, Session};
use crate::identifiers::{DeviceId, RoomId, UserId};
#[derive(Debug)]
pub struct SessionStore {
2020-03-30 15:07:36 +00:00
entries: HashMap<String, Arc<Mutex<Vec<Arc<Mutex<Session>>>>>>,
}
impl SessionStore {
pub fn new() -> Self {
SessionStore {
entries: HashMap::new(),
}
}
pub async fn add(&mut self, session: Session) -> Arc<Mutex<Session>> {
if !self.entries.contains_key(&session.sender_key) {
2020-03-30 15:07:36 +00:00
self.entries.insert(
session.sender_key.to_owned(),
Arc::new(Mutex::new(Vec::new())),
);
}
let sessions = self.entries.get_mut(&session.sender_key).unwrap();
let session = Arc::new(Mutex::new(session));
sessions.lock().await.push(session.clone());
session
}
2020-03-30 15:07:36 +00:00
pub fn get(&self, sender_key: &str) -> Option<Arc<Mutex<Vec<Arc<Mutex<Session>>>>>> {
self.entries.get(sender_key).cloned()
}
pub fn set_for_sender(&mut self, sender_key: &str, sessions: Vec<Arc<Mutex<Session>>>) {
2020-03-30 15:07:36 +00:00
self.entries
.insert(sender_key.to_owned(), Arc::new(Mutex::new(sessions)));
}
}
#[derive(Debug)]
pub struct GroupSessionStore {
entries: HashMap<RoomId, HashMap<String, HashMap<String, InboundGroupSession>>>,
}
impl GroupSessionStore {
pub fn new() -> Self {
GroupSessionStore {
entries: HashMap::new(),
}
}
2020-03-25 10:32:40 +00:00
pub fn add(&mut self, session: InboundGroupSession) -> bool {
if !self.entries.contains_key(&session.room_id) {
let room_id = &*session.room_id;
self.entries.insert(room_id.clone(), HashMap::new());
}
let room_map = self.entries.get_mut(&session.room_id).unwrap();
if !room_map.contains_key(&*session.sender_key) {
let sender_key = &*session.sender_key;
room_map.insert(sender_key.to_owned(), HashMap::new());
}
let sender_map = room_map.get_mut(&*session.sender_key).unwrap();
let ret = sender_map.insert(session.session_id().to_owned(), session);
ret.is_some()
}
pub fn get(
&self,
2020-04-03 15:00:37 +00:00
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)]
pub struct DeviceStore {
2020-04-03 15:00:37 +00:00
entries: Arc<DashMap<UserId, DashMap<String, Device>>>,
}
pub struct UserDevices {
entries: ReadOnlyView<DeviceId, Device>,
}
impl UserDevices {
pub fn get(&self, device_id: &str) -> Option<Device> {
self.entries.get(device_id).cloned()
}
pub fn keys(&self) -> impl Iterator<Item = &DeviceId> {
self.entries.keys()
}
pub fn devices(&self) -> impl Iterator<Item = &Device> {
self.entries.values()
}
}
impl DeviceStore {
pub fn new() -> Self {
DeviceStore {
entries: Arc::new(DashMap::new()),
}
}
pub fn add(&self, device: Device) -> bool {
let user_id = device.user_id();
2020-04-03 15:00:37 +00:00
if !self.entries.contains_key(&user_id) {
self.entries.insert(user_id.clone(), DashMap::new());
}
2020-04-03 15:00:37 +00:00
let device_map = self.entries.get_mut(&user_id).unwrap();
device_map
.insert(device.device_id().to_owned(), device)
.is_some()
}
2020-04-03 15:00:37 +00:00
pub fn get(&self, user_id: &UserId, device_id: &str) -> Option<Device> {
self.entries
.get(user_id)
.and_then(|m| m.get(device_id).map(|d| d.value().clone()))
}
2020-04-03 15:00:37 +00:00
pub fn user_devices(&self, user_id: &UserId) -> UserDevices {
if !self.entries.contains_key(user_id) {
2020-04-03 15:00:37 +00:00
self.entries.insert(user_id.clone(), DashMap::new());
}
UserDevices {
entries: self.entries.get(user_id).unwrap().clone().into_read_only(),
}
}
}