2020-03-24 16:25:01 +00:00
|
|
|
// 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;
|
2020-03-26 10:22:40 +00:00
|
|
|
use std::sync::Arc;
|
|
|
|
|
2020-04-02 12:52:15 +00:00
|
|
|
use dashmap::{DashMap, ReadOnlyView};
|
2020-03-26 10:22:40 +00:00
|
|
|
use tokio::sync::Mutex;
|
|
|
|
|
2020-04-02 12:52:15 +00:00
|
|
|
use super::device::Device;
|
2020-03-26 10:22:40 +00:00
|
|
|
use super::olm::{InboundGroupSession, Session};
|
2020-04-09 14:22:25 +00:00
|
|
|
use crate::identifiers::{DeviceId, RoomId, UserId};
|
2020-03-26 10:22:40 +00:00
|
|
|
|
2020-04-15 13:32:58 +00:00
|
|
|
/// In-memory store for Olm Sessions.
|
2020-03-26 10:22:40 +00:00
|
|
|
#[derive(Debug)]
|
|
|
|
pub struct SessionStore {
|
2020-04-14 12:05:18 +00:00
|
|
|
entries: HashMap<String, Arc<Mutex<Vec<Session>>>>,
|
2020-03-26 10:22:40 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
impl SessionStore {
|
2020-04-15 13:32:58 +00:00
|
|
|
/// Create a new empty Session store.
|
2020-03-26 10:22:40 +00:00
|
|
|
pub fn new() -> Self {
|
|
|
|
SessionStore {
|
2020-03-27 11:09:54 +00:00
|
|
|
entries: HashMap::new(),
|
2020-03-26 10:22:40 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2020-04-15 13:32:58 +00:00
|
|
|
/// Add a session to the store.
|
2020-04-16 08:09:53 +00:00
|
|
|
///
|
2020-04-16 09:05:35 +00:00
|
|
|
/// Returns true if the the session was added, false if the session was
|
|
|
|
/// already in the store.
|
2020-04-16 08:09:53 +00:00
|
|
|
pub async fn add(&mut self, session: Session) -> bool {
|
2020-04-14 12:05:18 +00:00
|
|
|
if !self.entries.contains_key(&*session.sender_key) {
|
2020-03-30 15:07:36 +00:00
|
|
|
self.entries.insert(
|
2020-04-14 12:05:18 +00:00
|
|
|
session.sender_key.to_string(),
|
2020-03-30 15:07:36 +00:00
|
|
|
Arc::new(Mutex::new(Vec::new())),
|
|
|
|
);
|
2020-03-26 10:22:40 +00:00
|
|
|
}
|
2020-04-14 12:05:18 +00:00
|
|
|
let sessions = self.entries.get_mut(&*session.sender_key).unwrap();
|
2020-04-16 08:09:53 +00:00
|
|
|
|
|
|
|
if !sessions.lock().await.contains(&session) {
|
|
|
|
sessions.lock().await.push(session);
|
|
|
|
true
|
2020-04-16 09:05:35 +00:00
|
|
|
} else {
|
|
|
|
false
|
2020-04-16 08:09:53 +00:00
|
|
|
}
|
2020-03-27 11:09:54 +00:00
|
|
|
}
|
|
|
|
|
2020-04-15 13:32:58 +00:00
|
|
|
/// Get all the sessions that belong to the given sender key.
|
2020-04-14 12:05:18 +00:00
|
|
|
pub fn get(&self, sender_key: &str) -> Option<Arc<Mutex<Vec<Session>>>> {
|
2020-03-30 15:07:36 +00:00
|
|
|
self.entries.get(sender_key).cloned()
|
2020-03-26 10:22:40 +00:00
|
|
|
}
|
|
|
|
|
2020-04-15 13:32:58 +00:00
|
|
|
/// Add a list of sessions belonging to the sender key.
|
2020-04-14 12:05:18 +00:00
|
|
|
pub fn set_for_sender(&mut self, sender_key: &str, sessions: Vec<Session>) {
|
2020-03-30 15:07:36 +00:00
|
|
|
self.entries
|
|
|
|
.insert(sender_key.to_owned(), Arc::new(Mutex::new(sessions)));
|
2020-03-26 10:22:40 +00:00
|
|
|
}
|
|
|
|
}
|
2020-03-24 16:25:01 +00:00
|
|
|
|
|
|
|
#[derive(Debug)]
|
2020-04-15 13:32:58 +00:00
|
|
|
/// In-memory store that houlds inbound group sessions.
|
2020-03-24 16:25:01 +00:00
|
|
|
pub struct GroupSessionStore {
|
2020-04-10 14:08:47 +00:00
|
|
|
entries: HashMap<RoomId, HashMap<String, HashMap<String, InboundGroupSession>>>,
|
2020-03-24 16:25:01 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
impl GroupSessionStore {
|
2020-04-15 13:32:58 +00:00
|
|
|
/// Create a new empty store.
|
2020-03-24 16:25:01 +00:00
|
|
|
pub fn new() -> Self {
|
|
|
|
GroupSessionStore {
|
|
|
|
entries: HashMap::new(),
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2020-04-15 13:32:58 +00:00
|
|
|
/// Add a inbound group session to the store.
|
2020-04-16 08:09:53 +00:00
|
|
|
///
|
2020-04-16 09:05:35 +00:00
|
|
|
/// Returns true if the the session was added, false if the session was
|
|
|
|
/// already in the store.
|
2020-03-25 10:32:40 +00:00
|
|
|
pub fn add(&mut self, session: InboundGroupSession) -> bool {
|
|
|
|
if !self.entries.contains_key(&session.room_id) {
|
2020-04-10 14:08:47 +00:00
|
|
|
let room_id = &*session.room_id;
|
|
|
|
self.entries.insert(room_id.clone(), HashMap::new());
|
2020-03-24 16:25:01 +00:00
|
|
|
}
|
|
|
|
|
2020-04-03 10:34:05 +00:00
|
|
|
let room_map = self.entries.get_mut(&session.room_id).unwrap();
|
2020-03-24 16:25:01 +00:00
|
|
|
|
2020-04-10 14:08:47 +00:00
|
|
|
if !room_map.contains_key(&*session.sender_key) {
|
|
|
|
let sender_key = &*session.sender_key;
|
|
|
|
room_map.insert(sender_key.to_owned(), HashMap::new());
|
2020-03-24 16:25:01 +00:00
|
|
|
}
|
|
|
|
|
2020-04-10 14:08:47 +00:00
|
|
|
let sender_map = room_map.get_mut(&*session.sender_key).unwrap();
|
|
|
|
let ret = sender_map.insert(session.session_id().to_owned(), session);
|
2020-03-24 16:25:01 +00:00
|
|
|
|
2020-04-16 09:05:35 +00:00
|
|
|
ret.is_none()
|
2020-03-24 16:25:01 +00:00
|
|
|
}
|
|
|
|
|
2020-04-15 13:32:58 +00:00
|
|
|
/// Get a inbound group session from our store.
|
|
|
|
///
|
|
|
|
/// # Arguments
|
|
|
|
/// * `room_id` - The room id of the room that the session belongs to.
|
|
|
|
///
|
|
|
|
/// * `sender_key` - The sender key that sent us the session.
|
|
|
|
///
|
|
|
|
/// * `session_id` - The unique id of the session.
|
2020-03-24 16:25:01 +00:00
|
|
|
pub fn get(
|
|
|
|
&self,
|
2020-04-03 15:00:37 +00:00
|
|
|
room_id: &RoomId,
|
2020-03-24 16:25:01 +00:00
|
|
|
sender_key: &str,
|
|
|
|
session_id: &str,
|
2020-04-10 14:08:47 +00:00
|
|
|
) -> Option<InboundGroupSession> {
|
2020-03-24 16:25:01 +00:00
|
|
|
self.entries
|
|
|
|
.get(room_id)
|
2020-03-26 10:22:40 +00:00
|
|
|
.and_then(|m| m.get(sender_key).and_then(|m| m.get(session_id).cloned()))
|
2020-03-24 16:25:01 +00:00
|
|
|
}
|
|
|
|
}
|
2020-04-02 12:52:15 +00:00
|
|
|
|
2020-04-15 15:51:51 +00:00
|
|
|
/// In-memory store holding the devices of users.
|
2020-04-02 15:27:39 +00:00
|
|
|
#[derive(Clone, Debug)]
|
2020-04-02 12:52:15 +00:00
|
|
|
pub struct DeviceStore {
|
2020-04-03 15:00:37 +00:00
|
|
|
entries: Arc<DashMap<UserId, DashMap<String, Device>>>,
|
2020-04-02 12:52:15 +00:00
|
|
|
}
|
|
|
|
|
2020-04-15 15:51:51 +00:00
|
|
|
/// A read only view over all devices belonging to a user.
|
2020-04-02 12:52:15 +00:00
|
|
|
pub struct UserDevices {
|
2020-04-09 14:22:25 +00:00
|
|
|
entries: ReadOnlyView<DeviceId, Device>,
|
2020-04-02 12:52:15 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
impl UserDevices {
|
2020-04-15 15:51:51 +00:00
|
|
|
/// Get the specific device with the given device id.
|
2020-04-02 12:52:15 +00:00
|
|
|
pub fn get(&self, device_id: &str) -> Option<Device> {
|
|
|
|
self.entries.get(device_id).cloned()
|
|
|
|
}
|
|
|
|
|
2020-04-15 15:51:51 +00:00
|
|
|
/// Iterator over all the device ids of the user devices.
|
2020-04-09 14:22:25 +00:00
|
|
|
pub fn keys(&self) -> impl Iterator<Item = &DeviceId> {
|
2020-04-02 12:52:15 +00:00
|
|
|
self.entries.keys()
|
|
|
|
}
|
2020-04-03 08:27:30 +00:00
|
|
|
|
2020-04-15 15:51:51 +00:00
|
|
|
/// Iterator over all the devices of the user devices.
|
2020-04-03 08:27:30 +00:00
|
|
|
pub fn devices(&self) -> impl Iterator<Item = &Device> {
|
|
|
|
self.entries.values()
|
|
|
|
}
|
2020-04-02 12:52:15 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
impl DeviceStore {
|
2020-04-15 15:51:51 +00:00
|
|
|
/// Create a new empty device store.
|
2020-04-02 12:52:15 +00:00
|
|
|
pub fn new() -> Self {
|
|
|
|
DeviceStore {
|
2020-04-02 15:27:39 +00:00
|
|
|
entries: Arc::new(DashMap::new()),
|
2020-04-02 12:52:15 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2020-04-15 15:51:51 +00:00
|
|
|
/// Add a device to the store.
|
|
|
|
///
|
|
|
|
/// Returns true if the device was already in the store, false otherwise.
|
2020-04-02 12:52:15 +00:00
|
|
|
pub fn add(&self, device: Device) -> bool {
|
2020-04-09 14:22:25 +00:00
|
|
|
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-02 12:52:15 +00:00
|
|
|
}
|
2020-04-03 15:00:37 +00:00
|
|
|
let device_map = self.entries.get_mut(&user_id).unwrap();
|
2020-04-02 12:52:15 +00:00
|
|
|
|
|
|
|
device_map
|
|
|
|
.insert(device.device_id().to_owned(), device)
|
2020-04-16 09:05:35 +00:00
|
|
|
.is_none()
|
2020-04-02 12:52:15 +00:00
|
|
|
}
|
|
|
|
|
2020-04-15 15:51:51 +00:00
|
|
|
/// Get the device with the given device_id and belonging to the given user.
|
2020-04-03 15:00:37 +00:00
|
|
|
pub fn get(&self, user_id: &UserId, device_id: &str) -> Option<Device> {
|
2020-04-02 12:52:15 +00:00
|
|
|
self.entries
|
|
|
|
.get(user_id)
|
|
|
|
.and_then(|m| m.get(device_id).map(|d| d.value().clone()))
|
|
|
|
}
|
|
|
|
|
2020-04-15 15:51:51 +00:00
|
|
|
/// Get a read-only view over all devices of the given user.
|
2020-04-03 15:00:37 +00:00
|
|
|
pub fn user_devices(&self, user_id: &UserId) -> UserDevices {
|
2020-04-02 12:52:15 +00:00
|
|
|
if !self.entries.contains_key(user_id) {
|
2020-04-03 15:00:37 +00:00
|
|
|
self.entries.insert(user_id.clone(), DashMap::new());
|
2020-04-02 12:52:15 +00:00
|
|
|
}
|
|
|
|
UserDevices {
|
|
|
|
entries: self.entries.get(user_id).unwrap().clone().into_read_only(),
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
2020-04-15 13:32:58 +00:00
|
|
|
|
|
|
|
#[cfg(test)]
|
|
|
|
mod test {
|
|
|
|
use std::collections::HashMap;
|
|
|
|
use std::convert::TryFrom;
|
|
|
|
|
|
|
|
use crate::api::r0::keys::SignedKey;
|
2020-04-15 15:51:51 +00:00
|
|
|
use crate::crypto::device::test::get_device;
|
2020-04-15 13:32:58 +00:00
|
|
|
use crate::crypto::memory_stores::{DeviceStore, GroupSessionStore, SessionStore};
|
|
|
|
use crate::crypto::olm::{Account, InboundGroupSession, OutboundGroupSession, Session};
|
|
|
|
use crate::identifiers::RoomId;
|
|
|
|
|
|
|
|
async fn get_account_and_session() -> (Account, Session) {
|
|
|
|
let alice = Account::new();
|
|
|
|
|
|
|
|
let bob = Account::new();
|
|
|
|
|
|
|
|
bob.generate_one_time_keys(1).await;
|
|
|
|
let one_time_key = bob
|
|
|
|
.one_time_keys()
|
|
|
|
.await
|
|
|
|
.curve25519()
|
|
|
|
.iter()
|
|
|
|
.nth(0)
|
|
|
|
.unwrap()
|
|
|
|
.1
|
|
|
|
.to_owned();
|
|
|
|
let one_time_key = SignedKey {
|
|
|
|
key: one_time_key,
|
|
|
|
signatures: HashMap::new(),
|
|
|
|
};
|
|
|
|
let sender_key = bob.identity_keys().curve25519().to_owned();
|
|
|
|
let session = alice
|
|
|
|
.create_outbound_session(&sender_key, &one_time_key)
|
|
|
|
.await
|
|
|
|
.unwrap();
|
|
|
|
|
|
|
|
(alice, session)
|
|
|
|
}
|
|
|
|
|
|
|
|
#[tokio::test]
|
|
|
|
async fn test_session_store() {
|
|
|
|
let (account, session) = get_account_and_session().await;
|
|
|
|
|
|
|
|
let mut store = SessionStore::new();
|
2020-04-16 08:09:53 +00:00
|
|
|
|
|
|
|
assert!(store.add(session.clone()).await);
|
2020-04-16 09:05:35 +00:00
|
|
|
assert!(!store.add(session.clone()).await);
|
2020-04-15 13:32:58 +00:00
|
|
|
|
|
|
|
let sessions = store.get(&session.sender_key).unwrap();
|
|
|
|
let sessions = sessions.lock().await;
|
|
|
|
|
|
|
|
let loaded_session = &sessions[0];
|
|
|
|
|
|
|
|
assert_eq!(&session, loaded_session);
|
|
|
|
}
|
|
|
|
|
2020-04-15 15:51:51 +00:00
|
|
|
#[tokio::test]
|
|
|
|
async fn test_session_store_bulk_storing() {
|
|
|
|
let (account, session) = get_account_and_session().await;
|
|
|
|
|
|
|
|
let mut store = SessionStore::new();
|
|
|
|
store.set_for_sender(&session.sender_key, vec![session.clone()]);
|
|
|
|
|
|
|
|
let sessions = store.get(&session.sender_key).unwrap();
|
|
|
|
let sessions = sessions.lock().await;
|
|
|
|
|
|
|
|
let loaded_session = &sessions[0];
|
|
|
|
|
|
|
|
assert_eq!(&session, loaded_session);
|
|
|
|
}
|
|
|
|
|
2020-04-15 13:32:58 +00:00
|
|
|
#[tokio::test]
|
|
|
|
async fn test_group_session_store() {
|
|
|
|
let alice = Account::new();
|
|
|
|
let room_id = RoomId::try_from("!test:localhost").unwrap();
|
|
|
|
|
|
|
|
let outbound = OutboundGroupSession::new(&room_id);
|
|
|
|
|
|
|
|
assert_eq!(0, outbound.message_index().await);
|
|
|
|
assert!(!outbound.shared());
|
|
|
|
outbound.mark_as_shared();
|
|
|
|
assert!(outbound.shared());
|
|
|
|
|
|
|
|
let inbound = InboundGroupSession::new(
|
|
|
|
"test_key",
|
|
|
|
"test_key",
|
|
|
|
&room_id,
|
|
|
|
outbound.session_key().await,
|
|
|
|
)
|
|
|
|
.unwrap();
|
|
|
|
|
|
|
|
let mut store = GroupSessionStore::new();
|
|
|
|
store.add(inbound.clone());
|
|
|
|
|
|
|
|
let loaded_session = store
|
|
|
|
.get(&room_id, "test_key", outbound.session_id())
|
|
|
|
.unwrap();
|
|
|
|
assert_eq!(inbound, loaded_session);
|
|
|
|
}
|
2020-04-15 15:51:51 +00:00
|
|
|
|
|
|
|
#[tokio::test]
|
|
|
|
async fn test_device_store() {
|
|
|
|
let device = get_device();
|
|
|
|
let store = DeviceStore::new();
|
|
|
|
|
|
|
|
assert!(store.add(device.clone()));
|
2020-04-16 09:05:35 +00:00
|
|
|
assert!(!store.add(device.clone()));
|
2020-04-15 15:51:51 +00:00
|
|
|
|
|
|
|
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().nth(0).unwrap(), device.device_id());
|
|
|
|
assert_eq!(user_devices.devices().nth(0).unwrap(), &device);
|
|
|
|
|
|
|
|
let loaded_device = user_devices.get(device.device_id()).unwrap();
|
|
|
|
|
|
|
|
assert_eq!(device, loaded_device);
|
|
|
|
}
|
2020-04-15 13:32:58 +00:00
|
|
|
}
|