Write websocket server.

src/guards/channel.rs

@@ -2,7 +2,7 @@ use rocket::http::{ RawStr };
 use rocket::request::{ FromParam };
 use bson::{ bson, doc, from_bson };
 
-use crate::database::{ self, user::User };
+use crate::database;
 
 use database::channel::Channel;
 use database::message::Message;

src/main.rs

@@ -2,16 +2,22 @@
 #[macro_use] extern crate rocket;
 #[macro_use] extern crate rocket_contrib;
 
+pub mod websocket;
 pub mod database;
 pub mod guards;
 pub mod routes;
 pub mod email;
 
 use dotenv;
+use std::thread;
 
 fn main() {
 	dotenv::dotenv().ok();
 	database::connect();
 
+	thread::spawn(|| {
+		websocket::launch_server();
+	});
+
 	routes::mount(rocket::ignite()).launch();
 }

src/routes/account.rs

@@ -1,7 +1,7 @@
 use crate::database;
 use crate::email;
 
-use bson::{ bson, doc, Bson::UtcDatetime, from_bson};
+use bson::{ bson, doc, Bson::UtcDatetime, from_bson };
 use rand::{ Rng, distributions::Alphanumeric };
 use rocket_contrib::json::{ Json, JsonValue };
 use serde::{ Serialize, Deserialize };

src/routes/channel.rs

@@ -11,8 +11,8 @@ use ulid::Ulid;
 #[repr(usize)]
 pub enum ChannelType {
 	DM = 0,
-	GROUP_DM = 1,
-	GUILD_CHANNEL = 2,
+	GROUPDM = 1,
+	GUILDCHANNEL = 2,
 }
 
 fn has_permission(user: &User, target: &Channel) -> bool {

src/routes/user.rs

@@ -20,8 +20,8 @@ pub fn me(user: User) -> JsonValue {
 }
 
 /// retrieve another user's information
-#[get("/<target>")]
-pub fn user(user: User, target: User) -> JsonValue {
+#[get("/<_target>")]
+pub fn user(_user: User, _target: User) -> JsonValue {
 	json!([])
 }
 
@@ -67,7 +67,7 @@ pub fn dms(user: User) -> JsonValue {
 					"type": channel::ChannelType::DM as i32
 				},
 				{
-					"type": channel::ChannelType::GROUP_DM as i32
+					"type": channel::ChannelType::GROUPDM as i32
 				}
 			],
 			"recipients": user.id
@@ -130,7 +130,7 @@ enum Relationship {
 	OUTGOING = 1,
 	INCOMING = 2,
 	BLOCKED = 3,
-	BLOCKED_OTHER = 4,
+	BLOCKEDOTHER = 4,
 	NONE = 5,
 	SELF = 6,
 }
@@ -154,8 +154,11 @@ fn get_relationship(a: &User, b: &User) -> Relationship {
 						return Relationship::OUTGOING
 					},
 					3 => {
-						return Relationship::BLOCKED_OTHER
-					}
+						return Relationship::BLOCKEDOTHER
+					},
+					4 => {
+						return Relationship::BLOCKED
+					},
 					_ => {
 						return Relationship::NONE
 					}
@@ -249,7 +252,7 @@ pub fn add_friend(user: User, target: User) -> JsonValue {
 				"success": false,
 				"error": "You have blocked this person."
 			}),
-		Relationship::BLOCKED_OTHER =>
+		Relationship::BLOCKEDOTHER =>
 			json!({
 				"success": false,
 				"error": "You have been blocked by this person."
@@ -340,7 +343,7 @@ pub fn remove_friend(user: User, target: User) -> JsonValue {
 			})
 		},
 		Relationship::BLOCKED |
