Skip to content
Snippets Groups Projects

Compare revisions

Changes are shown as if the source revision was being merged into the target revision. Learn more about comparing revisions.

Source

Select target project
No results found

Target

Select target project
No results found
Show changes
Hide whitespace changes
Inline Side-by-side
Showing
with 1183 additions and 1026 deletions
src/database/migrations/mod.rs
View file @ 33262cc7
......@@ -3,16 +3,17 @@ use super::get_connection;
pub mod init;
pub mod scripts;
pub fn run_migrations() {
pub async fn run_migrations() {
let client = get_connection();
let list = client
.list_database_names(None, None)
.await
.expect("Failed to fetch database names.");
if list.iter().position(|x| x == "revolt").is_none() {
init::create_database();
init::create_database().await;
} else {
scripts::migrate_database();
scripts::migrate_database().await;
}
}
src/database/migrations/scripts.rs
View file @ 33262cc7
use super::super::get_collection;
use crate::database::{permissions, get_collection, get_db, PermissionTuple};
use futures::StreamExt;
use log::info;
use mongodb::bson::{doc, from_bson, Bson};
use mongodb::options::FindOptions;
use mongodb::{bson::{doc, from_document, to_document}, options::FindOptions};
use serde::{Deserialize, Serialize};
#[derive(Serialize, Deserialize)]
......@@ -11,19 +11,20 @@ struct MigrationInfo {
revision: i32,
}
pub const LATEST_REVISION: i32 = 2;
pub const LATEST_REVISION: i32 = 7;
pub fn migrate_database() {
pub async fn migrate_database() {
let migrations = get_collection("migrations");
let data = migrations
.find_one(None, None)
.await
.expect("Failed to fetch migration data.");
if let Some(doc) = data {
let info: MigrationInfo =
from_bson(Bson::Document(doc)).expect("Failed to read migration information.");
from_document(doc).expect("Failed to read migration information.");
let revision = run_migrations(info.revision);
let revision = run_migrations(info.revision).await;
migrations
.update_one(
......@@ -37,6 +38,7 @@ pub fn migrate_database() {
},
None,
)
.await
.expect("Failed to commit migration information.");
info!("Migration complete. Currently at revision {}.", revision);
......@@ -45,7 +47,7 @@ pub fn migrate_database() {
}
}
pub fn run_migrations(revision: i32) -> i32 {
pub async fn run_migrations(revision: i32) -> i32 {
info!("Starting database migration.");
if revision <= 0 {
......@@ -53,65 +55,152 @@ pub fn run_migrations(revision: i32) -> i32 {
}
if revision <= 1 {
info!("Running migration [revision 1]: Add channels to guild object.");
info!("Running migration [revision 1 / 2021-04-24]: Migrate to Autumn v1.0.0.");
let col = get_collection("guilds");
let guilds = col
.find(
let messages = get_collection("messages");
let attachments = get_collection("attachments");
messages
.update_many(
doc! { "attachment": { "$exists": 1 } },
doc! { "$set": { "attachment.tag": "attachments", "attachment.size": 0 } },
None,
)
.await
.expect("Failed to update messages.");
attachments
.update_many(
doc! {},
doc! { "$set": { "tag": "attachments", "size": 0 } },
None,
FindOptions::builder().projection(doc! { "_id": 1 }).build(),
)
.expect("Failed to fetch guilds.");
.await
.expect("Failed to update attachments.");
}
if revision <= 2 {
info!("Running migration [revision 2 / 2021-05-08]: Add servers collection.");
get_db()
.create_collection("servers", None)
.await
.expect("Failed to create servers collection.");
}
if revision <= 3 {
info!("Running migration [revision 3 / 2021-05-25]: Support multiple file uploads, add channel_unreads and user_settings.");
let result = get_collection("channels")
let messages = get_collection("messages");
let mut cursor = messages
.find(
doc! {
"type": 2
"attachment": {
"$exists": 1
}
},
FindOptions::builder()
.projection(doc! { "_id": 1, "guild": 1 })
.projection(doc! {
"_id": 1,
"attachments": [ "$attachment" ]
})
.build(),
)
.expect("Failed to fetch channels.");
let mut channels = vec![];
for doc in result {
let channel = doc.expect("Failed to fetch channel.");
let id = channel
.get_str("_id")
.expect("Failed to get channel id.")
.to_string();
let gid = channel
.get_str("guild")
.expect("Failed to get guild id.")
.to_string();
channels.push((id, gid));
.await
.expect("Failed to fetch messages.");
while let Some(result) = cursor.next().await {
let doc = result.unwrap();
let id = doc.get_str("_id").unwrap();
let attachments = doc.get_array("attachments").unwrap();
messages
.update_one(
doc! { "_id": id },
doc! { "$unset": { "attachment": 1 }, "$set": { "attachments": attachments } },
None,
)
.await
.unwrap();
}
for doc in guilds {
let guild = doc.expect("Failed to fetch guild.");
let id = guild.get_str("_id").expect("Failed to get guild id.");
get_db()
.create_collection("channel_unreads", None)
.await
.expect("Failed to create channel_unreads collection.");
get_db()
.create_collection("user_settings", None)
.await
.expect("Failed to create user_settings collection.");
}
let list: Vec<String> = channels
.iter()
.filter(|x| x.1 == id)
.map(|x| x.0.clone())
.collect();
if revision <= 4 {
info!("Running migration [revision 4 / 2021-06-01]: Add more server collections.");
col.update_one(
doc! {
"_id": id
},
get_db()
.create_collection("server_members", None)
.await
.expect("Failed to create server_members collection.");
get_db()
.create_collection("server_bans", None)
.await
.expect("Failed to create server_bans collection.");
get_db()
.create_collection("channel_invites", None)
.await
.expect("Failed to create channel_invites collection.");
}
if revision <= 5 {
info!("Running migration [revision 5 / 2021-06-26]: Add permissions.");
#[derive(Serialize)]
struct Server {
pub default_permissions: PermissionTuple,
}
let server = Server {
default_permissions: (
*permissions::server::DEFAULT_PERMISSION as i32,
*permissions::channel::DEFAULT_PERMISSION_SERVER as i32
)
};
get_collection("servers")
.update_many(
doc! { },
doc! {
"$set": {
"channels": list
}
"$set": to_document(&server).unwrap()
},
None,
None
)
.expect("Failed to update guild.");
}
.await
.expect("Failed to migrate servers.");
}
if revision <= 6 {
info!("Running migration [revision 6 / 2021-07-09]: Add message text index.");
get_db()
.run_command(
doc! {
"createIndexes": "messages",
"indexes": [
{
"key": {
"content": "text"
},
"name": "content"
}
]
},
None,
)
.await
.expect("Failed to create message index.");
}
// Reminder to update LATEST_REVISION when adding new migrations.
......
src/database/mod.rs
View file @ 33262cc7
use mongodb::sync::{Client, Collection, Database};
use std::env;
use crate::util::variables::MONGO_URI;
use mongodb::{Client, Collection, Database};
use once_cell::sync::OnceCell;
static DBCONN: OnceCell<Client> = OnceCell::new();
pub fn connect() {
let client =
Client::with_uri_str(&env::var("DB_URI").expect("DB_URI not in environment variables!"))
.expect("Failed to init db connection.");
pub async fn connect() {
let client = Client::with_uri_str(&MONGO_URI)
.await
.expect("Failed to init db connection.");
DBCONN.set(client).unwrap();
migrations::run_migrations();
migrations::run_migrations().await;
}
pub fn get_connection() -> &'static Client {
......@@ -25,13 +26,11 @@ pub fn get_collection(collection: &str) -> Collection {
get_db().collection(collection)
}
pub mod entities;
pub mod guards;
pub mod migrations;
pub mod channel;
pub mod guild;
pub mod message;
pub mod mutual;
pub mod permissions;
pub mod user;
pub use entities::*;
pub use guards::*;
pub use permissions::*;
src/database/mutual.rs deleted 100644 → 0
View file @ cbac8029
use super::{get_collection, MemberPermissions};
use mongodb::bson::doc;
use mongodb::options::FindOptions;
pub fn find_mutual_guilds(user_id: &str, target_id: &str) -> Vec<String> {
let col = get_collection("members");
if let Ok(result) = col.find(
doc! {
"$and": [
{ "id": user_id },
{ "id": target_id },
]
},
FindOptions::builder().projection(doc! { "_id": 1 }).build(),
) {
let mut results = vec![];
for doc in result {
if let Ok(guild) = doc {
results.push(guild.get_str("_id").unwrap().to_string());
}
}
results
} else {
vec![]
}
}
pub fn find_mutual_friends(user_id: &str, target_id: &str) -> Vec<String> {
let col = get_collection("users");
if let Ok(result) = col.find(
doc! {
"$and": [
{ "relations": { "$elemMatch": { "id": user_id, "status": 0 } } },
{ "relations": { "$elemMatch": { "id": target_id, "status": 0 } } },
]
},
FindOptions::builder().projection(doc! { "_id": 1 }).build(),
) {
let mut results = vec![];
for doc in result {
if let Ok(user) = doc {
results.push(user.get_str("_id").unwrap().to_string());
}
}
results
} else {
vec![]
}
}
pub fn find_mutual_groups(user_id: &str, target_id: &str) -> Vec<String> {
let col = get_collection("channels");
if let Ok(result) = col.find(
doc! {
"type": 1,
"$and": [
{ "recipients": user_id },
{ "recipients": target_id },
]
},
FindOptions::builder().projection(doc! { "_id": 1 }).build(),
) {
let mut results = vec![];
for doc in result {
if let Ok(group) = doc {
results.push(group.get_str("_id").unwrap().to_string());
}
}
results
} else {
vec![]
}
}
pub fn has_mutual_connection(user_id: &str, target_id: &str, with_permission: bool) -> bool {
let mut doc = doc! { "_id": 1 };
if with_permission {
doc.insert("default_permissions", 1);
}
let opt = FindOptions::builder().projection(doc);
if let Ok(result) = get_collection("guilds").find(
doc! {
"$and": [
{ "members": { "$elemMatch": { "id": user_id } } },
{ "members": { "$elemMatch": { "id": target_id } } },
]
},
if with_permission {
opt.build()
} else {
opt.limit(1).build()
},
) {
if with_permission {
for item in result {
// ? logic should match permissions.rs#calculate
if let Ok(guild) = item {
if guild.get_str("owner").unwrap() == user_id {
return true;
}
let permissions = guild.get_i32("default_permissions").unwrap() as u32;
if MemberPermissions([permissions]).get_send_direct_messages() {
return true;
}
}
}
false
} else if result.count() > 0 {
true
} else {
false
}
} else {
false
}
}
src/database/permissions.rs deleted 100644 → 0
View file @ cbac8029
use super::mutual::has_mutual_connection;
use crate::database::channel::Channel;
use crate::database::guild::{fetch_guild, fetch_member, Guild, Member, MemberKey};
use crate::database::user::{User, UserRelationship};
use num_enum::TryFromPrimitive;
#[derive(Debug, PartialEq, Eq, TryFromPrimitive)]
#[repr(u8)]
pub enum Relationship {
Friend = 0,
Outgoing = 1,
Incoming = 2,
Blocked = 3,
BlockedOther = 4,
NONE = 5,
SELF = 6,
}
#[derive(Debug, PartialEq, Eq, TryFromPrimitive, Copy, Clone)]
#[repr(u32)]
pub enum Permission {
Access = 1,
CreateInvite = 2,
KickMembers = 4,
BanMembers = 8,
ReadMessages = 16,
SendMessages = 32,
ManageMessages = 64,
ManageChannels = 128,
ManageServer = 256,
ManageRoles = 512,
SendDirectMessages = 1024,
}
bitfield! {
pub struct MemberPermissions(MSB0 [u32]);
u32;
pub get_access, set_access: 31;
pub get_create_invite, set_create_invite: 30;
pub get_kick_members, set_kick_members: 29;
pub get_ban_members, set_ban_members: 28;
pub get_read_messages, set_read_messages: 27;
pub get_send_messages, set_send_messages: 26;
pub get_manage_messages, set_manage_messages: 25;
pub get_manage_channels, set_manage_channels: 24;
pub get_manage_server, set_manage_server: 23;
pub get_manage_roles, set_manage_roles: 22;
pub get_send_direct_messages, set_send_direct_messages: 21;
}
pub fn get_relationship_internal(
user_id: &str,
target_id: &str,
relationships: &Option<Vec<UserRelationship>>,
) -> Relationship {
if user_id == target_id {
return Relationship::SELF;
}
if let Some(arr) = &relationships {
for entry in arr {
if entry.id == target_id {
match entry.status {
0 => return Relationship::Friend,
1 => return Relationship::Outgoing,
2 => return Relationship::Incoming,
3 => return Relationship::Blocked,
4 => return Relationship::BlockedOther,
_ => return Relationship::NONE,
}
}
}
}
Relationship::NONE
}
pub fn get_relationship(a: &User, b: &User) -> Relationship {
if a.id == b.id {
return Relationship::SELF;
}
get_relationship_internal(&a.id, &b.id, &a.relations)
}
pub struct PermissionCalculator {
pub user: User,
pub channel: Option<Channel>,
pub guild: Option<Guild>,
pub member: Option<Member>,
}
impl PermissionCalculator {
pub fn new(user: User) -> PermissionCalculator {
PermissionCalculator {
user,
channel: None,
guild: None,
member: None,
}
}
pub fn channel(self, channel: Channel) -> PermissionCalculator {
PermissionCalculator {
channel: Some(channel),
..self
}
}
pub fn guild(self, guild: Guild) -> PermissionCalculator {
PermissionCalculator {
guild: Some(guild),
..self
}
}
pub fn fetch_data(mut self) -> PermissionCalculator {
let guild = if let Some(value) = self.guild {
Some(value)
} else if let Some(channel) = &self.channel {
match channel.channel_type {
0..=1 => None,
2 => {
if let Some(id) = &channel.guild {
if let Ok(result) = fetch_guild(id) {
result
} else {
None
}
} else {
None
}
}
_ => None,
}
} else {
None
};
if let Some(guild) = &guild {
if let Ok(result) = fetch_member(MemberKey(guild.id.clone(), self.user.id.clone())) {
self.member = result;
}
}
self.guild = guild;
self
}
pub fn calculate(&self) -> u32 {
let mut permissions: u32 = 0;
if let Some(guild) = &self.guild {
if let Some(_member) = &self.member {
// ? logic should match mutual.rs#has_mutual_connection
if guild.owner == self.user.id {
return u32::MAX;
}
permissions = guild.default_permissions as u32;
}
}
if let Some(channel) = &self.channel {
match channel.channel_type {
0 => {
if let Some(arr) = &channel.recipients {
let mut other_user = None;
for item in arr {
if item != &self.user.id {
other_user = Some(item);
}
}
if let Some(other) = other_user {
let relationship = get_relationship_internal(
&self.user.id,
&other,
&self.user.relations,
);
if relationship == Relationship::Friend {
permissions = 1024 + 128 + 32 + 16 + 1;
} else if relationship == Relationship::Blocked
|| relationship == Relationship::BlockedOther
{
permissions = 1;
} else if has_mutual_connection(&self.user.id, other, true) {
permissions = 1024 + 128 + 32 + 16 + 1;
} else {
permissions = 1;
}
} else {
// ? In this case, it is a "self DM".
return 1024 + 128 + 32 + 16 + 1;
}
}
}
1 => {
if let Some(id) = &channel.owner {
if &self.user.id == id {
return u32::MAX;
}
}
if let Some(arr) = &channel.recipients {
for item in arr {
if item == &self.user.id {
permissions = 177;
break;
}
}
}
}
2 => {
// nothing implemented yet
}
_ => {}
}
}
permissions
}
pub fn as_permission(&self) -> MemberPermissions<[u32; 1]> {
MemberPermissions([self.calculate()])
}
}
src/database/permissions/channel.rs 0 → 100644
View file @ 33262cc7
use crate::database::*;
use crate::util::result::{Error, Result};
use super::PermissionCalculator;
use num_enum::TryFromPrimitive;
use std::ops;
#[derive(Debug, PartialEq, Eq, TryFromPrimitive, Copy, Clone)]
#[repr(u32)]
pub enum ChannelPermission {
View = 0b00000000000000000000000000000001, // 1
SendMessage = 0b00000000000000000000000000000010, // 2
ManageMessages = 0b00000000000000000000000000000100, // 4
ManageChannel = 0b00000000000000000000000000001000, // 8
VoiceCall = 0b00000000000000000000000000010000, // 16
InviteOthers = 0b00000000000000000000000000100000, // 32
EmbedLinks = 0b00000000000000000000000001000000, // 64
UploadFiles = 0b00000000000000000000000010000000, // 128
}
lazy_static! {
pub static ref DEFAULT_PERMISSION_DM: u32 =
ChannelPermission::View
+ ChannelPermission::SendMessage
+ ChannelPermission::ManageChannel
+ ChannelPermission::VoiceCall
+ ChannelPermission::InviteOthers
+ ChannelPermission::EmbedLinks
+ ChannelPermission::UploadFiles;
pub static ref DEFAULT_PERMISSION_SERVER: u32 =
ChannelPermission::View
+ ChannelPermission::SendMessage
+ ChannelPermission::VoiceCall
+ ChannelPermission::InviteOthers
+ ChannelPermission::EmbedLinks
+ ChannelPermission::UploadFiles;
}
impl_op_ex!(+ |a: &ChannelPermission, b: &ChannelPermission| -> u32 { *a as u32 | *b as u32 });
impl_op_ex_commutative!(+ |a: &u32, b: &ChannelPermission| -> u32 { *a | *b as u32 });
bitfield! {
pub struct ChannelPermissions(MSB0 [u32]);
u32;
pub get_view, _: 31;
pub get_send_message, _: 30;
pub get_manage_messages, _: 29;
pub get_manage_channel, _: 28;
pub get_voice_call, _: 27;
pub get_invite_others, _: 26;
pub get_embed_links, _: 25;
pub get_upload_files, _: 24;
}
impl<'a> PermissionCalculator<'a> {
pub async fn calculate_channel(self) -> Result<u32> {
let channel = if let Some(channel) = self.channel {
channel
} else {
unreachable!()
};
match channel {
Channel::SavedMessages { user: owner, .. } => {
if &self.perspective.id == owner {
Ok(u32::MAX)
} else {
Ok(0)
}
}
Channel::DirectMessage { recipients, .. } => {
if recipients
.iter()
.find(|x| *x == &self.perspective.id)
.is_some()
{
if let Some(recipient) = recipients.iter().find(|x| *x != &self.perspective.id)
{
let perms = self.for_user(recipient).await?;
if perms.get_send_message() {
return Ok(*DEFAULT_PERMISSION_DM);
}
return Ok(ChannelPermission::View as u32);
}
}
Ok(0)
}
Channel::Group { recipients, permissions, owner, .. } => {
if &self.perspective.id == owner {
return Ok(*DEFAULT_PERMISSION_DM)
}
if recipients
.iter()
.find(|x| *x == &self.perspective.id)
.is_some()
{
if let Some(permissions) = permissions {
Ok(permissions.clone() as u32)
} else {
Ok(*DEFAULT_PERMISSION_DM)
}
} else {
Ok(0)
}
}
Channel::TextChannel { server, default_permissions, role_permissions, .. }
| Channel::VoiceChannel { server, default_permissions, role_permissions, .. } => {
let server = Ref::from_unchecked(server.clone()).fetch_server().await?;
if self.perspective.id == server.owner {
Ok(u32::MAX)
} else {
match Ref::from_unchecked(self.perspective.id.clone()).fetch_member(&server.id).await {
Ok(member) => {
let mut perm = if let Some(permission) = default_permissions {
*permission as u32
} else {
server.default_permissions.1 as u32
};
if let Some(roles) = member.roles {
for role in roles {
if let Some(permission) = role_permissions.get(&role) {
perm |= *permission as u32;
}
if let Some(server_role) = server.roles.get(&role) {
perm |= server_role.permissions.1 as u32;
}
}
}
Ok(perm)
}
Err(error) => {
match &error {
Error::NotFound => Ok(0),
_ => Err(error)
}
}
}
}
}
}
}
pub async fn for_channel(self) -> Result<ChannelPermissions<[u32; 1]>> {
Ok(ChannelPermissions([self.calculate_channel().await?]))
}
}
src/database/permissions/mod.rs 0 → 100644
View file @ 33262cc7
pub use crate::database::*;
pub mod channel;
pub mod server;
pub mod user;
pub use user::get_relationship;
pub struct PermissionCalculator<'a> {
perspective: &'a User,
user: Option<&'a User>,
relationship: Option<&'a RelationshipStatus>,
channel: Option<&'a Channel>,
server: Option<&'a Server>,
// member: Option<&'a Member>,
has_mutual_connection: bool,
}
impl<'a> PermissionCalculator<'a> {
pub fn new(perspective: &'a User) -> PermissionCalculator {
PermissionCalculator {
perspective,
user: None,
relationship: None,
channel: None,
server: None,
// member: None,
has_mutual_connection: false,
}
}
pub fn with_user(self, user: &'a User) -> PermissionCalculator {
PermissionCalculator {
user: Some(&user),
..self
}
}
pub fn with_relationship(self, relationship: &'a RelationshipStatus) -> PermissionCalculator {
PermissionCalculator {
relationship: Some(&relationship),
..self
}
}
pub fn with_channel(self, channel: &'a Channel) -> PermissionCalculator {
PermissionCalculator {
channel: Some(&channel),
..self
}
}
pub fn with_server(self, server: &'a Server) -> PermissionCalculator {
PermissionCalculator {
server: Some(&server),
..self
}
}
/* pub fn with_member(self, member: &'a Member) -> PermissionCalculator {
PermissionCalculator {
member: Some(&member),
..self
}
} */
pub fn with_mutual_connection(self) -> PermissionCalculator<'a> {
PermissionCalculator {
has_mutual_connection: true,
..self
}
}
}
src/database/permissions/server.rs 0 → 100644
View file @ 33262cc7
use crate::util::result::{Error, Result};
use super::PermissionCalculator;
use super::Ref;
use num_enum::TryFromPrimitive;
use std::ops;
#[derive(Debug, PartialEq, Eq, TryFromPrimitive, Copy, Clone)]
#[repr(u32)]
pub enum ServerPermission {
View = 0b00000000000000000000000000000001, // 1
ManageRoles = 0b00000000000000000000000000000010, // 2
ManageChannels = 0b00000000000000000000000000000100, // 4
ManageServer = 0b00000000000000000000000000001000, // 8
KickMembers = 0b00000000000000000000000000010000, // 16
BanMembers = 0b00000000000000000000000000100000, // 32
// 6 bits of space
ChangeNickname = 0b00000000000000000001000000000000, // 4096
ManageNicknames = 0b00000000000000000010000000000000, // 8192
ChangeAvatar = 0b00000000000000000100000000000000, // 16382
RemoveAvatars = 0b00000000000000001000000000000000, // 32768
// 16 bits of space
}
lazy_static! {
pub static ref DEFAULT_PERMISSION: u32 =
ServerPermission::View
+ ServerPermission::ChangeNickname
+ ServerPermission::ChangeAvatar;
}
impl_op_ex!(+ |a: &ServerPermission, b: &ServerPermission| -> u32 { *a as u32 | *b as u32 });
impl_op_ex_commutative!(+ |a: &u32, b: &ServerPermission| -> u32 { *a | *b as u32 });
bitfield! {
pub struct ServerPermissions(MSB0 [u32]);
u32;
pub get_view, _: 31;
pub get_manage_roles, _: 30;
pub get_manage_channels, _: 29;
pub get_manage_server, _: 28;
pub get_kick_members, _: 27;
pub get_ban_members, _: 26;
pub get_change_nickname, _: 19;
pub get_manage_nicknames, _: 18;
pub get_change_avatar, _: 17;
pub get_remove_avatars, _: 16;
}
impl<'a> PermissionCalculator<'a> {
pub async fn calculate_server(self) -> Result<u32> {
let server = if let Some(server) = self.server {
server
} else {
unreachable!()
};
if self.perspective.id == server.owner {
Ok(u32::MAX)
} else {
match Ref::from_unchecked(self.perspective.id.clone()).fetch_member(&server.id).await {
Ok(member) => {
let mut perm = server.default_permissions.0 as u32;
if let Some(roles) = member.roles {
for role in roles {
if let Some(server_role) = server.roles.get(&role) {
perm |= server_role.permissions.0 as u32;
}
}
}
Ok(perm)
}
Err(error) => {
match &error {
Error::NotFound => Ok(0),
_ => Err(error)
}
}
}
}
}
pub async fn for_server(self) -> Result<ServerPermissions<[u32; 1]>> {
Ok(ServerPermissions([self.calculate_server().await?]))
}
}
src/database/permissions/user.rs 0 → 100644
View file @ 33262cc7
use crate::database::*;
use crate::util::result::{Error, Result};
use super::PermissionCalculator;
use mongodb::bson::doc;
use num_enum::TryFromPrimitive;
use std::ops;
#[derive(Debug, PartialEq, Eq, TryFromPrimitive, Copy, Clone)]
#[repr(u32)]
pub enum UserPermission {
Access = 0b00000000000000000000000000000001, // 1
ViewProfile = 0b00000000000000000000000000000010, // 2
SendMessage = 0b00000000000000000000000000000100, // 4
Invite = 0b00000000000000000000000000001000, // 8
}
bitfield! {
pub struct UserPermissions(MSB0 [u32]);
u32;
pub get_access, _: 31;
pub get_view_profile, _: 30;
pub get_send_message, _: 29;
pub get_invite, _: 28;
}
impl_op_ex!(+ |a: &UserPermission, b: &UserPermission| -> u32 { *a as u32 | *b as u32 });
impl_op_ex_commutative!(+ |a: &u32, b: &UserPermission| -> u32 { *a | *b as u32 });
pub fn get_relationship(a: &User, b: &str) -> RelationshipStatus {
if a.id == b {
return RelationshipStatus::User;
}
if let Some(relations) = &a.relations {
if let Some(relationship) = relations.iter().find(|x| x.id == b) {
return relationship.status.clone();
}
}
RelationshipStatus::None
}
impl<'a> PermissionCalculator<'a> {
pub async fn calculate_user(self, target: &str) -> Result<u32> {
if &self.perspective.id == target {
return Ok(u32::MAX);
}
let mut permissions: u32 = 0;
match self
.relationship
.clone()
.map(|v| v.to_owned())
.unwrap_or_else(|| get_relationship(&self.perspective, &target))
{
RelationshipStatus::Friend | RelationshipStatus::User => return Ok(u32::MAX),
RelationshipStatus::Blocked | RelationshipStatus::BlockedOther => {
return Ok(UserPermission::Access as u32)
}
RelationshipStatus::Incoming | RelationshipStatus::Outgoing => {
permissions = UserPermission::Access as u32;
// ! INFO: if we add boolean switch for permission to
// ! message people who have mutual, we need to get
// ! rid of this return statement.
// return Ok(permissions);
}
_ => {}
}
let check_server_overlap = async || {
let server_ids = User::fetch_server_ids(&self.perspective.id).await?;
Ok(
get_collection("server_members")
.find_one(
doc! {
"_id.user": &target,
"_id.server": {
"$in": server_ids
}
},
None
)
.await
.map_err(|_| Error::DatabaseError {
operation: "find_one",
with: "server_members",
})?
.is_some()
)
};
if self.has_mutual_connection
|| check_server_overlap().await?
|| get_collection("channels")
.find_one(
doc! {
"channel_type": {
"$in": ["Group", "DirectMessage"]
},
"recipients": {
"$all": [ &self.perspective.id, target ]
}
},
None,
)
.await
.map_err(|_| Error::DatabaseError {
operation: "find_one",
with: "channels",
})?
.is_some()
{
// ! FIXME: add privacy settings
return Ok(UserPermission::Access + UserPermission::ViewProfile);
}
Ok(permissions)
}
pub async fn for_user(self, target: &str) -> Result<UserPermissions<[u32; 1]>> {
Ok(UserPermissions([self.calculate_user(&target).await?]))
}
pub async fn for_user_given(self) -> Result<UserPermissions<[u32; 1]>> {
let id = &self.user.unwrap().id;
Ok(UserPermissions([self.calculate_user(&id).await?]))
}
}
src/database/user.rs deleted 100644 → 0
View file @ cbac8029
use super::channel::fetch_channels;
use super::get_collection;
use super::guild::serialise_guilds_with_channels;
use lru::LruCache;
use mongodb::bson::{doc, from_bson, Bson, DateTime};
use mongodb::options::FindOptions;
use rocket::http::{RawStr, Status};
use rocket::request::{self, FromParam, FromRequest, Request};
use rocket::Outcome;
use rocket_contrib::json::JsonValue;
use serde::{Deserialize, Serialize};
use std::sync::{Arc, Mutex};
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct UserEmailVerification {
pub verified: bool,
pub target: Option<String>,
pub expiry: Option<DateTime>,
pub rate_limit: Option<DateTime>,
pub code: Option<String>,
}
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct UserRelationship {
pub id: String,
pub status: u8,
}
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct User {
#[serde(rename = "_id")]
pub id: String,
pub email: String,
pub username: String,
pub password: String,
pub display_name: String,
pub access_token: Option<String>,
pub email_verification: UserEmailVerification,
pub relations: Option<Vec<UserRelationship>>,
}
impl User {
pub fn serialise(self, relationship: i32) -> JsonValue {
if relationship == super::Relationship::SELF as i32 {
json!({
"id": self.id,
"username": self.username,
"display_name": self.display_name,
"email": self.email,
"verified": self.email_verification.verified,
})
} else {
json!({
"id": self.id,
"username": self.username,
"display_name": self.display_name,
"relationship": relationship
})
}
}
pub fn find_guilds(&self) -> Result<Vec<String>, String> {
let members = get_collection("members")
.find(
doc! {
"_id.user": &self.id
},
None,
)
.map_err(|_| "Failed to fetch members.")?;
Ok(members
.into_iter()
.filter_map(|x| match x {
Ok(doc) => match doc.get_document("_id") {
Ok(id) => match id.get_str("guild") {
Ok(value) => Some(value.to_string()),
Err(_) => None,
},
Err(_) => None,
},
Err(_) => None,
})
.collect())
}
pub fn find_dms(&self) -> Result<Vec<String>, String> {
let channels = get_collection("channels")
.find(
doc! {
"recipients": &self.id
},
FindOptions::builder().projection(doc! { "_id": 1 }).build(),
)
.map_err(|_| "Failed to fetch channel ids.")?;
Ok(channels
.into_iter()
.filter_map(|x| x.ok())
.filter_map(|x| match x.get_str("_id") {
Ok(value) => Some(value.to_string()),
Err(_) => None,
})
.collect())
}
pub fn create_payload(self) -> Result<JsonValue, String> {
let v = vec![];
let relations = self.relations.as_ref().unwrap_or(&v);
let users: Vec<JsonValue> = fetch_users(&relations.iter().map(|x| x.id.clone()).collect())?
.into_iter()
.map(|x| {
let id = x.id.clone();
x.serialise(relations.iter().find(|y| y.id == id).unwrap().status as i32)
})
.collect();
let channels: Vec<JsonValue> = fetch_channels(&self.find_dms()?)?
.into_iter()
.map(|x| x.serialise())
.collect();
Ok(json!({
"users": users,
"channels": channels,
"guilds": serialise_guilds_with_channels(&self.find_guilds()?)?,
"user": self.serialise(super::Relationship::SELF as i32)
}))
}
}
lazy_static! {
static ref CACHE: Arc<Mutex<LruCache<String, User>>> =
Arc::new(Mutex::new(LruCache::new(4_000_000)));
}
pub fn fetch_user(id: &str) -> Result<Option<User>, String> {
{
if let Ok(mut cache) = CACHE.lock() {
let existing = cache.get(&id.to_string());
if let Some(user) = existing {
return Ok(Some((*user).clone()));
}
} else {
return Err("Failed to lock cache.".to_string());
}
}
let col = get_collection("users");
if let Ok(result) = col.find_one(doc! { "_id": id }, None) {
if let Some(doc) = result {
if let Ok(user) = from_bson(Bson::Document(doc)) as Result<User, _> {
let mut cache = CACHE.lock().unwrap();
cache.put(id.to_string(), user.clone());
Ok(Some(user))
} else {
Err("Failed to deserialize user!".to_string())
}
} else {
Ok(None)
}
} else {
Err("Failed to fetch user from database.".to_string())
}
}
pub fn fetch_users(ids: &Vec<String>) -> Result<Vec<User>, String> {
let mut missing = vec![];
let mut users = vec![];
{
if let Ok(mut cache) = CACHE.lock() {
for id in ids {
let existing = cache.get(id);
if let Some(user) = existing {
users.push((*user).clone());
} else {
missing.push(id);
}
}
} else {
return Err("Failed to lock cache.".to_string());
}
}
if missing.len() == 0 {
return Ok(users);
}
let col = get_collection("users");
if let Ok(result) = col.find(doc! { "_id": { "$in": missing } }, None) {
for item in result {
let mut cache = CACHE.lock().unwrap();
if let Ok(doc) = item {
if let Ok(user) = from_bson(Bson::Document(doc)) as Result<User, _> {
cache.put(user.id.clone(), user.clone());
users.push(user);
} else {
return Err("Failed to deserialize user!".to_string());
}
} else {
return Err("Failed to fetch user.".to_string());
}
}
Ok(users)
} else {
Err("Failed to fetch user from database.".to_string())
}
}
#[derive(Debug)]
pub enum AuthError {
Failed,
Missing,
Invalid,
}
impl<'a, 'r> FromRequest<'a, 'r> for User {
type Error = AuthError;
fn from_request(request: &'a Request<'r>) -> request::Outcome<Self, Self::Error> {
let u = request.headers().get("x-user").next();
let t = request.headers().get("x-auth-token").next();
if let Some(uid) = u {
if let Some(token) = t {
if let Ok(result) = fetch_user(uid) {
if let Some(user) = result {
if let Some(access_token) = &user.access_token {
if access_token == token {
Outcome::Success(user)
} else {
Outcome::Failure((Status::Forbidden, AuthError::Invalid))
}
} else {
Outcome::Failure((Status::Forbidden, AuthError::Invalid))
}
} else {
Outcome::Failure((Status::Forbidden, AuthError::Invalid))
}
} else {
Outcome::Failure((Status::Forbidden, AuthError::Failed))
}
} else {
Outcome::Failure((Status::Forbidden, AuthError::Missing))
}
} else {
Outcome::Failure((Status::Forbidden, AuthError::Missing))
}
}
}
impl<'r> FromParam<'r> for User {
type Error = &'r RawStr;
fn from_param(param: &'r RawStr) -> Result<Self, Self::Error> {
if let Ok(result) = fetch_user(&param.to_string()) {
if let Some(user) = result {
Ok(user)
} else {
Err(param)
}
} else {
Err(param)
}
}
}
use crate::notifications::events::Notification;
pub fn process_event(event: &Notification) {
match event {
Notification::user_friend_status(ev) => {
let mut cache = CACHE.lock().unwrap();
if let Some(user) = cache.peek_mut(&ev.id) {
if let Some(relations) = user.relations.as_mut() {
if ev.status == 0 {
if let Some(pos) = relations.iter().position(|x| x.id == ev.user) {
relations.remove(pos);
}
} else if let Some(entry) = relations.iter_mut().find(|x| x.id == ev.user) {
entry.status = ev.status as u8;
} else {
relations.push(UserRelationship {
id: ev.id.clone(),
status: ev.status as u8,
});
}
}
}
}
_ => {}
}
}
src/main.rs
View file @ 33262cc7
#![feature(proc_macro_hygiene, decl_macro)]
#![feature(async_closure)]
#[macro_use]
extern crate rocket;
#[macro_use]
extern crate rocket_contrib;
#[macro_use]
extern crate bitfield;
#[macro_use]
extern crate lazy_static;
#[macro_use]
extern crate impl_ops;
#[macro_use]
extern crate bitfield;
extern crate ctrlc;
pub mod database;
pub mod notifications;
pub mod routes;
pub mod util;
pub mod version;
use async_std::task;
use chrono::Duration;
use futures::join;
use log::info;
use rauth::options::{EmailVerification, Options, SMTP};
use rauth::{
auth::Auth,
options::{Template, Templates},
};
use rocket_cors::AllowedOrigins;
use std::thread;
use rocket_prometheus::PrometheusMetrics;
use util::variables::{
APP_URL, HCAPTCHA_KEY, INVITE_ONLY, PUBLIC_URL, SMTP_FROM, SMTP_HOST, SMTP_PASSWORD,
SMTP_USERNAME, USE_EMAIL, USE_HCAPTCHA, USE_PROMETHEUS,
};
fn main() {
#[async_std::main]
async fn main() {
dotenv::dotenv().ok();
env_logger::init();
database::connect();
notifications::start_worker();
env_logger::init_from_env(env_logger::Env::default().filter_or("RUST_LOG", "info"));
info!(
"Starting REVOLT server [version {}].",
crate::version::VERSION
);
thread::spawn(|| {
notifications::pubsub::launch_subscriber();
});
util::variables::preflight_checks();
database::connect().await;
notifications::hive::init_hive().await;
thread::spawn(|| {
notifications::ws::launch_server();
});
ctrlc::set_handler(move || {
// Force ungraceful exit to avoid hang.
std::process::exit(0);
})
.expect("Error setting Ctrl-C handler");
let web_task = task::spawn(launch_web());
let hive_task = task::spawn(notifications::hive::listen());
join!(
web_task,
hive_task,
notifications::websocket::launch_server()
);
}
async fn launch_web() {
let cors = rocket_cors::CorsOptions {
allowed_origins: AllowedOrigins::All,
..Default::default()
}
.to_cors()
.unwrap();
.expect("Failed to create CORS.");
let mut options = Options::new()
.base_url(format!("{}/auth", *PUBLIC_URL))
.email_verification(if *USE_EMAIL {
EmailVerification::Enabled {
success_redirect_uri: format!("{}/login", *APP_URL),
welcome_redirect_uri: format!("{}/welcome", *APP_URL),
password_reset_url: Some(format!("{}/login/reset", *APP_URL)),
verification_expiry: Duration::days(1),
password_reset_expiry: Duration::hours(1),
templates: Templates {
verify_email: Template {
title: "Verify your Revolt account.",
text: "You're almost there!
If you did not perform this action you can safely ignore this email.
Please verify your account here: {{url}}",
html: None,
},
reset_password: Template {
title: "Reset your Revolt password.",
text: "You requested a password reset, if you did not perform this action you can safely ignore this email.
Reset your password here: {{url}}",
html: None,
},
welcome: None,
},
smtp: SMTP {
from: (*SMTP_FROM).to_string(),
host: (*SMTP_HOST).to_string(),
username: (*SMTP_USERNAME).to_string(),
password: (*SMTP_PASSWORD).to_string(),
},
}
} else {
EmailVerification::Disabled
});
if *INVITE_ONLY {
options = options.invite_only_collection(database::get_collection("invites"))
}
if *USE_HCAPTCHA {
options = options.hcaptcha_secret(HCAPTCHA_KEY.clone());
}
let auth = Auth::new(database::get_collection("accounts"), options);
let mut rocket = rocket::ignite();
if *USE_PROMETHEUS {
info!("Enabled Prometheus metrics!");
let prometheus = PrometheusMetrics::new();
rocket = rocket
.attach(prometheus.clone())
.mount("/metrics", prometheus);
}
routes::mount(rocket::ignite()).attach(cors).launch();
routes::mount(rocket)
.mount("/", rocket_cors::catch_all_options_routes())
.mount("/auth", rauth::routes::routes())
.manage(auth)
.manage(cors.clone())
.attach(cors)
.launch()
.await
.unwrap();
}
src/notifications/events.rs 0 → 100644
View file @ 33262cc7
use hive_pubsub::PubSub;
use mongodb::bson::doc;
use rauth::auth::Session;
use rocket_contrib::json::JsonValue;
use serde::{Deserialize, Serialize};
use super::hive::{get_hive, subscribe_if_exists};
use crate::{database::*, util::result::Result};
#[derive(Serialize, Deserialize, Debug, Clone)]
#[serde(tag = "error")]
pub enum WebSocketError {
LabelMe,
InternalError { at: String },
InvalidSession,
OnboardingNotFinished,
AlreadyAuthenticated,
}
#[derive(Deserialize, Debug)]
#[serde(tag = "type")]
pub enum ServerboundNotification {
Authenticate(Session),
BeginTyping { channel: String },
EndTyping { channel: String },
}
#[derive(Serialize, Deserialize, Debug, Clone)]
pub enum RemoveUserField {
ProfileContent,
ProfileBackground,
StatusText,
Avatar,
}
#[derive(Serialize, Deserialize, Debug, Clone)]
pub enum RemoveChannelField {
Icon,
Description
}
#[derive(Serialize, Deserialize, Debug, Clone)]
pub enum RemoveServerField {
Icon,
Banner,
Description,
}
#[derive(Serialize, Deserialize, Debug, Clone)]
pub enum RemoveRoleField {
Colour,
}
#[derive(Serialize, Deserialize, Debug, Clone)]
pub enum RemoveMemberField {
Nickname,
Avatar,
}
#[derive(Serialize, Deserialize, Debug, Clone)]
#[serde(tag = "type")]
pub enum ClientboundNotification {
Error(WebSocketError),
Authenticated,
Ready {
users: Vec<User>,
servers: Vec<Server>,
channels: Vec<Channel>,
members: Vec<Member>
},
Message(Message),
MessageUpdate {
id: String,
channel: String,
data: JsonValue,
},
MessageDelete {
id: String,
channel: String,
},
ChannelCreate(Channel),
ChannelUpdate {
id: String,
data: JsonValue,
#[serde(skip_serializing_if = "Option::is_none")]
clear: Option<RemoveChannelField>,
},
ChannelDelete {
id: String,
},
ChannelGroupJoin {
id: String,
user: String,
},
ChannelGroupLeave {
id: String,
user: String,
},
ChannelStartTyping {
id: String,
user: String,
},
ChannelStopTyping {
id: String,
user: String,
},
ChannelAck {
id: String,
user: String,
message_id: String,
},
ServerUpdate {
id: String,
data: JsonValue,
#[serde(skip_serializing_if = "Option::is_none")]
clear: Option<RemoveServerField>,
},
ServerDelete {
id: String,
},
ServerMemberUpdate {
id: MemberCompositeKey,
data: JsonValue,
#[serde(skip_serializing_if = "Option::is_none")]
clear: Option<RemoveMemberField>,
},
ServerMemberJoin {
id: String,
user: String,
},
ServerMemberLeave {
id: String,
user: String,
},
ServerRoleUpdate {
id: String,
role_id: String,
data: JsonValue,
#[serde(skip_serializing_if = "Option::is_none")]
clear: Option<RemoveRoleField>
},
ServerRoleDelete {
id: String,
role_id: String
},
UserUpdate {
id: String,
data: JsonValue,
#[serde(skip_serializing_if = "Option::is_none")]
clear: Option<RemoveUserField>,
},
UserRelationship {
id: String,
user: User,
status: RelationshipStatus,
},
UserSettingsUpdate {
id: String,
update: JsonValue,
},
}
impl ClientboundNotification {
pub fn publish(self, topic: String) {
async_std::task::spawn(async move {
prehandle_hook(&self).await.ok(); // ! FIXME: this should be moved to pubsub
hive_pubsub::backend::mongo::publish(get_hive(), &topic, self)
.await
.ok();
});
}
pub fn publish_as_user(self, user: String) {
self.clone().publish(user.clone());
async_std::task::spawn(async move {
if let Ok(server_ids) = User::fetch_server_ids(&user).await {
for server in server_ids {
self.clone().publish(server.clone());
}
}
});
}
}
pub async fn prehandle_hook(notification: &ClientboundNotification) -> Result<()> {
match &notification {
ClientboundNotification::ChannelGroupJoin { id, user } => {
subscribe_if_exists(user.clone(), id.clone()).ok();
}
ClientboundNotification::ChannelCreate(channel) => {
let channel_id = channel.id();
match &channel {
Channel::SavedMessages { user, .. } => {
subscribe_if_exists(user.clone(), channel_id.to_string()).ok();
}
Channel::DirectMessage { recipients, .. } | Channel::Group { recipients, .. } => {
for recipient in recipients {
subscribe_if_exists(recipient.clone(), channel_id.to_string()).ok();
}
}
Channel::TextChannel { server, .. }
| Channel::VoiceChannel { server, .. } => {
// ! FIXME: write a better algorithm?
let members = Server::fetch_member_ids(server).await?;
for member in members {
subscribe_if_exists(member.clone(), channel_id.to_string()).ok();
}
}
}
}
ClientboundNotification::ServerMemberJoin { id, user } => {
let server = Ref::from_unchecked(id.clone()).fetch_server().await?;
subscribe_if_exists(user.clone(), id.clone()).ok();
for channel in server.channels {
subscribe_if_exists(user.clone(), channel).ok();
}
}
ClientboundNotification::UserRelationship { id, user, status } => {
if status != &RelationshipStatus::None {
subscribe_if_exists(id.clone(), user.id.clone()).ok();
}
}
_ => {}
}
Ok(())
}
pub async fn posthandle_hook(notification: &ClientboundNotification) {
match &notification {
ClientboundNotification::ChannelDelete { id } => {
get_hive().hive.drop_topic(&id).ok();
}
ClientboundNotification::ChannelGroupLeave { id, user } => {
get_hive().hive.unsubscribe(user, id).ok();
}
ClientboundNotification::ServerDelete { id } => {
get_hive().hive.drop_topic(&id).ok();
}
ClientboundNotification::ServerMemberLeave { id, user } => {
get_hive().hive.unsubscribe(user, id).ok();
if let Ok(server) = Ref::from_unchecked(id.clone()).fetch_server().await {
for channel in server.channels {
get_hive().hive.unsubscribe(user, &channel).ok();
}
}
}
ClientboundNotification::UserRelationship { id, user, status } => {
if status == &RelationshipStatus::None {
get_hive().hive.unsubscribe(id, &user.id).ok();
}
}
_ => {}
}
}
src/notifications/events/guilds.rs deleted 100644 → 0
View file @ cbac8029
use serde::{Deserialize, Serialize};
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct UserJoin {
pub id: String,
pub user: String,
}
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct UserLeave {
pub id: String,
pub user: String,
pub banned: bool,
}
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct ChannelCreate {
pub id: String,
pub channel: String,
pub name: String,
pub description: String,
}
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct ChannelDelete {
pub id: String,
pub channel: String,
}
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Delete {
pub id: String,
}
src/notifications/events/message.rs deleted 100644 → 0
View file @ cbac8029
use serde::{Deserialize, Serialize};
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Create {
pub id: String,
pub nonce: Option<String>,
pub channel: String,
pub author: String,
pub content: String,
}
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Edit {
pub id: String,
pub channel: String,
pub author: String,
pub content: String,
}
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Delete {
pub id: String,
}
src/notifications/events/mod.rs deleted 100644 → 0
View file @ cbac8029
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
pub mod groups;
pub mod guilds;
pub mod message;
pub mod users;
#[allow(non_camel_case_types)]
#[derive(Serialize, Deserialize, Debug, Clone)]
pub enum Notification {
message_create(message::Create),
message_edit(message::Edit),
message_delete(message::Delete),
group_user_join(groups::UserJoin),
group_user_leave(groups::UserLeave),
guild_user_join(guilds::UserJoin),
guild_user_leave(guilds::UserLeave),
guild_channel_create(guilds::ChannelCreate),
guild_channel_delete(guilds::ChannelDelete),
guild_delete(guilds::Delete),
user_friend_status(users::FriendStatus),
}
impl Notification {
pub fn serialize(self) -> String {
if let Value::Object(obj) = json!(self) {
let (key, value) = obj.iter().next().unwrap();
if let Value::Object(data) = value {
let mut data = data.clone();
data.insert("type".to_string(), Value::String(key.to_string()));
json!(data).to_string()
} else {
unreachable!()
}
} else {
unreachable!()
}
}
pub fn push_to_cache(&self) {
crate::database::channel::process_event(&self);
crate::database::guild::process_event(&self);
crate::database::user::process_event(&self);
}
}
src/notifications/events/users.rs deleted 100644 → 0
View file @ cbac8029
use serde::{Deserialize, Serialize};
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct FriendStatus {
pub id: String,
pub user: String,
pub status: i32,
}
src/notifications/hive.rs 0 → 100644
View file @ 33262cc7
use super::{events::ClientboundNotification, websocket};
use crate::database::*;
use futures::FutureExt;
use hive_pubsub::backend::mongo::MongodbPubSub;
use hive_pubsub::PubSub;
use log::{debug, error};
use once_cell::sync::OnceCell;
use serde_json::to_string;
type Hive = MongodbPubSub<String, String, ClientboundNotification>;
static HIVE: OnceCell<Hive> = OnceCell::new();
pub async fn init_hive() {
let hive = MongodbPubSub::new(
|ids, notification: ClientboundNotification| {
let notif = notification.clone();
async_std::task::spawn(async move {
super::events::posthandle_hook(&notif).await;
});
if let Ok(data) = to_string(&notification) {
debug!("Pushing out notification. {}", data);
websocket::publish(ids, notification);
} else {
error!("Failed to serialise notification.");
}
},
get_collection("pubsub"),
);
if HIVE.set(hive).is_err() {
panic!("Failed to set global pubsub instance.");
}
}
pub async fn listen() {
HIVE.get()
.unwrap()
.listen()
.fuse()
.await
.expect("Hive hit an error");
}
pub fn subscribe_multiple(user: String, topics: Vec<String>) -> Result<(), String> {
let hive = HIVE.get().unwrap();
for topic in topics {
hive.subscribe(user.clone(), topic)?;
}
Ok(())
}
pub fn subscribe_if_exists(user: String, topic: String) -> Result<(), String> {
let hive = HIVE.get().unwrap();
if hive.hive.map.lock().unwrap().get_left(&user).is_some() {
hive.subscribe(user, topic)?;
}
Ok(())
}
pub fn get_hive() -> &'static Hive {
HIVE.get().unwrap()
}
src/notifications/mod.rs
View file @ 33262cc7
use crate::database::channel::Channel;
use once_cell::sync::OnceCell;
use std::sync::mpsc::{channel, Sender};
use std::thread;
pub mod events;
pub mod pubsub;
pub mod state;
pub mod ws;
pub fn send_message<U: Into<Option<Vec<String>>>, G: Into<Option<String>>>(
users: U,
guild: G,
data: events::Notification,
) -> bool {
let users = users.into();
let guild = guild.into();
data.push_to_cache();
if pubsub::send_message(users.clone(), guild.clone(), data.clone()) {
state::send_message(users, guild, data.serialize());
true
} else {
false
}
}
struct NotificationArguments {
users: Option<Vec<String>>,
guild: Option<String>,
data: events::Notification,
}
static mut SENDER: OnceCell<Sender<NotificationArguments>> = OnceCell::new();
pub fn start_worker() {
let (sender, receiver) = channel();
unsafe {
SENDER.set(sender).unwrap();
}
thread::spawn(move || {
while let Ok(data) = receiver.recv() {
send_message(data.users, data.guild, data.data);
}
});
}
pub fn send_message_threaded<U: Into<Option<Vec<String>>>, G: Into<Option<String>>>(
users: U,
guild: G,
data: events::Notification,
) -> bool {
unsafe {
SENDER
.get()
.unwrap()
.send(NotificationArguments {
users: users.into(),
guild: guild.into(),
data,
})
.is_ok()
}
}
pub fn send_message_given_channel(data: events::Notification, channel: &Channel) {
match channel.channel_type {
0..=1 => send_message_threaded(channel.recipients.clone(), None, data),
2 => send_message_threaded(None, channel.guild.clone(), data),
_ => unreachable!(),
};
}
pub mod hive;
pub mod payload;
pub mod subscriptions;
pub mod websocket;
src/notifications/payload.rs 0 → 100644
View file @ 33262cc7
use std::collections::HashSet;
use crate::{database::*, notifications::events::ClientboundNotification};
use crate::{
database::{entities::User, get_collection},
util::result::{Error, Result},
};
use futures::StreamExt;
use mongodb::bson::{doc, from_document};
pub async fn generate_ready(mut user: User) -> Result<ClientboundNotification> {
let mut user_ids: HashSet<String> = HashSet::new();
if let Some(relationships) = &user.relations {
user_ids.extend(
relationships
.iter()
.map(|relationship| relationship.id.clone()),
);
}
let members = User::fetch_memberships(&user.id).await?;
let server_ids: Vec<String> = members.iter()
.map(|x| x.id.server.clone())
.collect();
let mut cursor = get_collection("servers")
.find(
doc! {
"_id": {
"$in": server_ids
}
},
None,
)
.await
.map_err(|_| Error::DatabaseError {
operation: "find",
with: "servers",
})?;
let mut servers = vec![];
let mut channel_ids = vec![];
while let Some(result) = cursor.next().await {
if let Ok(doc) = result {
let server: Server = from_document(doc).map_err(|_| Error::DatabaseError {
operation: "from_document",
with: "server",
})?;
channel_ids.extend(server.channels.iter().cloned());
servers.push(server);
}
}
let mut cursor = get_collection("channels")
.find(
doc! {
"$or": [
{
"_id": {
"$in": channel_ids
}
},
{
"channel_type": "SavedMessages",
"user": &user.id
},
{
"channel_type": "DirectMessage",
"recipients": &user.id
},
{
"channel_type": "Group",
"recipients": &user.id
}
]
},
None,
)
.await
.map_err(|_| Error::DatabaseError {
operation: "find",
with: "channels",
})?;
let mut channels = vec![];
while let Some(result) = cursor.next().await {
if let Ok(doc) = result {
let channel = from_document(doc).map_err(|_| Error::DatabaseError {
operation: "from_document",
with: "channel",
})?;
if let Channel::Group { recipients, .. } = &channel {
user_ids.extend(recipients.iter().cloned());
} else if let Channel::DirectMessage { recipients, .. } = &channel {
user_ids.extend(recipients.iter().cloned());
}
channels.push(channel);
}
}
user_ids.remove(&user.id);
let mut users = if user_ids.len() > 0 {
user.fetch_multiple_users(user_ids.into_iter().collect::<Vec<String>>())
.await?
} else {
vec![]
};
user.relationship = Some(RelationshipStatus::User);
user.online = Some(true);
users.push(user);
Ok(ClientboundNotification::Ready {
users,
servers,
channels,
members
})
}
src/notifications/pubsub.rs deleted 100644 → 0
View file @ cbac8029
use super::events::Notification;
use crate::database::get_collection;
use mongodb::bson::{doc, from_bson, to_bson, Bson};
use mongodb::options::{CursorType, FindOneOptions, FindOptions};
use serde::{Deserialize, Serialize};
use std::time::Duration;
use ulid::Ulid;
use once_cell::sync::OnceCell;
static SOURCEID: OnceCell<String> = OnceCell::new();
#[derive(Serialize, Deserialize, Debug)]
pub struct PubSubMessage {
#[serde(rename = "_id")]
id: String,
source: String,
user_recipients: Option<Vec<String>>,
target_guild: Option<String>,
data: Notification,
}
pub fn send_message(users: Option<Vec<String>>, guild: Option<String>, data: Notification) -> bool {
let message = PubSubMessage {
id: Ulid::new().to_string(),
source: SOURCEID.get().unwrap().to_string(),
user_recipients: users.into(),
target_guild: guild.into(),
data,
};
if get_collection("pubsub")
.insert_one(
to_bson(&message)
.expect("Failed to serialize pubsub message.")
.as_document()
.expect("Failed to convert to a document.")
.clone(),
None,
)
.is_ok()
{
true
} else {
false
}
}
pub fn launch_subscriber() {
let source = Ulid::new().to_string();
SOURCEID
.set(source.clone())
.expect("Failed to create and set source ID.");
let pubsub = get_collection("pubsub");
if let Ok(result) = pubsub.find_one(
doc! {},
FindOneOptions::builder().sort(doc! { "_id": -1 }).build(),
) {
let query = if let Some(doc) = result {
doc! { "_id": { "$gt": doc.get_str("_id").unwrap() } }
} else {
doc! {}
};
if let Ok(mut cursor) = pubsub.find(
query,
FindOptions::builder()
.cursor_type(CursorType::TailableAwait)
.no_cursor_timeout(true)
.max_await_time(Duration::from_secs(1200))
.build(),
) {
loop {
while let Some(item) = cursor.next() {
if let Ok(doc) = item {
if let Ok(message) =
from_bson(Bson::Document(doc)) as Result<PubSubMessage, _>
{
if &message.source != &source {
super::state::send_message(
message.user_recipients,
message.target_guild,
message.data.serialize(),
);
}
} else {
eprintln!("Failed to deserialize pubsub message.");
}
} else {
eprintln!("Failed to unwrap a document from pubsub.");
}
}
}
} else {
eprintln!("Failed to open subscriber cursor.");
}
} else {
eprintln!("Failed to fetch latest document from pubsub collection.");
}
}