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