2020-11-06 21:52:49 +00:00
|
|
|
use crate::db::errors::DataError;
|
2020-11-07 14:37:56 +00:00
|
|
|
use sled::transaction::TransactionalTree;
|
2020-11-06 21:52:49 +00:00
|
|
|
use sled::Transactional;
|
|
|
|
use sled::Tree;
|
|
|
|
use std::collections::HashSet;
|
|
|
|
use std::str;
|
|
|
|
|
|
|
|
#[derive(Clone)]
|
|
|
|
pub struct Rooms {
|
|
|
|
/// Room ID -> RoomInfo struct (single entries)
|
|
|
|
pub(in crate::db) roomid_roominfo: Tree,
|
|
|
|
|
|
|
|
/// Room ID -> list of usernames in room.
|
|
|
|
pub(in crate::db) roomid_usernames: Tree,
|
|
|
|
|
|
|
|
/// Username -> list of room IDs user is in.
|
|
|
|
pub(in crate::db) username_roomids: Tree,
|
|
|
|
}
|
|
|
|
|
2020-11-07 20:47:49 +00:00
|
|
|
#[derive(Clone, Copy)]
|
2020-11-07 14:37:56 +00:00
|
|
|
enum TxableTree<'a> {
|
|
|
|
Tree(&'a Tree),
|
|
|
|
Tx(&'a TransactionalTree),
|
|
|
|
}
|
|
|
|
|
|
|
|
impl<'a> Into<TxableTree<'a>> for &'a Tree {
|
|
|
|
fn into(self) -> TxableTree<'a> {
|
|
|
|
TxableTree::Tree(self)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
impl<'a> Into<TxableTree<'a>> for &'a TransactionalTree {
|
|
|
|
fn into(self) -> TxableTree<'a> {
|
|
|
|
TxableTree::Tx(self)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
fn get_set<'a, T: Into<TxableTree<'a>>>(tree: T, key: &[u8]) -> Result<HashSet<String>, DataError> {
|
|
|
|
let set: HashSet<String> = match tree.into() {
|
|
|
|
TxableTree::Tree(tree) => tree.get(key)?,
|
|
|
|
TxableTree::Tx(tx) => tx.get(key)?,
|
|
|
|
}
|
|
|
|
.map(|bytes| bincode::deserialize::<HashSet<String>>(&bytes))
|
|
|
|
.unwrap_or(Ok(HashSet::new()))?;
|
|
|
|
|
|
|
|
Ok(set)
|
|
|
|
}
|
|
|
|
|
2020-11-07 20:47:49 +00:00
|
|
|
fn remove_from_set<'a, T: Into<TxableTree<'a>> + Copy>(
|
|
|
|
tree: T,
|
|
|
|
key: &[u8],
|
|
|
|
value_to_remove: &str,
|
|
|
|
) -> Result<(), DataError> {
|
|
|
|
let mut set = get_set(tree, key)?;
|
|
|
|
set.remove(value_to_remove);
|
|
|
|
insert_set(tree, key, set)?;
|
|
|
|
Ok(())
|
|
|
|
}
|
|
|
|
|
2020-11-07 14:37:56 +00:00
|
|
|
fn insert_set<'a, T: Into<TxableTree<'a>>>(
|
|
|
|
tree: T,
|
|
|
|
key: &[u8],
|
|
|
|
set: HashSet<String>,
|
|
|
|
) -> Result<(), DataError> {
|
|
|
|
let serialized = bincode::serialize(&set)?;
|
|
|
|
match tree.into() {
|
|
|
|
TxableTree::Tree(tree) => tree.insert(key, serialized)?,
|
|
|
|
TxableTree::Tx(tx) => tx.insert(key, serialized)?,
|
|
|
|
};
|
|
|
|
Ok(())
|
|
|
|
}
|
2020-11-06 21:52:49 +00:00
|
|
|
|
|
|
|
impl Rooms {
|
|
|
|
pub(in crate::db) fn new(db: &sled::Db) -> Result<Rooms, sled::Error> {
|
|
|
|
Ok(Rooms {
|
|
|
|
roomid_roominfo: db.open_tree("roomid_roominfo")?,
|
|
|
|
roomid_usernames: db.open_tree("roomid_usernames")?,
|
|
|
|
username_roomids: db.open_tree("username_roomids")?,
|
|
|
|
})
|
|
|
|
}
|
|
|
|
|
2020-11-07 14:37:56 +00:00
|
|
|
pub fn get_rooms_for_user(&self, username: &str) -> Result<HashSet<String>, DataError> {
|
|
|
|
get_set(&self.username_roomids, username.as_bytes())
|
|
|
|
}
|
|
|
|
|
|
|
|
pub fn get_users_in_room(&self, room_id: &str) -> Result<HashSet<String>, DataError> {
|
|
|
|
get_set(&self.roomid_usernames, room_id.as_bytes())
|
|
|
|
}
|
|
|
|
|
2020-11-06 21:52:49 +00:00
|
|
|
pub fn add_user_to_room(&self, username: &str, room_id: &str) -> Result<(), DataError> {
|
2020-11-07 14:37:56 +00:00
|
|
|
(&self.username_roomids, &self.roomid_usernames).transaction(
|
|
|
|
|(tx_username_rooms, tx_room_usernames)| {
|
|
|
|
let username_key = &username.as_bytes();
|
|
|
|
let mut user_to_rooms = get_set(tx_username_rooms, username_key)?;
|
|
|
|
user_to_rooms.insert(room_id.to_string());
|
|
|
|
insert_set(tx_username_rooms, username_key, user_to_rooms)?;
|
|
|
|
|
|
|
|
let roomid_key = &room_id.as_bytes();
|
|
|
|
let mut room_to_users = get_set(tx_room_usernames, roomid_key)?;
|
|
|
|
room_to_users.insert(username.to_string());
|
|
|
|
insert_set(tx_room_usernames, roomid_key, room_to_users)?;
|
|
|
|
|
|
|
|
Ok(())
|
|
|
|
},
|
|
|
|
)?;
|
|
|
|
|
2020-11-06 21:52:49 +00:00
|
|
|
Ok(())
|
|
|
|
}
|
|
|
|
|
|
|
|
pub fn remove_user_from_room(&self, username: &str, room_id: &str) -> Result<(), DataError> {
|
2020-11-07 14:37:56 +00:00
|
|
|
(&self.username_roomids, &self.roomid_usernames).transaction(
|
|
|
|
|(tx_username_rooms, tx_room_usernames)| {
|
|
|
|
let username_key = &username.as_bytes();
|
|
|
|
let mut user_to_rooms = get_set(tx_username_rooms, username_key)?;
|
|
|
|
user_to_rooms.remove(room_id);
|
|
|
|
insert_set(tx_username_rooms, username_key, user_to_rooms)?;
|
|
|
|
|
|
|
|
let roomid_key = &room_id.as_bytes();
|
|
|
|
let mut room_to_users = get_set(tx_room_usernames, roomid_key)?;
|
|
|
|
room_to_users.remove(username);
|
|
|
|
insert_set(tx_room_usernames, roomid_key, room_to_users)?;
|
|
|
|
|
|
|
|
Ok(())
|
|
|
|
},
|
|
|
|
)?;
|
|
|
|
|
|
|
|
Ok(())
|
|
|
|
}
|
|
|
|
|
2020-11-07 20:47:49 +00:00
|
|
|
pub fn clear_info(&self, room_id: &str) -> Result<(), DataError> {
|
|
|
|
(&self.username_roomids, &self.roomid_usernames).transaction(
|
|
|
|
|(tx_username_roomids, tx_roomid_usernames)| {
|
|
|
|
let roomid_key = room_id.as_bytes();
|
|
|
|
let users_in_room = get_set(tx_roomid_usernames, roomid_key)?;
|
|
|
|
|
|
|
|
//Remove the room ID from every user's room ID list.
|
|
|
|
for username in users_in_room {
|
|
|
|
remove_from_set(tx_username_roomids, username.as_bytes(), room_id)?;
|
|
|
|
}
|
|
|
|
|
|
|
|
//Remove this room entry for the room ID -> username tree.
|
|
|
|
tx_roomid_usernames.remove(roomid_key)?;
|
|
|
|
|
|
|
|
//TODO: delete roominfo struct from room info tree.
|
|
|
|
Ok(())
|
|
|
|
},
|
|
|
|
)?;
|
|
|
|
|
2020-11-06 21:52:49 +00:00
|
|
|
Ok(())
|
|
|
|
}
|
|
|
|
}
|
2020-11-07 14:37:56 +00:00
|
|
|
|
|
|
|
#[cfg(test)]
|
|
|
|
mod tests {
|
|
|
|
use super::*;
|
|
|
|
use sled::Config;
|
|
|
|
|
|
|
|
fn create_test_instance() -> Rooms {
|
|
|
|
let config = Config::new().temporary(true);
|
|
|
|
let db = config.open().unwrap();
|
|
|
|
Rooms::new(&db).unwrap()
|
|
|
|
}
|
|
|
|
|
|
|
|
#[test]
|
|
|
|
fn add_user_to_room() {
|
|
|
|
let rooms = create_test_instance();
|
|
|
|
rooms
|
|
|
|
.add_user_to_room("testuser", "myroom")
|
|
|
|
.expect("Could not add user to room");
|
|
|
|
|
|
|
|
let users_in_room = rooms
|
|
|
|
.get_users_in_room("myroom")
|
|
|
|
.expect("Could not retrieve users in room");
|
|
|
|
|
|
|
|
let rooms_for_user = rooms
|
|
|
|
.get_rooms_for_user("testuser")
|
|
|
|
.expect("Could not get rooms for user");
|
|
|
|
|
|
|
|
let expected_users_in_room: HashSet<String> =
|
|
|
|
vec![String::from("testuser")].into_iter().collect();
|
|
|
|
|
|
|
|
let expected_rooms_for_user: HashSet<String> =
|
|
|
|
vec![String::from("myroom")].into_iter().collect();
|
|
|
|
|
|
|
|
assert_eq!(expected_users_in_room, users_in_room);
|
|
|
|
assert_eq!(expected_rooms_for_user, rooms_for_user);
|
|
|
|
}
|
|
|
|
|
|
|
|
#[test]
|
|
|
|
fn remove_user_from_room() {
|
|
|
|
let rooms = create_test_instance();
|
|
|
|
rooms
|
|
|
|
.add_user_to_room("testuser", "myroom")
|
|
|
|
.expect("Could not add user to room");
|
|
|
|
|
|
|
|
rooms
|
|
|
|
.remove_user_from_room("testuser", "myroom")
|
|
|
|
.expect("Could not remove user from room");
|
|
|
|
|
|
|
|
let users_in_room = rooms
|
|
|
|
.get_users_in_room("myroom")
|
|
|
|
.expect("Could not retrieve users in room");
|
|
|
|
|
|
|
|
let rooms_for_user = rooms
|
|
|
|
.get_rooms_for_user("testuser")
|
|
|
|
.expect("Could not get rooms for user");
|
|
|
|
|
|
|
|
assert_eq!(HashSet::new(), users_in_room);
|
|
|
|
assert_eq!(HashSet::new(), rooms_for_user);
|
|
|
|
}
|
2020-11-07 20:47:49 +00:00
|
|
|
|
|
|
|
#[test]
|
|
|
|
fn clear_info() {
|
|
|
|
let rooms = create_test_instance();
|
|
|
|
rooms
|
|
|
|
.add_user_to_room("testuser", "myroom1")
|
|
|
|
.expect("Could not add user to room1");
|
|
|
|
|
|
|
|
rooms
|
|
|
|
.add_user_to_room("testuser", "myroom2")
|
|
|
|
.expect("Could not add user to room2");
|
|
|
|
|
|
|
|
rooms
|
|
|
|
.clear_info("myroom1")
|
|
|
|
.expect("Could not clear room info");
|
|
|
|
|
|
|
|
let users_in_room1 = rooms
|
|
|
|
.get_users_in_room("myroom1")
|
|
|
|
.expect("Could not retrieve users in room1");
|
|
|
|
|
|
|
|
let users_in_room2 = rooms
|
|
|
|
.get_users_in_room("myroom2")
|
|
|
|
.expect("Could not retrieve users in room2");
|
|
|
|
|
|
|
|
let rooms_for_user = rooms
|
|
|
|
.get_rooms_for_user("testuser")
|
|
|
|
.expect("Could not get rooms for user");
|
|
|
|
|
|
|
|
let expected_users_in_room2: HashSet<String> =
|
|
|
|
vec![String::from("testuser")].into_iter().collect();
|
|
|
|
|
|
|
|
let expected_rooms_for_user: HashSet<String> =
|
|
|
|
vec![String::from("myroom2")].into_iter().collect();
|
|
|
|
|
|
|
|
assert_eq!(HashSet::new(), users_in_room1);
|
|
|
|
assert_eq!(expected_users_in_room2, users_in_room2);
|
|
|
|
assert_eq!(expected_rooms_for_user, rooms_for_user);
|
|
|
|
}
|
2020-11-07 14:37:56 +00:00
|
|
|
}
|