server/src/websocket/server.rs

   1 //! `ChatServer` is an actor. It maintains list of connection client session.
   2 //! And manages available rooms. Peers send messages to other peers in same
   3 //! room through `ChatServer`.
   4
   5 use actix::prelude::*;
   6 use diesel::r2d2::{ConnectionManager, Pool};
   7 use diesel::PgConnection;
   8 use failure::Error;
   9 use rand::{rngs::ThreadRng, Rng};
  10 use serde::{Deserialize, Serialize};
  11 use serde_json::Value;
  12 use std::collections::{HashMap, HashSet};
  13 use std::str::FromStr;
  14 use std::time::SystemTime;
  15
  16 use crate::api::comment::*;
  17 use crate::api::community::*;
  18 use crate::api::post::*;
  19 use crate::api::site::*;
  20 use crate::api::user::*;
  21 use crate::api::*;
  22 use crate::Settings;
  23
  24 /// Chat server sends this messages to session
  25 #[derive(Message)]
  26 #[rtype(result = "()")]
  27 pub struct WSMessage(pub String);
  28
  29 /// Message for chat server communications
  30
  31 /// New chat session is created
  32 #[derive(Message)]
  33 #[rtype(usize)]
  34 pub struct Connect {
  35   pub addr: Recipient<WSMessage>,
  36   pub ip: String,
  37 }
  38
  39 /// Session is disconnected
  40 #[derive(Message)]
  41 #[rtype(result = "()")]
  42 pub struct Disconnect {
  43   pub id: usize,
  44   pub ip: String,
  45 }
  46
  47 // TODO this is unused rn
  48 /// Send message to specific room
  49 #[derive(Message)]
  50 #[rtype(result = "()")]
  51 pub struct ClientMessage {
  52   /// Id of the client session
  53   pub id: usize,
  54   /// Peer message
  55   pub msg: String,
  56   /// Room name
  57   pub room: String,
  58 }
  59
  60 #[derive(Serialize, Deserialize, Message)]
  61 #[rtype(String)]
  62 pub struct StandardMessage {
  63   /// Id of the client session
  64   pub id: usize,
  65   /// Peer message
  66   pub msg: String,
  67 }
  68
  69 #[derive(Debug)]
  70 pub struct RateLimitBucket {
  71   last_checked: SystemTime,
  72   allowance: f64,
  73 }
  74
  75 pub struct SessionInfo {
  76   pub addr: Recipient<WSMessage>,
  77   pub ip: String,
  78 }
  79
  80 /// `ChatServer` manages chat rooms and responsible for coordinating chat
  81 /// session. implementation is super primitive
  82 pub struct ChatServer {
  83   sessions: HashMap<usize, SessionInfo>, // A map from generated random ID to session addr
  84   rate_limits: HashMap<String, RateLimitBucket>,
  85   rooms: HashMap<i32, HashSet<usize>>, // A map from room / post name to set of connectionIDs
  86   rng: ThreadRng,
  87   db: Pool<ConnectionManager<PgConnection>>,
  88 }
  89
  90 impl ChatServer {
  91   pub fn startup(db: Pool<ConnectionManager<PgConnection>>) -> ChatServer {
  92     // default room
  93     let rooms = HashMap::new();
  94
  95     ChatServer {
  96       sessions: HashMap::new(),
  97       rate_limits: HashMap::new(),
  98       rooms,
  99       rng: rand::thread_rng(),
 100       db,
 101     }
 102   }
 103
 104   /// Send message to all users in the room
 105   fn send_room_message(&self, room: i32, message: &str, skip_id: usize) {
 106     if let Some(sessions) = self.rooms.get(&room) {
 107       for id in sessions {
 108         if *id != skip_id {
 109           if let Some(info) = self.sessions.get(id) {
 110             let _ = info.addr.do_send(WSMessage(message.to_owned()));
 111           }
 112         }
 113       }
 114     }
 115   }
 116
 117   fn join_room(&mut self, room_id: i32, id: usize) {
 118     // remove session from all rooms
 119     for sessions in self.rooms.values_mut() {
 120       sessions.remove(&id);
 121     }
 122
 123     // If the room doesn't exist yet
 124     if self.rooms.get_mut(&room_id).is_none() {
 125       self.rooms.insert(room_id, HashSet::new());
 126     }
 127
 128     self.rooms.get_mut(&room_id).unwrap().insert(id);
 129   }
 130
 131   fn send_community_message(
 132     &self,
 133     community_id: i32,
 134     message: &str,
 135     skip_id: usize,
 136   ) -> Result<(), Error> {
 137     use crate::db::post_view::*;
 138     use crate::db::*;
 139
 140     let conn = self.db.get()?;
 141
 142     let posts = PostQueryBuilder::create(&conn)
 143       .listing_type(ListingType::Community)
 144       .sort(&SortType::New)
 145       .for_community_id(community_id)
 146       .limit(9999)
 147       .list()?;
 148
 149     for post in posts {
 150       self.send_room_message(post.id, message, skip_id);
 151     }
 152
 153     Ok(())
 154   }
 155
 156   fn check_rate_limit_register(&mut self, id: usize) -> Result<(), Error> {
 157     self.check_rate_limit_full(
 158       id,
 159       Settings::get().rate_limit.register,
 160       Settings::get().rate_limit.register_per_second,
 161     )
 162   }
 163
 164   fn check_rate_limit_post(&mut self, id: usize) -> Result<(), Error> {
 165     self.check_rate_limit_full(
 166       id,
 167       Settings::get().rate_limit.post,
 168       Settings::get().rate_limit.post_per_second,
 169     )
 170   }
 171
 172   fn check_rate_limit_message(&mut self, id: usize) -> Result<(), Error> {
 173     self.check_rate_limit_full(
 174       id,
 175       Settings::get().rate_limit.message,
 176       Settings::get().rate_limit.message_per_second,
 177     )
 178   }
 179
 180   #[allow(clippy::float_cmp)]
 181   fn check_rate_limit_full(&mut self, id: usize, rate: i32, per: i32) -> Result<(), Error> {
 182     if let Some(info) = self.sessions.get(&id) {
 183       if let Some(rate_limit) = self.rate_limits.get_mut(&info.ip) {
 184         // The initial value
 185         if rate_limit.allowance == -2f64 {
 186           rate_limit.allowance = rate as f64;
 187         };
 188
 189         let current = SystemTime::now();
 190         let time_passed = current.duration_since(rate_limit.last_checked)?.as_secs() as f64;
 191         rate_limit.last_checked = current;
 192         rate_limit.allowance += time_passed * (rate as f64 / per as f64);
 193         if rate_limit.allowance > rate as f64 {
 194           rate_limit.allowance = rate as f64;
 195         }
 196
 197         if rate_limit.allowance < 1.0 {
 198           println!(
 199             "Rate limited IP: {}, time_passed: {}, allowance: {}",
 200             &info.ip, time_passed, rate_limit.allowance
 201           );
 202           Err(
 203             APIError {
 204               op: "Rate Limit".to_string(),
 205               message: format!("Too many requests. {} per {} seconds", rate, per),
 206             }
 207             .into(),
 208           )
 209         } else {
 210           rate_limit.allowance -= 1.0;
 211           Ok(())
 212         }
 213       } else {
 214         Ok(())
 215       }
 216     } else {
 217       Ok(())
 218     }
 219   }
 220 }
 221
 222 /// Make actor from `ChatServer`
 223 impl Actor for ChatServer {
 224   /// We are going to use simple Context, we just need ability to communicate
 225   /// with other actors.
 226   type Context = Context<Self>;
 227 }
 228
 229 /// Handler for Connect message.
 230 ///
 231 /// Register new session and assign unique id to this session
 232 impl Handler<Connect> for ChatServer {
 233   type Result = usize;
 234
 235   fn handle(&mut self, msg: Connect, _ctx: &mut Context<Self>) -> Self::Result {
 236     // notify all users in same room
 237     // self.send_room_message(&"Main".to_owned(), "Someone joined", 0);
 238
 239     // register session with random id
 240     let id = self.rng.gen::<usize>();
 241     println!("{} joined", &msg.ip);
 242
 243     self.sessions.insert(
 244       id,
 245       SessionInfo {
 246         addr: msg.addr,
 247         ip: msg.ip.to_owned(),
 248       },
 249     );
 250
 251     if self.rate_limits.get(&msg.ip).is_none() {
 252       self.rate_limits.insert(
 253         msg.ip,
 254         RateLimitBucket {
 255           last_checked: SystemTime::now(),
 256           allowance: -2f64,
 257         },
 258       );
 259     }
 260
 261     id
 262   }
 263 }
 264
 265 /// Handler for Disconnect message.
 266 impl Handler<Disconnect> for ChatServer {
 267   type Result = ();
 268
 269   fn handle(&mut self, msg: Disconnect, _: &mut Context<Self>) {
 270     // let mut rooms: Vec<i32> = Vec::new();
 271
 272     // remove address
 273     if self.sessions.remove(&msg.id).is_some() {
 274       // remove session from all rooms
 275       for sessions in self.rooms.values_mut() {
 276         if sessions.remove(&msg.id) {
 277           // rooms.push(*id);
 278         }
 279       }
 280     }
 281   }
 282 }
 283
 284 /// Handler for Message message.
 285 impl Handler<StandardMessage> for ChatServer {
 286   type Result = MessageResult<StandardMessage>;
 287
 288   fn handle(&mut self, msg: StandardMessage, _: &mut Context<Self>) -> Self::Result {
 289     let msg_out = match parse_json_message(self, msg) {
 290       Ok(m) => m,
 291       Err(e) => e.to_string(),
 292     };
 293
 294     MessageResult(msg_out)
 295   }
 296 }
 297
 298 fn parse_json_message(chat: &mut ChatServer, msg: StandardMessage) -> Result<String, Error> {
 299   let json: Value = serde_json::from_str(&msg.msg)?;
 300   let data = &json["data"].to_string();
 301   let op = &json["op"].as_str().ok_or(APIError {
 302     op: "Unknown op type".to_string(),
 303     message: "Unknown op type".to_string(),
 304   })?;
 305
 306   let conn = chat.db.get()?;
 307
 308   let user_operation: UserOperation = UserOperation::from_str(&op)?;
 309
 310   match user_operation {
 311     UserOperation::Login => {
 312       let login: Login = serde_json::from_str(data)?;
 313       let res = Oper::new(user_operation, login).perform(&conn)?;
 314       Ok(serde_json::to_string(&res)?)
 315     }
 316     UserOperation::Register => {
 317       let register: Register = serde_json::from_str(data)?;
 318       let res = Oper::new(user_operation, register).perform(&conn);
 319       if res.is_ok() {
 320         chat.check_rate_limit_register(msg.id)?;
 321       }
 322       Ok(serde_json::to_string(&res?)?)
 323     }
 324     UserOperation::GetUserDetails => {
 325       let get_user_details: GetUserDetails = serde_json::from_str(data)?;
 326       let res = Oper::new(user_operation, get_user_details).perform(&conn)?;
 327       Ok(serde_json::to_string(&res)?)
 328     }
 329     UserOperation::SaveUserSettings => {
 330       let save_user_settings: SaveUserSettings = serde_json::from_str(data)?;
 331       let res = Oper::new(user_operation, save_user_settings).perform(&conn)?;
 332       Ok(serde_json::to_string(&res)?)
 333     }
 334     UserOperation::AddAdmin => {
 335       let add_admin: AddAdmin = serde_json::from_str(data)?;
 336       let res = Oper::new(user_operation, add_admin).perform(&conn)?;
 337       Ok(serde_json::to_string(&res)?)
 338     }
 339     UserOperation::BanUser => {
 340       let ban_user: BanUser = serde_json::from_str(data)?;
 341       let res = Oper::new(user_operation, ban_user).perform(&conn)?;
 342       Ok(serde_json::to_string(&res)?)
 343     }
 344     UserOperation::GetReplies => {
 345       let get_replies: GetReplies = serde_json::from_str(data)?;
 346       let res = Oper::new(user_operation, get_replies).perform(&conn)?;
 347       Ok(serde_json::to_string(&res)?)
 348     }
 349     UserOperation::GetUserMentions => {
 350       let get_user_mentions: GetUserMentions = serde_json::from_str(data)?;
 351       let res = Oper::new(user_operation, get_user_mentions).perform(&conn)?;
 352       Ok(serde_json::to_string(&res)?)
 353     }
 354     UserOperation::EditUserMention => {
 355       let edit_user_mention: EditUserMention = serde_json::from_str(data)?;
 356       let res = Oper::new(user_operation, edit_user_mention).perform(&conn)?;
 357       Ok(serde_json::to_string(&res)?)
 358     }
 359     UserOperation::MarkAllAsRead => {
 360       let mark_all_as_read: MarkAllAsRead = serde_json::from_str(data)?;
 361       let res = Oper::new(user_operation, mark_all_as_read).perform(&conn)?;
 362       Ok(serde_json::to_string(&res)?)
 363     }
 364     UserOperation::GetCommunity => {
 365       let get_community: GetCommunity = serde_json::from_str(data)?;
 366       let res = Oper::new(user_operation, get_community).perform(&conn)?;
 367       Ok(serde_json::to_string(&res)?)
 368     }
 369     UserOperation::ListCommunities => {
 370       let list_communities: ListCommunities = serde_json::from_str(data)?;
 371       let res = Oper::new(user_operation, list_communities).perform(&conn)?;
 372       Ok(serde_json::to_string(&res)?)
 373     }
 374     UserOperation::CreateCommunity => {
 375       chat.check_rate_limit_register(msg.id)?;
 376       let create_community: CreateCommunity = serde_json::from_str(data)?;
 377       let res = Oper::new(user_operation, create_community).perform(&conn)?;
 378       Ok(serde_json::to_string(&res)?)
 379     }
 380     UserOperation::EditCommunity => {
 381       let edit_community: EditCommunity = serde_json::from_str(data)?;
 382       let res = Oper::new(user_operation, edit_community).perform(&conn)?;
 383       let mut community_sent: CommunityResponse = res.clone();
 384       community_sent.community.user_id = None;
 385       community_sent.community.subscribed = None;
 386       let community_sent_str = serde_json::to_string(&community_sent)?;
 387       chat.send_community_message(community_sent.community.id, &community_sent_str, msg.id)?;
 388       Ok(serde_json::to_string(&res)?)
 389     }
 390     UserOperation::FollowCommunity => {
 391       let follow_community: FollowCommunity = serde_json::from_str(data)?;
 392       let res = Oper::new(user_operation, follow_community).perform(&conn)?;
 393       Ok(serde_json::to_string(&res)?)
 394     }
 395     UserOperation::GetFollowedCommunities => {
 396       let followed_communities: GetFollowedCommunities = serde_json::from_str(data)?;
 397       let res = Oper::new(user_operation, followed_communities).perform(&conn)?;
 398       Ok(serde_json::to_string(&res)?)
 399     }
 400     UserOperation::BanFromCommunity => {
 401       let ban_from_community: BanFromCommunity = serde_json::from_str(data)?;
 402       let community_id = ban_from_community.community_id;
 403       let res = Oper::new(user_operation, ban_from_community).perform(&conn)?;
 404       let res_str = serde_json::to_string(&res)?;
 405       chat.send_community_message(community_id, &res_str, msg.id)?;
 406       Ok(res_str)
 407     }
 408     UserOperation::AddModToCommunity => {
 409       let mod_add_to_community: AddModToCommunity = serde_json::from_str(data)?;
 410       let community_id = mod_add_to_community.community_id;
 411       let res = Oper::new(user_operation, mod_add_to_community).perform(&conn)?;
 412       let res_str = serde_json::to_string(&res)?;
 413       chat.send_community_message(community_id, &res_str, msg.id)?;
 414       Ok(res_str)
 415     }
 416     UserOperation::ListCategories => {
 417       let list_categories: ListCategories = ListCategories;
 418       let res = Oper::new(user_operation, list_categories).perform(&conn)?;
 419       Ok(serde_json::to_string(&res)?)
 420     }
 421     UserOperation::CreatePost => {
 422       chat.check_rate_limit_post(msg.id)?;
 423       let create_post: CreatePost = serde_json::from_str(data)?;
 424       let res = Oper::new(user_operation, create_post).perform(&conn)?;
 425       Ok(serde_json::to_string(&res)?)
 426     }
 427     UserOperation::GetPost => {
 428       let get_post: GetPost = serde_json::from_str(data)?;
 429       chat.join_room(get_post.id, msg.id);
 430       let res = Oper::new(user_operation, get_post).perform(&conn)?;
 431       Ok(serde_json::to_string(&res)?)
 432     }
 433     UserOperation::GetPosts => {
 434       let get_posts: GetPosts = serde_json::from_str(data)?;
 435       let res = Oper::new(user_operation, get_posts).perform(&conn)?;
 436       Ok(serde_json::to_string(&res)?)
 437     }
 438     UserOperation::CreatePostLike => {
 439       chat.check_rate_limit_message(msg.id)?;
 440       let create_post_like: CreatePostLike = serde_json::from_str(data)?;
 441       let res = Oper::new(user_operation, create_post_like).perform(&conn)?;
 442       Ok(serde_json::to_string(&res)?)
 443     }
 444     UserOperation::EditPost => {
 445       let edit_post: EditPost = serde_json::from_str(data)?;
 446       let res = Oper::new(user_operation, edit_post).perform(&conn)?;
 447       let mut post_sent = res.clone();
 448       post_sent.post.my_vote = None;
 449       let post_sent_str = serde_json::to_string(&post_sent)?;
 450       chat.send_room_message(post_sent.post.id, &post_sent_str, msg.id);
 451       Ok(serde_json::to_string(&res)?)
 452     }
 453     UserOperation::SavePost => {
 454       let save_post: SavePost = serde_json::from_str(data)?;
 455       let res = Oper::new(user_operation, save_post).perform(&conn)?;
 456       Ok(serde_json::to_string(&res)?)
 457     }
 458     UserOperation::CreateComment => {
 459       chat.check_rate_limit_message(msg.id)?;
 460       let create_comment: CreateComment = serde_json::from_str(data)?;
 461       let post_id = create_comment.post_id;
 462       let res = Oper::new(user_operation, create_comment).perform(&conn)?;
 463       let mut comment_sent = res.clone();
 464       comment_sent.comment.my_vote = None;
 465       comment_sent.comment.user_id = None;
 466       let comment_sent_str = serde_json::to_string(&comment_sent)?;
 467       chat.send_room_message(post_id, &comment_sent_str, msg.id);
 468       Ok(serde_json::to_string(&res)?)
 469     }
 470     UserOperation::EditComment => {
 471       let edit_comment: EditComment = serde_json::from_str(data)?;
 472       let post_id = edit_comment.post_id;
 473       let res = Oper::new(user_operation, edit_comment).perform(&conn)?;
 474       let mut comment_sent = res.clone();
 475       comment_sent.comment.my_vote = None;
 476       comment_sent.comment.user_id = None;
 477       let comment_sent_str = serde_json::to_string(&comment_sent)?;
 478       chat.send_room_message(post_id, &comment_sent_str, msg.id);
 479       Ok(serde_json::to_string(&res)?)
 480     }
 481     UserOperation::SaveComment => {
 482       let save_comment: SaveComment = serde_json::from_str(data)?;
 483       let res = Oper::new(user_operation, save_comment).perform(&conn)?;
 484       Ok(serde_json::to_string(&res)?)
 485     }
 486     UserOperation::CreateCommentLike => {
 487       chat.check_rate_limit_message(msg.id)?;
 488       let create_comment_like: CreateCommentLike = serde_json::from_str(data)?;
 489       let post_id = create_comment_like.post_id;
 490       let res = Oper::new(user_operation, create_comment_like).perform(&conn)?;
 491       let mut comment_sent = res.clone();
 492       comment_sent.comment.my_vote = None;
 493       comment_sent.comment.user_id = None;
 494       let comment_sent_str = serde_json::to_string(&comment_sent)?;
 495       chat.send_room_message(post_id, &comment_sent_str, msg.id);
 496       Ok(serde_json::to_string(&res)?)
 497     }
 498     UserOperation::GetModlog => {
 499       let get_modlog: GetModlog = serde_json::from_str(data)?;
 500       let res = Oper::new(user_operation, get_modlog).perform(&conn)?;
 501       Ok(serde_json::to_string(&res)?)
 502     }
 503     UserOperation::CreateSite => {
 504       let create_site: CreateSite = serde_json::from_str(data)?;
 505       let res = Oper::new(user_operation, create_site).perform(&conn)?;
 506       Ok(serde_json::to_string(&res)?)
 507     }
 508     UserOperation::EditSite => {
 509       let edit_site: EditSite = serde_json::from_str(data)?;
 510       let res = Oper::new(user_operation, edit_site).perform(&conn)?;
 511       Ok(serde_json::to_string(&res)?)
 512     }
 513     UserOperation::GetSite => {
 514       let online: usize = chat.sessions.len();
 515       let get_site: GetSite = serde_json::from_str(data)?;
 516       let mut res = Oper::new(user_operation, get_site).perform(&conn)?;
 517       res.online = online;
 518       Ok(serde_json::to_string(&res)?)
 519     }
 520     UserOperation::Search => {
 521       let search: Search = serde_json::from_str(data)?;
 522       let res = Oper::new(user_operation, search).perform(&conn)?;
 523       Ok(serde_json::to_string(&res)?)
 524     }
 525     UserOperation::TransferCommunity => {
 526       let transfer_community: TransferCommunity = serde_json::from_str(data)?;
 527       let res = Oper::new(user_operation, transfer_community).perform(&conn)?;
 528       Ok(serde_json::to_string(&res)?)
 529     }
 530     UserOperation::TransferSite => {
 531       let transfer_site: TransferSite = serde_json::from_str(data)?;
 532       let res = Oper::new(user_operation, transfer_site).perform(&conn)?;
 533       Ok(serde_json::to_string(&res)?)
 534     }
 535     UserOperation::DeleteAccount => {
 536       let delete_account: DeleteAccount = serde_json::from_str(data)?;
 537       let res = Oper::new(user_operation, delete_account).perform(&conn)?;
 538       Ok(serde_json::to_string(&res)?)
 539     }
 540     UserOperation::PasswordReset => {
 541       let password_reset: PasswordReset = serde_json::from_str(data)?;
 542       let res = Oper::new(user_operation, password_reset).perform(&conn)?;
 543       Ok(serde_json::to_string(&res)?)
 544     }
 545     UserOperation::PasswordChange => {
 546       let password_change: PasswordChange = serde_json::from_str(data)?;
 547       let res = Oper::new(user_operation, password_change).perform(&conn)?;
 548       Ok(serde_json::to_string(&res)?)
 549     }
 550   }
 551 }