timokoesters
6 years ago
5 changed files with 206 additions and 101 deletions
@ -1,134 +1,115 @@
@@ -1,134 +1,115 @@
|
||||
use crate::utils; |
||||
use directories::ProjectDirs; |
||||
use log::debug; |
||||
use crate::{utils, Database}; |
||||
use ruma_events::collections::all::Event; |
||||
use ruma_identifiers::{EventId, RoomId, UserId}; |
||||
use std::convert::TryInto; |
||||
|
||||
const USERID_PASSWORD: &str = "userid_password"; |
||||
const USERID_DEVICEIDS: &str = "userid_deviceids"; |
||||
const DEVICEID_TOKEN: &str = "deviceid_token"; |
||||
const TOKEN_USERID: &str = "token_userid"; |
||||
|
||||
pub struct Data(sled::Db); |
||||
pub struct Data { |
||||
hostname: String, |
||||
db: Database, |
||||
} |
||||
|
||||
impl Data { |
||||
/// Load an existing database or create a new one.
|
||||
pub fn load_or_create() -> Self { |
||||
Data( |
||||
sled::open( |
||||
ProjectDirs::from("xyz", "koesters", "matrixserver") |
||||
.unwrap() |
||||
.data_dir(), |
||||
) |
||||
.unwrap(), |
||||
) |
||||
} |
||||
|
||||
/// Set the hostname of the server. Warning: Hostname changes will likely break things.
|
||||
pub fn set_hostname(&self, hostname: &str) { |
||||
self.0.insert("hostname", hostname).unwrap(); |
||||
pub fn load_or_create(hostname: &str) -> Self { |
||||
Self { |
||||
hostname: hostname.to_owned(), |
||||
db: Database::load_or_create(hostname), |
||||
} |
||||
} |
||||
|
||||
/// Get the hostname of the server.
|
||||
pub fn hostname(&self) -> String { |
||||
utils::bytes_to_string(&self.0.get("hostname").unwrap().unwrap()) |
||||
pub fn hostname(&self) -> &str { |
||||
&self.hostname |
||||
} |
||||
|
||||
/// Check if a user has an account by looking for an assigned password.
|
||||
pub fn user_exists(&self, user_id: &UserId) -> bool { |
||||
self.0 |
||||
.open_tree(USERID_PASSWORD) |
||||
.unwrap() |
||||
self.db |
||||
.userid_password |
||||
.contains_key(user_id.to_string()) |
||||
.unwrap() |
||||
} |
||||
|
||||
/// Create a new user account by assigning them a password.
|
||||
pub fn user_add(&self, user_id: &UserId, password: Option<String>) { |
||||
self.0 |
||||
.open_tree(USERID_PASSWORD) |
||||
.unwrap() |
||||
self.db |
||||
.userid_password |
||||
.insert(user_id.to_string(), &*password.unwrap_or_default()) |
||||
.unwrap(); |
||||
} |
||||
|
||||
/// Find out which user an access token belongs to.
|
||||
pub fn user_from_token(&self, token: &str) -> Option<UserId> { |
||||
self.0 |
||||
.open_tree(TOKEN_USERID) |
||||
.unwrap() |
||||
self.db |
||||
.token_userid |
||||
.get(token) |
||||
.unwrap() |
||||
.and_then(|bytes| (*utils::bytes_to_string(&bytes)).try_into().ok()) |
||||
.and_then(|bytes| (*utils::string_from_bytes(&bytes)).try_into().ok()) |
||||
} |
||||
|
||||
/// Checks if the given password is equal to the one in the database.
|
||||
pub fn password_get(&self, user_id: &UserId) -> Option<String> { |
||||
self.0 |
||||
.open_tree(USERID_PASSWORD) |
||||
.unwrap() |
||||
self.db |
||||
.userid_password |
||||
.get(user_id.to_string()) |
||||
.unwrap() |
||||
.map(|bytes| utils::bytes_to_string(&bytes)) |
||||
.map(|bytes| utils::string_from_bytes(&bytes)) |
||||
} |
||||
|
||||
/// Add a new device to a user.
|
||||
pub fn device_add(&self, user_id: &UserId, device_id: &str) { |
||||
self.0 |
||||
.open_tree(USERID_DEVICEIDS) |
||||
.unwrap() |
||||
.insert(user_id.to_string(), device_id) |
||||
.unwrap(); |
||||
if self |
||||
.db |
||||
.userid_deviceids |
||||
.get_iter(&user_id.to_string().as_bytes()) |
||||
.filter_map(|item| item.ok()) |
||||
.map(|(_key, value)| value) |
||||
.all(|device| device != device_id) |
||||
{ |
||||
self.db |
||||
.userid_deviceids |
||||
.add(user_id.to_string().as_bytes(), device_id.into()); |
||||
} |
||||
} |
||||
|
||||
/// Replace the access token of one device.
|
||||
pub fn token_replace(&self, user_id: &UserId, device_id: &String, token: String) { |
||||
// Make sure the device id belongs to the user
|
||||
debug_assert!(self |
||||
.0 |
||||
.open_tree(USERID_DEVICEIDS) |
||||
.unwrap() |
||||
.get(&user_id.to_string()) // Does the user exist?
|
||||
.unwrap() |
||||
.map(|bytes| utils::bytes_to_vec(&bytes)) |
||||
.filter(|devices| devices.contains(device_id)) // Does the user have that device?
|
||||
.is_some()); |
||||
.db |
||||
.userid_deviceids |
||||
.get_iter(&user_id.to_string().as_bytes()) |
||||
.filter_map(|item| item.ok()) |
||||
.map(|(_key, value)| value) |
||||
.any(|device| device == device_id.as_bytes())); // Does the user have that device?
|
||||
|
||||
// Remove old token
|
||||
if let Some(old_token) = self |
||||
.0 |
||||
.open_tree(DEVICEID_TOKEN) |
||||
.unwrap() |
||||
.get(device_id) |
||||
.unwrap() |
||||
{ |
||||
self.0 |
||||
.open_tree(TOKEN_USERID) |
||||
.unwrap() |
||||
.remove(old_token) |
||||
.unwrap(); |
||||
// It will be removed from DEVICEID_TOKEN by the insert later
|
||||
if let Some(old_token) = self.db.deviceid_token.get(device_id).unwrap() { |
||||
self.db.token_userid.remove(old_token).unwrap(); |
||||
// It will be removed from deviceid_token by the insert later
|
||||
} |
||||
|
||||
// Assign token to device_id
|
||||
self.0 |
||||
.open_tree(DEVICEID_TOKEN) |
||||
.unwrap() |
||||
.insert(device_id, &*token) |
||||
.unwrap(); |
||||
self.db.deviceid_token.insert(device_id, &*token).unwrap(); |
||||
|
||||
// Assign token to user
|
||||
self.0 |
||||
.open_tree(TOKEN_USERID) |
||||
.unwrap() |
||||
self.db |
||||
.token_userid |
||||
.insert(token, &*user_id.to_string()) |
||||
.unwrap(); |
||||
} |
||||
|
||||
/// Create a new room event.
|
||||
pub fn event_add(&self, event: &Event, room_id: &RoomId, event_id: &EventId) { |
||||
debug!("{}", serde_json::to_string(event).unwrap()); |
||||
todo!(); |
||||
pub fn event_add(&self, room_id: &RoomId, event_id: &EventId, event: &Event) { |
||||
let mut key = room_id.to_string().as_bytes().to_vec(); |
||||
key.extend_from_slice(event_id.to_string().as_bytes()); |
||||
self.db |
||||
.roomid_eventid_event |
||||
.insert(&key, &*serde_json::to_string(event).unwrap()) |
||||
.unwrap(); |
||||
} |
||||
|
||||
pub fn debug(&self) { |
||||
self.db.debug(); |
||||
} |
||||
} |
||||
|
||||
@ -0,0 +1,117 @@
@@ -0,0 +1,117 @@
|
||||
use crate::utils; |
||||
use directories::ProjectDirs; |
||||
use sled::IVec; |
||||
|
||||
pub struct MultiValue(sled::Tree); |
||||
|
||||
impl MultiValue { |
||||
/// Get an iterator over all values.
|
||||
pub fn iter_all(&self) -> sled::Iter { |
||||
self.0.iter() |
||||
} |
||||
|
||||
/// Get an iterator over all values of this id.
|
||||
pub fn get_iter(&self, id: &[u8]) -> sled::Iter { |
||||
// Data keys start with d
|
||||
let mut key = vec![b'd']; |
||||
key.extend_from_slice(id.as_ref()); |
||||
key.push(0xff); // Add delimiter so we don't find usernames starting with the same id
|
||||
|
||||
self.0.scan_prefix(key) |
||||
} |
||||
|
||||
/// Add another value to the id.
|
||||
pub fn add(&self, id: &[u8], value: IVec) { |
||||
// The new value will need a new index. We store the last used index in 'n' + id
|
||||
let mut count_key: Vec<u8> = vec![b'n']; |
||||
count_key.extend_from_slice(id.as_ref()); |
||||
|
||||
// Increment the last index and use that
|
||||
let index = self |
||||
.0 |
||||
.update_and_fetch(&count_key, utils::increment) |
||||
.unwrap() |
||||
.unwrap(); |
||||
|
||||
// Data keys start with d
|
||||
let mut key = vec![b'd']; |
||||
key.extend_from_slice(id.as_ref()); |
||||
key.push(0xff); |
||||
key.extend_from_slice(&index); |
||||
|
||||
self.0.insert(key, value).unwrap(); |
||||
} |
||||
} |
||||
|
||||
pub struct Database { |
||||
pub userid_password: sled::Tree, |
||||
pub userid_deviceids: MultiValue, |
||||
pub deviceid_token: sled::Tree, |
||||
pub token_userid: sled::Tree, |
||||
pub roomid_eventid_event: sled::Tree, |
||||
_db: sled::Db, |
||||
} |
||||
|
||||
impl Database { |
||||
/// Load an existing database or create a new one.
|
||||
pub fn load_or_create(hostname: &str) -> Self { |
||||
let mut path = ProjectDirs::from("xyz", "koesters", "matrixserver") |
||||
.unwrap() |
||||
.data_dir() |
||||
.to_path_buf(); |
||||
path.push(hostname); |
||||
let db = sled::open(&path).unwrap(); |
||||
|
||||
Self { |
||||
userid_password: db.open_tree("userid_password").unwrap(), |
||||
userid_deviceids: MultiValue(db.open_tree("userid_deviceids").unwrap()), |
||||
deviceid_token: db.open_tree("deviceid_token").unwrap(), |
||||
token_userid: db.open_tree("token_userid").unwrap(), |
||||
roomid_eventid_event: db.open_tree("roomid_eventid_event").unwrap(), |
||||
_db: db, |
||||
} |
||||
} |
||||
|
||||
pub fn debug(&self) { |
||||
println!("# UserId -> Password:"); |
||||
for (k, v) in self.userid_password.iter().map(|r| r.unwrap()) { |
||||
println!( |
||||
"{} -> {}", |
||||
String::from_utf8_lossy(&k), |
||||
String::from_utf8_lossy(&v), |
||||
); |
||||
} |
||||
println!("# UserId -> DeviceIds:"); |
||||
for (k, v) in self.userid_deviceids.iter_all().map(|r| r.unwrap()) { |
||||
println!( |
||||
"{} -> {}", |
||||
String::from_utf8_lossy(&k), |
||||
String::from_utf8_lossy(&v), |
||||
); |
||||
} |
||||
println!("# DeviceId -> Token:"); |
||||
for (k, v) in self.deviceid_token.iter().map(|r| r.unwrap()) { |
||||
println!( |
||||
"{} -> {}", |
||||
String::from_utf8_lossy(&k), |
||||
String::from_utf8_lossy(&v), |
||||
); |
||||
} |
||||
println!("# Token -> UserId:"); |
||||
for (k, v) in self.token_userid.iter().map(|r| r.unwrap()) { |
||||
println!( |
||||
"{} -> {}", |
||||
String::from_utf8_lossy(&k), |
||||
String::from_utf8_lossy(&v), |
||||
); |
||||
} |
||||
println!("# RoomId + EventId -> Event:"); |
||||
for (k, v) in self.roomid_eventid_event.iter().map(|r| r.unwrap()) { |
||||
println!( |
||||
"{} -> {}", |
||||
String::from_utf8_lossy(&k), |
||||
String::from_utf8_lossy(&v), |
||||
); |
||||
} |
||||
} |
||||
} |
||||
Loading…
Reference in new issue