jmug/stoatchat
1
0
Fork 1
Code Issues Packages Projects Releases Wiki Activity

Switch to async-std, refractor notifications code.

Browse Source
This commit is contained in:
Paul Makles
2020-12-28 17:52:22 +00:00
parent 6253a91276
commit c748b4349b
16 changed files with 498 additions and 809 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
src/main.rs
View File

@@ -6,21 +6,25 @@ extern crate rocket;
#[macro_use]
extern crate rocket_contrib;
#[macro_use]
extern crate bitfield;
#[macro_use]
extern crate lazy_static;
extern crate ctrlc;
pub mod notifications;
pub mod database;
pub mod pubsub;
pub mod routes;
pub mod util;
use rauth;
use log::info;
use futures::join;
use async_std::task;
use rocket_cors::AllowedOrigins;
#[tokio::main]
async fn main() {
fn main() {
task::block_on(entry())
}
async fn entry() {
dotenv::dotenv().ok();
env_logger::init_from_env(env_logger::Env::default().filter_or("RUST_LOG", "info"));
@@ -28,10 +32,16 @@ async fn main() {
util::variables::preflight_checks();
database::connect().await;
ctrlc::set_handler(move || {
// Force ungraceful exit to avoid hang.
std::process::exit(0);
}).expect("Error setting Ctrl-C handler");
pubsub::hive::init_hive();
//pubsub::websocket::launch_server();
join!(launch_web(), notifications::websocket::launch_server());
}
async fn launch_web() {
let cors = rocket_cors::CorsOptions {
allowed_origins: AllowedOrigins::All,
..Default::default()

67
src/notifications/events.rs Normal file
View File

@@ -0,0 +1,67 @@
use serde::{Deserialize, Serialize};
#[derive(Serialize, Deserialize, Debug, Clone)]
#[serde(tag = "type")]
pub enum Notification {
MessageCreate {
id: String,
nonce: Option<String>,
channel: String,
author: String,
content: String,
},
MessageEdit {
id: String,
channel: String,
author: String,
content: String,
},
MessageDelete {
id: String,
},
GroupUserJoin {
id: String,
user: String,
},
GroupUserLeave {
id: String,
user: String,
},
GuildUserJoin {
id: String,
user: String,
},
GuildUserLeave {
id: String,
user: String,
banned: bool,
},
GuildChannelCreate {
id: String,
channel: String,
name: String,
description: String,
},
GuildChannelDelete {
id: String,
channel: String,
},
GuildDelete {
id: String,
},
UserRelationship {
id: String,
user: String,
status: i32,
}
}

16
src/pubsub/hive.rs → src/notifications/hive.rs
View File

@@ -10,14 +10,12 @@ use log::{error, debug};
static HIVE: OnceCell<MongodbPubSub<String, String, Notification>> = OnceCell::new();
pub fn init_hive() {
pub async fn init_hive() {
let hive = MongodbPubSub::new(
|_ids, notification| {
if let Ok(data) = to_string(&notification) {
debug!("Pushing out notification. {}", data);
// if let Err(err) = websocket::publish(ids, data) {
// error!("Failed to publish notification through WebSocket! {}", err);
// }
// ! FIXME: push to websocket
} else {
error!("Failed to serialise notification.");
}
@@ -25,20 +23,18 @@ pub fn init_hive() {
get_collection("hive"),
);
//listen_thread(hive.clone());
if HIVE.set(hive).is_err() {
panic!("Failed to set global pubsub instance.");
}
}
pub fn publish(topic: &String, data: Notification) -> Result<(), String> {
let hive = HIVE.get().expect("Global pubsub instance not available.");
let hive = HIVE.get().unwrap();
hive.publish(topic, data)
}
pub fn subscribe(user: String, topics: Vec<String>) -> Result<(), String> {
let hive = HIVE.get().expect("Global pubsub instance not available.");
let hive = HIVE.get().unwrap();
for topic in topics {
hive.subscribe(user.clone(), topic)?;
}
@@ -47,14 +43,14 @@ pub fn subscribe(user: String, topics: Vec<String>) -> Result<(), String> {
}
pub fn drop_user(user: &String) -> Result<(), String> {
let hive = HIVE.get().expect("Global pubsub instance not available.");
let hive = HIVE.get().unwrap();
hive.drop_client(user)?;
Ok(())
}
pub fn drop_topic(topic: &String) -> Result<(), String> {
let hive = HIVE.get().expect("Global pubsub instance not available.");
let hive = HIVE.get().unwrap();
hive.drop_topic(topic)?;
Ok(())

2
src/pubsub/mod.rs → src/notifications/mod.rs
View File

@@ -1,3 +1,3 @@
pub mod websocket;
pub mod events;
pub mod hive;
// pub mod websocket;

28
src/notifications/websocket.rs Normal file
View File

@@ -0,0 +1,28 @@
use crate::util::variables::WS_HOST;
use log::info;
use async_std::task;
use futures::prelude::*;
use async_std::net::{TcpListener, TcpStream};
pub async fn launch_server() {
let try_socket = TcpListener::bind(WS_HOST.to_string()).await;
let listener = try_socket.expect("Failed to bind");
info!("Listening on: {}", *WS_HOST);
while let Ok((stream, _)) = listener.accept().await {
task::spawn(accept(stream));
}
}
async fn accept(stream: TcpStream) {
let addr = stream.peer_addr().expect("Connected streams should have a peer address.");
let ws_stream = async_tungstenite::accept_async(stream)
.await
.expect("Error during websocket handshake.");
info!("User established WebSocket connection from {}.", addr);
let (write, read) = ws_stream.split();
read.forward(write).await.expect("Failed to forward message")
}

13
src/pubsub/events/groups.rs
View File

@@ -1,13 +0,0 @@
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,
}

33
src/pubsub/events/guilds.rs
View File

@@ -1,33 +0,0 @@
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,
}

23
src/pubsub/events/message.rs
View File

@@ -1,23 +0,0 @@
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,
}

31
src/pubsub/events/mod.rs
View File

@@ -1,31 +0,0 @@
use serde::{Deserialize, Serialize};
pub mod groups;
pub mod guilds;
pub mod message;
pub mod users;
#[allow(non_camel_case_types)]
#[derive(Serialize, Deserialize, Debug, Clone)]
#[serde(tag = "type")]
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 push_to_cache(&self) {
//crate::database::channel::process_event(&self);
//crate::database::guild::process_event(&self);
//crate::database::user::process_event(&self);
}
}

8
src/pubsub/events/users.rs
View File

@@ -1,8 +0,0 @@
use serde::{Deserialize, Serialize};
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct FriendStatus {
pub id: String,
pub user: String,
pub status: i32,
}

200
src/pubsub/websocket.disabled/client.rs
View File

@@ -1,200 +0,0 @@
use super::state;
use crate::database::get_collection;
use crate::pubsub::hive;
use log::{error, info};
use mongodb::bson::doc;
use mongodb::options::FindOneOptions;
use serde_json::{from_str, json, Value};
use ulid::Ulid;
use ws::{CloseCode, Error, Handler, Handshake, Message, Result, Sender};
pub struct Client {
id: String,
sender: Sender,
user_id: Option<String>,
}
impl Client {
pub fn new(sender: Sender) -> Client {
Client {
id: Ulid::new().to_string(),
user_id: None,
sender,
}
}
}
impl Handler for Client {
fn on_open(&mut self, handshake: Handshake) -> Result<()> {
info!("Client connected. [{}] {:?}", self.id, handshake.peer_addr);
Ok(())
}
// Client sends { "type": "authenticate", "token": token }.
// Receives { "type": "authorised" } and waits.
// Client then receives { "type": "ready", "data": payload }.
// If at any point we hit an error, send { "type": "error", "error": error }.
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() {
return self.sender.send(
json!({
"type": "error",
"error": "Already authenticated!"
})
.to_string(),
);
} else if let Value::String(token) = &data["token"] {
let user = get_collection("users").find_one(
doc! {
"access_token": token
},
FindOneOptions::builder()
.projection(doc! { "_id": 1 })
.build(),
).await;
if let Ok(result) = user {
if let Some(doc) = result {
self.sender.send(
json!({
"type": "authorised"
})
.to_string(),
)?;
// FIXME: fetch above when we switch to new token system
// or auth cache system, something like that
let user = crate::database::user::fetch_user(
doc.get_str("_id").unwrap(),
)
.unwrap()
.unwrap(); // this should be guranteed, I think, maybe? I'm getting rid of it later. FIXME
self.user_id = Some(user.id.clone());
match user.create_payload() {
Ok(payload) => {
// ! Grab the ids from the payload,
// ! there's probably a better way to
// ! do this. I'll rewrite it at some point.
let mut ids = vec![
self.user_id.as_ref().unwrap().clone()
];
{
// This is bad code. But to be fair
// it should work just fine.
for user in payload.get("users").unwrap().as_array().unwrap() {
ids.push(user.as_object().unwrap().get("id").unwrap().as_str().unwrap().to_string());
}
for channel in payload.get("channels").unwrap().as_array().unwrap() {
ids.push(channel.as_object().unwrap().get("id").unwrap().as_str().unwrap().to_string());
}
for guild in payload.get("guilds").unwrap().as_array().unwrap() {
ids.push(guild.as_object().unwrap().get("id").unwrap().as_str().unwrap().to_string());
}
}
if let Err(err) = hive::subscribe(self.user_id.as_ref().unwrap().clone(), ids) {
error!("Failed to subscribe someone to the Hive! {}", err);
self.sender.send(
json!({
"type": "warn",
"error": "Failed to subscribe you to the Hive. You may not receive all notifications."
})
.to_string(),
)?;
}
self.sender.send(
json!({
"type": "ready",
"data": payload
})
.to_string(),
)?;
if state::accept(
self.id.clone(),
self.user_id.as_ref().unwrap().clone(),
self.sender.clone(),
)
.is_err()
{
self.sender.send(
json!({
"type": "warn",
"error": "Failed to accept your connection. You will not receive any notifications."
})
.to_string(),
)?;
}
}
Err(error) => {
error!("Failed to create payload! {}", error);
self.sender.send(
json!({
"type": "error",
"error": "Failed to create payload."
})
.to_string(),
)?;
}
}
} else {
self.sender.send(
json!({
"type": "error",
"error": "Invalid token."
})
.to_string(),
)?;
}
} else {
self.sender.send(
json!({
"type": "error",
"error": "Failed to fetch from database."
})
.to_string(),
)?;
}
} else {
self.sender.send(
json!({
"type": "error",
"error": "Missing token."
})
.to_string(),
)?;
}
}
}
}
}
Ok(())
}
fn on_close(&mut self, _code: CloseCode, reason: &str) {
info!("Client disconnected. [{}] {}", self.id, reason);
if let Err(error) = state::drop(&self.id) {
error!("Also failed to drop client from state! {}", error);
}
}
fn on_error(&mut self, err: Error) {
error!(
"A client disconnected due to an error. [{}] {}",
self.id, err
);
}
}

