Re-write notifications system.

src/notifications/events/message.rs

@@ -0,0 +1,10 @@
+use serde::{Deserialize, Serialize};
+
+#[derive(Serialize, Deserialize, Debug)]
+pub struct Create {
+    pub id: String,
+    pub nonce: Option<String>,
+    pub channel: String,
+    pub author: String,
+    pub content: String,
+}

src/notifications/events/mod.rs

@@ -0,0 +1,8 @@
+use serde::{Deserialize, Serialize};
+
+pub mod message;
+
+#[derive(Serialize, Deserialize, Debug)]
+pub enum Notification {
+    MessageCreate(message::Create),
+}

src/notifications/mod.rs

@@ -0,0 +1,21 @@
+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();
+
+    if pubsub::send_message(users.clone(), guild.clone(), data) {
+        state::send_message(users, guild, "bruh".to_string());
+
+        true
+    } else {
+        false
+    }
+}

src/notifications/pubsub.rs

@@ -0,0 +1,112 @@
+use super::events::Notification;
+use crate::database::get_collection;
+
+use 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>,
+
+    notification_type: 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(),
+        notification_type: match data {
+            Notification::MessageCreate(_) => "message_create",
+        }
+        .to_string(),
+        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,
+                                    json!(message.data).to_string(),
+                                );
+                            }
+                        } 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.");
+    }
+}

src/notifications/state.rs

@@ -0,0 +1,210 @@
+use super::events::Notification;
+use crate::database;
+use crate::util::vec_to_set;
+
+use bson::doc;
+use hashbrown::{HashMap, HashSet};
+use mongodb::options::FindOneOptions;
+use once_cell::sync::OnceCell;
+use std::sync::RwLock;
+use ws::Sender;
+
+pub enum StateResult {
+    DatabaseError,
+    InvalidToken,
+    Success(String),
+}
+
+static mut CONNECTIONS: OnceCell<RwLock<HashMap<String, Sender>>> = OnceCell::new();
+
+pub fn add_connection(id: String, sender: Sender) {
+    unsafe {
+        CONNECTIONS
+            .get()
+            .unwrap()
+            .write()
+            .unwrap()
+            .insert(id, sender);
+    }
+}
+
+pub struct User {
+    connections: HashSet<String>,
+    guilds: HashSet<String>,
+}
+
+impl User {
+    pub fn new() -> User {
+        User {
+            connections: HashSet::new(),
+            guilds: HashSet::new(),
+        }
+    }
+}
+
+pub struct Guild {
+    users: HashSet<String>,
+}
+
+impl Guild {
+    pub fn new() -> Guild {
+        Guild {
+            users: HashSet::new(),
+        }
+    }
+}
+
+pub struct GlobalState {
+    users: HashMap<String, User>,
+    guilds: HashMap<String, Guild>,
+}
+
+impl GlobalState {
+    pub fn new() -> GlobalState {
+        GlobalState {
+            users: HashMap::new(),
+            guilds: HashMap::new(),
+        }
+    }
+
+    pub fn push_to_guild(&mut self, guild: String, user: String) {
+        if !self.guilds.contains_key(&guild) {
+            self.guilds.insert(guild.clone(), Guild::new());
+        }
+
+        self.guilds.get_mut(&guild).unwrap().users.insert(user);
+    }
+
+    pub fn try_authenticate(&mut self, connection: String, access_token: String) -> StateResult {
+        if let Ok(result) = database::get_collection("users").find_one(
+            doc! {
+                "access_token": access_token,
+            },
+            FindOneOptions::builder()
+                .projection(doc! { "_id": 1 })
+                .build(),
+        ) {
+            if let Some(user) = result {
+                let user_id = user.get_str("_id").unwrap();
+
+                if self.users.contains_key(user_id) {
+                    self.users
+                        .get_mut(user_id)
+                        .unwrap()
+                        .connections
+                        .insert(connection);
+
+                    return StateResult::Success(user_id.to_string());
+                }
+
+                if let Ok(results) =
+                    database::get_collection("members").find(doc! { "_id.user": &user_id }, None)
+                {
+                    let mut guilds = vec![];
+                    for result in results {
+                        if let Ok(entry) = result {
+                            guilds.push(
+                                entry
+                                    .get_document("_id")
+                                    .unwrap()
+                                    .get_str("guild")
+                                    .unwrap()
+                                    .to_string(),
+                            );
+                        }
+                    }
+
+                    let mut user = User::new();
+                    for guild in guilds {
+                        user.guilds.insert(guild.clone());
+                        self.push_to_guild(guild, user_id.to_string());
+                    }
+
+                    user.connections.insert(connection);
+                    self.users.insert(user_id.to_string(), user);
+
+                    StateResult::Success(user_id.to_string())
+                } else {
+                    StateResult::DatabaseError
+                }
+            } else {
+                StateResult::InvalidToken
+            }
+        } else {
+            StateResult::DatabaseError
+        }
+    }
+
+    pub fn disconnect<U: Into<Option<String>>>(&mut self, user_id: U, connection: String) {
+        if let Some(user_id) = user_id.into() {
+            let user = self.users.get_mut(&user_id).unwrap();
+            user.connections.remove(&connection);
+
+            if user.connections.len() == 0 {
+                for guild in &user.guilds {
+                    self.guilds.get_mut(guild).unwrap().users.remove(&user_id);
+                }
+
+                self.users.remove(&user_id);
+            }
+        }
+
+        unsafe {
+            CONNECTIONS
+                .get()
+                .unwrap()
+                .write()
+                .unwrap()
+                .remove(&connection);
+        }
+    }
+}
+
+pub static mut DATA: OnceCell<RwLock<GlobalState>> = OnceCell::new();
+
+pub fn init() {
+    unsafe {
+        if CONNECTIONS.set(RwLock::new(HashMap::new())).is_err() {
+            panic!("Failed to set global connections map.");
+        }
+
+        if DATA.set(RwLock::new(GlobalState::new())).is_err() {
+            panic!("Failed to set global state.");
+        }
+    }
+}
+
+pub fn send_message(
+    users: Option<Vec<String>>,
+    guild: Option<String>,
+    data: String,
+) {
+    let state = unsafe { DATA.get().unwrap().read().unwrap() };
+    let mut connections = HashSet::new();
+    
+    let mut users = vec_to_set(&users.unwrap_or(vec![]));
+    if let Some(guild) = guild {
+        if let Some(entry) = state.guilds.get(&guild) {
+            for user in &entry.users {
+                users.insert(user.to_string());
+            }
+        }
+    }
+
+    for user in users {
+        if let Some(entry) = state.users.get(&user) {
+            for connection in &entry.connections {
+                connections.insert(connection.clone());
+            }
+        }
+    }
+
+    let targets = unsafe { CONNECTIONS.get().unwrap().read().unwrap() };
+    for conn in connections {
+        if let Some(sender) = targets.get(&conn) {
+            if sender.send(data.clone()).is_err() {
+                eprintln!("Failed to send a notification to a websocket. [{}]", &conn);
+            }
+        }
+    }
+}

