path: root/server/src/websocket/server.rs

//! `ChatServer` is an actor. It maintains list of connection client session.
//! And manages available rooms. Peers send messages to other peers in same
//! room through `ChatServer`.

use actix::prelude::*;
use rand::{rngs::ThreadRng, Rng};
use std::collections::{HashMap, HashSet};
use serde::{Deserialize, Serialize};
use serde_json::{Value};
use std::str::FromStr;
use failure::Error;
use std::time::{SystemTime};

use crate::api::*;
use crate::api::user::*;
use crate::api::community::*;
use crate::api::post::*;
use crate::api::comment::*;
use crate::api::site::*;

const RATE_LIMIT_MESSAGES: i32 = 30;
const RATE_LIMIT_PER_SECOND: i32 = 60;
const RATE_LIMIT_REGISTER_MESSAGES: i32 = 3;
const RATE_LIMIT_REGISTER_PER_SECOND: i32 = 60*3;


/// Chat server sends this messages to session
#[derive(Message)]
pub struct WSMessage(pub String);

/// Message for chat server communications

/// New chat session is created
#[derive(Message)]
#[rtype(usize)]
pub struct Connect {
  pub addr: Recipient<WSMessage>,
  pub ip: String,
}

/// Session is disconnected
#[derive(Message)]
pub struct Disconnect {
  pub id: usize,
  pub ip: String,
}

/// Send message to specific room
#[derive(Message)]
pub struct ClientMessage {
  /// Id of the client session
  pub id: usize,
  /// Peer message
  pub msg: String,
  /// Room name
  pub room: String,
}

#[derive(Serialize, Deserialize)]
pub struct StandardMessage {
  /// Id of the client session
  pub id: usize,
  /// Peer message
  pub msg: String,
}

impl actix::Message for StandardMessage {
  type Result = String;
}

#[derive(Debug)]
pub struct RateLimitBucket {
  last_checked: SystemTime,
  allowance: f64
}

pub struct SessionInfo {
  pub addr: Recipient<WSMessage>,
  pub ip: String,
}

/// `ChatServer` manages chat rooms and responsible for coordinating chat
/// session. implementation is super primitive
pub struct ChatServer {
  sessions: HashMap<usize, SessionInfo>, // A map from generated random ID to session addr
  rate_limits: HashMap<String, RateLimitBucket>,
  rooms: HashMap<i32, HashSet<usize>>, // A map from room / post name to set of connectionIDs
  rng: ThreadRng,
}

impl Default for ChatServer {
  fn default() -> ChatServer {
    // default room
    let rooms = HashMap::new();

    ChatServer {
      sessions: HashMap::new(),
      rate_limits: HashMap::new(),
      rooms: rooms,
      rng: rand::thread_rng(),
    }
  }
}

impl ChatServer {
  /// Send message to all users in the room
  fn send_room_message(&self, room: &i32, message: &str, skip_id: usize) {
    if let Some(sessions) = self.rooms.get(room) {
      for id in sessions {
        if *id != skip_id {
          if let Some(info) = self.sessions.get(id) {
            let _ = info.addr.do_send(WSMessage(message.to_owned()));
          }
        }
      }
    }
  }

  fn join_room(&mut self, room_id: i32, id: usize) {
    // remove session from all rooms
    for (_n, sessions) in &mut self.rooms {
      sessions.remove(&id);
    }

    // If the room doesn't exist yet
    if self.rooms.get_mut(&room_id).is_none() {
      self.rooms.insert(room_id, HashSet::new());
    }

    &self.rooms.get_mut(&room_id).unwrap().insert(id);
  }

  fn send_community_message(&self, community_id: &i32, message: &str, skip_id: usize) -> Result<(), Error> {
    use crate::db::*;
    use crate::db::post_view::*;
    let conn = establish_connection();
    let posts = PostView::list(
      &conn,
      PostListingType::Community, 
      &SortType::New, 
      Some(*community_id), 
      None,
      None, 
      None,
      false,
      false,
      false,
      None,
      Some(9999))?;
    for post in posts {
      self.send_room_message(&post.id, message, skip_id);
    }

    Ok(())
  }

  fn check_rate_limit_register(&mut self, id: usize) -> Result<(), Error> {
    self.check_rate_limit_full(id, RATE_LIMIT_REGISTER_MESSAGES, RATE_LIMIT_REGISTER_PER_SECOND)
  }

  fn check_rate_limit(&mut self, id: usize) -> Result<(), Error> {
    self.check_rate_limit_full(id, RATE_LIMIT_MESSAGES, RATE_LIMIT_PER_SECOND)
  }

  fn check_rate_limit_full(&mut self, id: usize, rate: i32, per: i32) -> Result<(), Error> {
    if let Some(info) = self.sessions.get(&id) {
      if let Some(rate_limit) = self.rate_limits.get_mut(&info.ip) {
        // The initial value
        if rate_limit.allowance == -2f64 {
          rate_limit.allowance = rate as f64;
        };

        let current = SystemTime::now();
        let time_passed = current.duration_since(rate_limit.last_checked)?.as_secs() as f64;
        rate_limit.last_checked = current;
        rate_limit.allowance += time_passed * (rate as f64 / per as f64);
        if rate_limit.allowance > rate as f64 {
          rate_limit.allowance = rate as f64;
        }

        if rate_limit.allowance < 1.0 {
          println!("Rate limited IP: {}, time_passed: {}, allowance: {}", &info.ip, time_passed, rate_limit.allowance);
          Err(APIError {
            op: "Rate Limit".to_string(), 
            message: format!("Too many requests. {} per {} seconds", rate, per),
          })?
        } else {
          rate_limit.allowance -= 1.0;
          Ok(())
        }
      } else {
        Ok(())
      }
    } else {
      Ok(())
    }
  }
}


/// Make actor from `ChatServer`
impl Actor for ChatServer {
  /// We are going to use simple Context, we just need ability to communicate
  /// with other actors.
  type Context = Context<Self>;
}

/// Handler for Connect message.
///
/// Register new session and assign unique id to this session
impl Handler<Connect> for ChatServer {
  type Result = usize;

  fn handle(&mut self, msg: Connect, _ctx: &mut Context<Self>) -> Self::Result {

    // notify all users in same room
    // self.send_room_message(&"Main".to_owned(), "Someone joined", 0);

    // register session with random id
    let id = self.rng.gen::<usize>();
    println!("{} joined", &msg.ip);

    self.sessions.insert(id, SessionInfo {
      addr: msg.addr,
      ip: msg.ip.to_owned(),
    });

    if self.rate_limits.get(&msg.ip).is_none() {
      self.rate_limits.insert(msg.ip, RateLimitBucket {
        last_checked: SystemTime::now(),
        allowance: -2f64,
      });
    }

    // for (k,v) in &self.rate_limits {
    //   println!("{}: {:?}", k,v);
    // }

    // auto join session to Main room
    // self.rooms.get_mut(&"Main".to_owned()).unwrap().insert(id);

    // send id back
    id
  }
}

/// Handler for Disconnect message.
impl Handler<Disconnect> for ChatServer {
  type Result = ();

  fn handle(&mut self, msg: Disconnect, _: &mut Context<Self>