forked from jmug/stoatchat
Run rust fmt.
This commit is contained in:
Paul Makles
@@ -2,138 +2,147 @@ extern crate ws;
|
||||
|
||||
use crate::database;
|
||||
|
||||
use ulid::Ulid;
|
||||
use std::sync::RwLock;
|
||||
use hashbrown::HashMap;
|
||||
use std::sync::RwLock;
|
||||
use ulid::Ulid;
|
||||
|
||||
use bson::{ bson, doc };
|
||||
use serde_json::{ Value, from_str, json };
|
||||
use bson::{bson, doc};
|
||||
use serde_json::{from_str, json, Value};
|
||||
|
||||
use ws::{ listen, Handler, Sender, Result, Message, Handshake, CloseCode, Error };
|
||||
use ws::{listen, CloseCode, Error, Handler, Handshake, Message, Result, Sender};
|
||||
|
||||
struct Cell {
|
||||
id: String,
|
||||
out: Sender,
|
||||
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,
|
||||
out: Sender,
|
||||
id: Option<String>,
|
||||
internal: String,
|
||||
}
|
||||
|
||||
impl Handler for Server {
|
||||
fn on_open(&mut self, _: Handshake) -> Result<()> {
|
||||
Ok(())
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn on_message(&mut self, msg: Message) -> Result<()> {
|
||||
if let Message::Text(text) = msg {
|
||||
let data: Value = from_str(&text).unwrap();
|
||||
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!({
|
||||
"type": "authenticate",
|
||||
"success": false,
|
||||
"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.");
|
||||
if let Value::String(packet_type) = &data["type"] {
|
||||
match packet_type.as_str() {
|
||||
"authenticate" => {
|
||||
if let Some(_) = self.id {
|
||||
self.out.send(
|
||||
json!({
|
||||
"type": "authenticate",
|
||||
"success": false,
|
||||
"error": "Already authenticated!"
|
||||
})
|
||||
.to_string(),
|
||||
)
|
||||
} else if let Value::String(token) = &data["token"] {
|
||||
let col = database::get_collection("users");
|
||||
|
||||
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]);
|
||||
}
|
||||
}
|
||||
match col.find_one(doc! { "access_token": token }, None).unwrap() {
|
||||
Some(u) => {
|
||||
let id = u.get_str("_id").expect("Missing id.");
|
||||
|
||||
println!("Websocket client connected. [ID: {} // {}]", id.to_string(), self.internal);
|
||||
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]);
|
||||
}
|
||||
}
|
||||
|
||||
self.id = Some(id.to_string());
|
||||
self.out.send(
|
||||
json!({
|
||||
"type": "authenticate",
|
||||
"success": true
|
||||
})
|
||||
.to_string()
|
||||
)
|
||||
},
|
||||
None =>
|
||||
self.out.send(
|
||||
json!({
|
||||
"type": "authenticate",
|
||||
"success": false,
|
||||
"error": "Invalid authentication token."
|
||||
})
|
||||
.to_string()
|
||||
)
|
||||
}
|
||||
} else {
|
||||
self.out.send(
|
||||
json!({
|
||||
"type": "authenticate",
|
||||
"success": false,
|
||||
"error": "Missing authentication token."
|
||||
})
|
||||
.to_string()
|
||||
)
|
||||
}
|
||||
},
|
||||
_ => Ok(())
|
||||
}
|
||||
} else {
|
||||
Ok(())
|
||||
}
|
||||
} else {
|
||||
Ok(())
|
||||
}
|
||||
println!(
|
||||
"Websocket client connected. [ID: {} // {}]",
|
||||
id.to_string(),
|
||||
self.internal
|
||||
);
|
||||
|
||||
self.id = Some(id.to_string());
|
||||
self.out.send(
|
||||
json!({
|
||||
"type": "authenticate",
|
||||
"success": true
|
||||
})
|
||||
.to_string(),
|
||||
)
|
||||
}
|
||||
None => self.out.send(
|
||||
json!({
|
||||
"type": "authenticate",
|
||||
"success": false,
|
||||
"error": "Invalid authentication token."
|
||||
})
|
||||
.to_string(),
|
||||
),
|
||||
}
|
||||
} else {
|
||||
self.out.send(
|
||||
json!({
|
||||
"type": "authenticate",
|
||||
"success": false,
|
||||
"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."),
|
||||
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 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());
|
||||
}
|
||||
}
|
||||
}
|
||||
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) {
|
||||
@@ -142,37 +151,43 @@ impl Handler for Server {
|
||||
}
|
||||
|
||||
pub fn launch_server() {
|
||||
unsafe {
|
||||
if let Err(_) = CLIENTS.set(RwLock::new(HashMap::new())) {
|
||||
panic!("Failed to set CLIENTS map!");
|
||||
}
|
||||
}
|
||||
unsafe {
|
||||
if let Err(_) = CLIENTS.set(RwLock::new(HashMap::new())) {
|
||||
panic!("Failed to set CLIENTS map!");
|
||||
}
|
||||
}
|
||||
|
||||
listen("192.168.0.10:9999", |out| { Server { out: out, id: None, internal: Ulid::new().to_string() } }).unwrap()
|
||||
listen("192.168.0.10:9999", |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();
|
||||
if map.contains_key(&id) {
|
||||
let arr = map.get(&id).unwrap();
|
||||
unsafe {
|
||||
let map = CLIENTS.get().unwrap().read().unwrap();
|
||||
if map.contains_key(&id) {
|
||||
let arr = map.get(&id).unwrap();
|
||||
|
||||
for item in arr {
|
||||
if let Err(_) = item.out.send(message.clone()) {
|
||||
return Err(());
|
||||
}
|
||||
}
|
||||
}
|
||||
for item in arr {
|
||||
if let Err(_) = item.out.send(message.clone()) {
|
||||
return Err(());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
// ! TODO: WRITE THREADED QUEUE SYSTEM
|
||||
// ! FETCH RECIPIENTS HERE INSTEAD OF IN METHOD
|
||||
|
||||
pub fn queue_message(ids: Vec<String>, message: String) {
|
||||
for id in ids {
|
||||
send_message(id, message.clone()).expect("uhhhhhhhhhh can i get uhhhhhhhhhhhhhhhhhh mcdonald cheese burger with fries");
|
||||
}
|
||||
for id in ids {
|
||||
send_message(id, message.clone())
|
||||
.expect("uhhhhhhhhhh can i get uhhhhhhhhhhhhhhhhhh mcdonald cheese burger with fries");
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user