src/notifications/ws.rs

@@ -0,0 +1,119 @@
+use super::state::{self, StateResult};
+
+use serde_json::{from_str, json, Value};
+use std::env;
+use ulid::Ulid;
+use ws::{listen, CloseCode, Error, Handler, Handshake, Message, Result, Sender};
+
+struct Server {
+    sender: Sender,
+    user_id: Option<String>,
+    id: String,
+}
+
+impl Handler for Server {
+    fn on_open(&mut self, _: Handshake) -> Result<()> {
+        state::add_connection(self.id.clone(), self.sender.clone());
+        Ok(())
+    }
+
+    fn on_message(&mut self, msg: Message) -> Result<()> {
+        if let Message::Text(text) = msg {
+            if let Ok(data) = from_str(&text) as std::result::Result<Value, _> {
+                if let Value::String(packet_type) = &data["type"] {
+                    if packet_type == "authenticate" {
+                        if self.user_id.is_some() {
+                            self.sender.send(
+                                json!({
+                                    "type": "authenticate",
+                                    "success": false,
+                                    "error": "Already authenticated!"
+                                })
+                                .to_string(),
+                            )
+                        } else if let Value::String(token) = &data["token"] {
+                            let mut state = unsafe { state::DATA.get().unwrap().write().unwrap() };
+
+                            match state.try_authenticate(self.id.clone(), token.to_string()) {
+                                StateResult::Success(user_id) => {
+                                    self.user_id = Some(user_id);
+                                    self.sender.send(
+                                        json!({
+                                            "type": "authenticate",
+                                            "success": true,
+                                        })
+                                        .to_string(),
+                                    )
+                                }
+                                StateResult::DatabaseError => self.sender.send(
+                                    json!({
+                                        "type": "authenticate",
+                                        "success": false,
+                                        "error": "Had database error."
+                                    })
+                                    .to_string(),
+                                ),
+                                StateResult::InvalidToken => self.sender.send(
+                                    json!({
+                                        "type": "authenticate",
+                                        "success": false,
+                                        "error": "Invalid token."
+                                    })
+                                    .to_string(),
+                                ),
+                            }
+                        } else {
+                            self.sender.send(
+                                json!({
+                                    "type": "authenticate",
+                                    "success": false,
+                                    "error": "Token not present."
+                                })
+                                .to_string(),
+                            )
+                        }
+                    } else {
+                        Ok(())
+                    }
+                } else {
+                    Ok(())
+                }
+            } else {
+                Ok(())
+            }
+        } else {
+            Ok(())
+        }
+    }
+
+    fn on_close(&mut self, _code: CloseCode, _reason: &str) {
+        unsafe {
+            state::DATA
+                .get()
+                .unwrap()
+                .write()
+                .unwrap()
+                .disconnect(self.user_id.clone(), self.id.clone());
+        }
+
+        println!("User disconnected. [{}]", self.id);
+    }
+
+    fn on_error(&mut self, err: Error) {
+        println!("The server encountered an error: {:?}", err);
+    }
+}
+
+pub fn launch_server() {
+    state::init();
+
+    listen(
+        env::var("WS_HOST").unwrap_or("0.0.0.0:9999".to_string()),
+        |sender| Server {
+            sender,
+            user_id: None,
+            id: Ulid::new().to_string(),
+        },
+    )
+    .unwrap()
+}