-		Relationship::BLOCKED_OTHER |
+		Relationship::BLOCKEDOTHER |
 		Relationship::NONE |
 		Relationship::SELF =>
 			json!({

src/websocket/mod.rs

@@ -0,0 +1,161 @@
+extern crate ws;
+
+use crate::database;
+
+use ulid::Ulid;
+use std::sync::RwLock;
+use hashbrown::HashMap;
+
+use bson::{ bson, doc };
+use serde_json::{ Value, from_str, json };
+
+use ws::{ listen, Handler, Sender, Result, Message, Handshake, CloseCode, Error };
+
+struct Cell {
+	id: String,
+	out: Sender,
+}
+
+use once_cell::sync::OnceCell;
+static mut CLIENTS: OnceCell<RwLock<HashMap<String, Vec<Cell>>>> = OnceCell::new();
+
+struct Server {
+	out: Sender,
+	id: Option<String>,
+	internal: String,
+}
+
+impl Handler for Server {
+    fn on_open(&mut self, _: Handshake) -> Result<()> {
+		Ok(())
+    }
+
+    fn on_message(&mut self, msg: Message) -> Result<()> {
+		if let Message::Text(text) = msg {
+			let data: Value = from_str(&text).unwrap();
+
+			if let Value::String(packet_type) = &data["type"] {
+				match packet_type.as_str() {
+					"authenticate" => {
+						if let Some(_) = self.id {
+							self.out.send(
+								json!({
+									"error": "Already authenticated!"
+								})
+								.to_string()
+							)
+						} else if let Value::String(token) = &data["token"] {
+							let col = database::get_collection("users");
+							
+							match col.find_one(
+								doc! { "access_token": token },
+								None
+							).unwrap() {
+								Some(u) => {
+									let id = u.get_str("_id").expect("Missing id.");
+
+									unsafe {
+										let mut map = CLIENTS.get_mut().unwrap().write().unwrap();
+										let cell = Cell { id: self.internal.clone(), out: self.out.clone() };
+										if map.contains_key(&id.to_string()) {
+											map.get_mut(&id.to_string())
+												.unwrap()
+												.push(cell);
+										} else {
+											map.insert(id.to_string(), vec![cell]);
+										}
+									}
+
+									println!("Websocket client connected. [ID: {} // {}]", id.to_string(), self.internal);
+
+									self.id = Some(id.to_string());
+									self.out.send(
+										json!({
+											"success": true
+										})
+										.to_string()
+									)
+								},
+								None =>
+									self.out.send(
+										json!({
+											"error": "Invalid authentication token."
+										})
+										.to_string()
+									)
+							}
+						} else {
+							self.out.send(
+								json!({
+									"error": "Missing authentication token."
+								})
+								.to_string()
+							)
+						}
+					},
+					_ => Ok(())
+				}
+			} else {
+				Ok(())
+			}
+		} else {
+			Ok(())
+		}
+    }
+
+    fn on_close(&mut self, code: CloseCode, reason: &str) {
+        match code {
+            CloseCode::Normal => println!("The client is done with the connection."),
+            CloseCode::Away   => println!("The client is leaving the site."),
+            CloseCode::Abnormal => println!(
+                "Closing handshake failed! Unable to obtain closing status from client."),
+            _ => println!("The client encountered an error: {}", reason),
+        }
+
+		if let Some(id) = &self.id {
+			println!("Websocket client disconnected. [ID: {} // {}]", id, self.internal);
+			unsafe {
+				let mut map = CLIENTS.get_mut().unwrap().write().unwrap();
+				let arr = map.get_mut(&id.clone()).unwrap();
+
+				if arr.len() == 1 {
+					map.remove(&id.clone());
+				} else {
+					let index = arr.iter().position(|x| x.id == self.internal).unwrap();
+					arr.remove(index);
+					println!("User [{}] is still connected {} times", self.id.as_ref().unwrap(), arr.len());
+				}
+			}
+		}
+    }
+
+    fn on_error(&mut self, err: Error) {
+        println!("The server encountered an error: {:?}", err);
+    }
+}
+
+pub fn launch_server() {
+	unsafe {
+		if let Err(_) = CLIENTS.set(RwLock::new(HashMap::new())) {
+			panic!("Failed to set CLIENTS map!");
+		}
+	}
+
+	listen("127.0.0.1:3012", |out| { Server { out: out, id: None, internal: Ulid::new().to_string() } }).unwrap()
+}
+
+pub fn send_message(id: String, message: String) -> std::result::Result<(), ()> {
+	unsafe {
+		let map = CLIENTS.get().unwrap().read().unwrap();
+		let arr = map.get(&id).unwrap();
+
+		Ok(for item in arr {
+			match item.out.send(message.clone()) {
+				Ok(_) => (),
+				Err(_) => {
+					return Err(());
+				}
+			}
+		})
+	}
+}