23
src/pubsub/websocket.disabled/mod.rs
View File

@@ -1,23 +0,0 @@
use crate::util::variables::WS_HOST;
use log::{error, info};
use std::thread;
use ws::listen;
mod client;
mod state;
pub use state::publish;
pub fn launch_server() {
thread::spawn(|| {
if listen(WS_HOST.to_string(), |sender| client::Client::new(sender)).is_err() {
error!(
"Failed to listen for WebSocket connections on {:?}!",
*WS_HOST
);
} else {
info!("Listening for WebSocket connections on {:?}", *WS_HOST);
}
});
}

89
src/pubsub/websocket.disabled/state.rs
View File

@@ -1,89 +0,0 @@
use crate::pubsub::hive;
use many_to_many::ManyToMany;
use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use log::{error, info};
use ws::Sender;
lazy_static! {
static ref CONNECTIONS: Arc<RwLock<HashMap<String, Sender>>> =
Arc::new(RwLock::new(HashMap::new()));
static ref CLIENTS: Arc<RwLock<ManyToMany<String, String>>> =
Arc::new(RwLock::new(ManyToMany::new()));
}
pub fn accept(id: String, user_id: String, sender: Sender) -> Result<(), String> {
let mut conns = CONNECTIONS
.write()
.map_err(|_| "Failed to lock connections for writing.")?;
conns.insert(id.clone(), sender);
let mut clients = CLIENTS
.write()
.map_err(|_| "Failed to lock clients for writing.")?;
clients.insert(user_id.clone(), id.clone());
info!("Accepted user [{}] for connection {}.", user_id, id);
Ok(())
}
pub fn drop(id: &String) -> Result<(), String> {
let mut conns = CONNECTIONS
.write()
.map_err(|_| "Failed to lock connections for writing.")?;
conns.remove(id);
let mut clients = CLIENTS
.write()
.map_err(|_| "Failed to lock clients for writing.")?;
let uid = if let Some(ids) = clients.get_right(id) {
let user_id: String = ids.into_iter().next().unwrap();
info!("Dropped user [{}] for connection {}.", user_id, id);
Some(user_id)
} else {
None
};
clients.remove_right(id);
if let Some(user_id) = &uid {
if let None = clients.get_left(user_id) {
if let Err(error) = hive::drop_user(user_id) {
error!("Failed to drop user from hive! {}", error);
} else {
info!("User [{}] has completed disconnected from node.", user_id);
}
}
}
Ok(())
}
pub fn publish(clients: Vec<String>, data: String) -> Result<(), String> {
let conns = CONNECTIONS
.read()
.map_err(|_| "Failed to lock connections for reading.")?;
let client_map = CLIENTS
.read()
.map_err(|_| "Failed to lock clients for reading.")?;
for client in clients {
if let Some(targets) = client_map.get_left(&client) {
for target in &targets {
if let Some(connection) = conns.get(target) {
if let Err(err) = connection.send(data.clone()) {
error!("Failed to publish notification to client [{}]! {}", target, err);
}
}
}
}
}
Ok(())
}

344
src/routes/account.disabled.rs
View File

@@ -1,344 +0,0 @@
use super::Response;
use crate::database;
use crate::util::variables::{DISABLE_REGISTRATION, USE_EMAIL};
use crate::util::{captcha, email, gen_token};
use bcrypt::{hash, verify};
use chrono::prelude::*;
use database::user::User;
use log::error;
use mongodb::bson::{doc, from_bson, Bson};
use rocket_contrib::json::Json;
use serde::{Deserialize, Serialize};
use ulid::Ulid;
use validator::validate_email;
#[derive(Serialize, Deserialize)]
pub struct Create {
username: String,
password: String,
email: String,
captcha: Option<String>,
}
/// create a new Revolt account
/// (1) validate input
/// [username] 2 to 32 characters
/// [password] 8 to 72 characters
/// [email] validate against RFC
/// (2) check email existence
/// (3) add user and send email verification
#[post("/create", data = "<info>")]
pub async fn create(info: Json<Create>) -> Response {
if let Err(error) = captcha::verify(&info.captcha).await {
return Response::BadRequest(json!({ "error": error }));
}
if *DISABLE_REGISTRATION {
return Response::BadRequest(json!({ "error": "Registration disabled." }));
}
let col = database::get_collection("users");
if info.username.len() < 2 || info.username.len() > 32 {
return Response::NotAcceptable(
json!({ "error": "Username needs to be at least 2 chars and less than 32 chars." }),
);
}
if info.password.len() < 8 || info.password.len() > 72 {
return Response::NotAcceptable(
json!({ "error": "Password needs to be at least 8 chars and at most 72." }),
);
}
if !validate_email(info.email.clone()) {
return Response::UnprocessableEntity(json!({ "error": "Invalid email." }));
}
if let Some(_) = col
.find_one(doc! { "email": info.email.clone() }, None)
.await
.expect("Failed user lookup")
{
return Response::Conflict(json!({ "error": "Email already in use!" }));
}
if let Some(_) = col
.find_one(doc! { "username": info.username.clone() }, None)
.await
.expect("Failed user lookup")
{
return Response::Conflict(json!({ "error": "Username already in use!" }));
}
if let Ok(hashed) = hash(info.password.clone(), 10) {
let access_token = gen_token(92);
let code = gen_token(48);
let email_verification = match *USE_EMAIL {
true => doc! {
"verified": false,
"target": info.email.clone(),
"expiry": Bson::DateTime(Utc::now() + chrono::Duration::days(1)),
"rate_limit": Bson::DateTime(Utc::now() + chrono::Duration::minutes(1)),
"code": code.clone(),
},
false => doc! {
"verified": true
},
};
let id = Ulid::new().to_string();
match col.insert_one(
doc! {
"_id": &id,
"email": info.email.clone(),
"username": info.username.clone(),
"display_name": info.username.clone(),
"password": hashed,
"access_token": &access_token,
"email_verification": email_verification
},
None,
)
.await {
Ok(_) => {
if *USE_EMAIL {
let sent = email::send_verification_email(info.email.clone(), code);
Response::Success(json!({
"email_sent": sent,
}))
} else {
Response::Success(json!({
"id": id,
"access_token": access_token
}))
}
}
Err(_) => {
Response::InternalServerError(json!({ "error": "Failed to create account." }))
}
}
} else {
Response::InternalServerError(json!({ "error": "Failed to hash." }))
}
}
/// verify an email for a Revolt account
/// (1) check if code is valid
/// (2) check if it expired yet
/// (3) set account as verified
#[get("/verify/<code>")]
pub async fn verify_email(code: String) -> Response {
let col = database::get_collection("users");
if let Some(u) = col
.find_one(doc! { "email_verification.code": code.clone() }, None)
.await
.expect("Failed user lookup")
{
let user: User = from_bson(Bson::Document(u)).expect("Failed to unwrap user.");
let ev = user.email_verification;
if Utc::now() > *ev.expiry.unwrap() {
Response::Gone(json!({
"success": false,
"error": "Token has expired!",
}))
} else {
let target = ev.target.unwrap();
col.update_one(
doc! { "_id": user.id },
doc! {
"$unset": {
"email_verification.code": "",
"email_verification.expiry": "",
"email_verification.target": "",
"email_verification.rate_limit": "",
},
"$set": {
"email_verification.verified": true,
"email": target.clone(),
},
},
None,
)
.await
.expect("Failed to update user!");
if *USE_EMAIL {
if let Err(err) = email::send_welcome_email(target.to_string(), user.username) {
error!("Failed to send welcome email! {}", err);
}
}
Response::Redirect(super::Redirect::to("https://app.revolt.chat"))
}
} else {
Response::BadRequest(json!({ "error": "Invalid code." }))
}
}
#[derive(Serialize, Deserialize)]
pub struct Resend {
email: String,
captcha: Option<String>,
}
/// resend a verification email
/// (1) check if verification is pending for x email
/// (2) check for rate limit
/// (3) resend the email
#[post("/resend", data = "<info>")]
pub async fn resend_email(info: Json<Resend>) -> Response {
if let Err(error) = captcha::verify(&info.captcha).await {
return Response::BadRequest(json!({ "error": error }));
}
let col = database::get_collection("users");
if let Some(u) = col
.find_one(
doc! { "email_verification.target": info.email.clone() },
None,
)
.await
.expect("Failed user lookup.")
{
let user: User = from_bson(Bson::Document(u)).expect("Failed to unwrap user.");
let ev = user.email_verification;
let expiry = ev.expiry.unwrap();
let rate_limit = ev.rate_limit.unwrap();
if Utc::now() < *rate_limit {
Response::TooManyRequests(
json!({ "error": "You are being rate limited, please try again in a while." }),
)
} else {
let mut new_expiry = Bson::DateTime(Utc::now() + chrono::Duration::days(1));
if info.email.clone() != user.email {
if Utc::now() > *expiry {
return Response::Gone(
json!({ "error": "To help protect your account, please login and change your email again. The original request was made over one day ago." }),
);
}
new_expiry = Bson::DateTime(*expiry);
}
let code = gen_token(48);
col.update_one(
doc! { "_id": user.id },
doc! {
"$set": {
"email_verification.code": code.clone(),
"email_verification.expiry": new_expiry,
"email_verification.rate_limit": Bson::DateTime(Utc::now() + chrono::Duration::minutes(1)),
},
},
None,
)
.await
.expect("Failed to update user!");
if let Err(err) = email::send_verification_email(info.email.clone(), code) {
return Response::InternalServerError(json!({ "error": err }));
}
Response::Result(super::Status::Ok)
}
} else {
Response::NotFound(json!({ "error": "Email not found or pending verification!" }))
}
}
#[derive(Serialize, Deserialize)]
pub struct Login {
email: String,
password: String,
captcha: Option<String>,
}
/// login to a Revolt account
/// (1) find user by email
/// (2) verify password
/// (3) return access token
#[post("/login", data = "<info>")]
pub async fn login(info: Json<Login>) -> Response {
if let Err(error) = captcha::verify(&info.captcha).await {
return Response::BadRequest(json!({ "error": error }));
}
let col = database::get_collection("users");
if let Some(u) = col
.find_one(doc! { "email": info.email.clone() }, None)
.await
.expect("Failed user lookup")
{
let user: User = from_bson(Bson::Document(u)).expect("Failed to unwrap user.");
match verify(info.password.clone(), &user.password)
.expect("Failed to check hash of password.")
{
true => {
let token = match user.access_token {
Some(t) => t.to_string(),
None => {
let token = gen_token(92);
if col
.update_one(
doc! { "_id": &user.id },
doc! { "$set": { "access_token": token.clone() } },
None,
)
.await
.is_err()
{
return Response::InternalServerError(
json!({ "error": "Failed database operation." }),
);
}
token
}
};
Response::Success(json!({ "access_token": token, "id": user.id }))
}
false => Response::Unauthorized(json!({ "error": "Invalid password." })),
}
} else {
Response::NotFound(json!({ "error": "Email is not registered." }))
}
}
#[derive(Serialize, Deserialize)]
pub struct Token {
token: String,
}
/// login to a Revolt account via token
#[post("/token", data = "<info>")]
pub async fn token(info: Json<Token>) -> Response {
let col = database::get_collection("users");
if let Ok(result) = col.find_one(doc! { "access_token": info.token.clone() }, None).await {
if let Some(user) = result {
Response::Success(json!({
"id": user.get_str("_id").unwrap(),
}))
} else {
Response::Unauthorized(json!({
"error": "Invalid token!",
}))
}
} else {
Response::InternalServerError(json!({
"error": "Failed database query.",
}))
}
}