Adding a register, community, and forum rate limit, based on IP.

- Fixes #124

2 files changed, 114 insertions, 16 deletions

diff --git a/server/src/bin/main.rs b/server/src/bin/main.rs
index 96f8181d..701a4e36 100644
--- a/server/src/bin/main.rs
+++ b/server/src/bin/main.rs
@@ -29,7 +29,8 @@ fn chat_route(req: &HttpRequest<WsChatSessionState>) -> Result<HttpResponse, Err
     req,
     WSSession {
       id: 0,
-      hb: Instant::now()
+      hb: Instant::now(),
+      ip: req.connection_info().host().to_string(),
     },
     )
 }
@@ -37,6 +38,7 @@ fn chat_route(req: &HttpRequest<WsChatSessionState>) -> Result<HttpResponse, Err
 struct WSSession {
   /// unique session id
   id: usize,
+  ip: String,
   /// Client must send ping at least once per 10 seconds (CLIENT_TIMEOUT),
   /// otherwise we drop connection.
   hb: Instant
@@ -61,6 +63,7 @@ impl Actor for WSSession {
       .addr
       .send(Connect {
         addr: addr.recipient(),
+        ip: self.ip.to_owned(),
       })
     .into_actor(self)
       .then(|res, act, ctx| {
@@ -76,7 +79,10 @@ impl Actor for WSSession {
 
   fn stopping(&mut self, ctx: &mut Self::Context) -> Running {
     // notify chat server
-    ctx.state().addr.do_send(Disconnect { id: self.id });
+    ctx.state().addr.do_send(Disconnect { 
+      id: self.id,
+      ip: self.ip.to_owned(),
+    });
     Running::Stop
   }
 }
@@ -111,7 +117,7 @@ impl StreamHandler<ws::Message, ws::ProtocolError> for WSSession {
           .addr
           .send(StandardMessage {
             id: self.id,
-            msg: m
+            msg: m,
           })
         .into_actor(self)
           .then(|res, _, ctx| {
@@ -215,7 +221,7 @@ impl WSSession {
         // notify chat server
         ctx.state()
           .addr
-          .do_send(Disconnect { id: act.id });
+          .do_send(Disconnect { id: act.id, ip: act.ip.to_owned() });
 
         // stop actor
         ctx.stop();
diff --git a/server/src/websocket_server/server.rs b/server/src/websocket_server/server.rs
index fef60e8b..b2416951 100644
--- a/server/src/websocket_server/server.rs
+++ b/server/src/websocket_server/server.rs
@@ -11,6 +11,7 @@ use bcrypt::{verify};
 use std::str::FromStr;
 use diesel::PgConnection;
 use failure::Error;
+use std::time::{SystemTime};
 
 use {Crud, Joinable, Likeable, Followable, Bannable, Saveable, establish_connection, naive_now, naive_from_unix, SortType, SearchType, has_slurs, remove_slurs};
 use actions::community::*;
@@ -25,6 +26,11 @@ use actions::user_view::*;
 use actions::moderator_views::*;
 use actions::moderator::*;
 
+const RATE_LIMIT_MESSAGES: i32 = 30;
+const RATE_LIMIT_PER_SECOND: i32 = 60;
+const RATE_LIMIT_REGISTER_MESSAGES: i32 = 1;
+const RATE_LIMIT_REGISTER_PER_SECOND: i32 = 60;
+
 #[derive(EnumString,ToString,Debug)]
 pub enum UserOperation {
   Login, Register, CreateCommunity, CreatePost, ListCommunities, ListCategories, GetPost, GetCommunity, CreateComment, EditComment, SaveComment, CreateCommentLike, GetPosts, CreatePostLike, EditPost, SavePost, EditCommunity, FollowCommunity, GetFollowedCommunities, GetUserDetails, GetReplies, GetModlog, BanFromCommunity, AddModToCommunity, CreateSite, EditSite, GetSite, AddAdmin, BanUser, Search, MarkAllAsRead
@@ -48,12 +54,14 @@ pub struct WSMessage(pub String);
 #[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
@@ -487,10 +495,22 @@ pub struct MarkAllAsRead {
   auth: 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, Recipient<WSMessage>>, // A map from generated random ID to session addr
+  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,
 }
@@ -502,6 +522,7 @@ impl Default for ChatServer {
 
     ChatServer {
       sessions: HashMap::new(),
+      rate_limits: HashMap::new(),
       rooms: rooms,
       rng: rand::thread_rng(),
     }
@@ -514,8 +535,8 @@ impl ChatServer {
     if let Some(sessions) = self.rooms.get(&room) {
       for id in sessions {
         if *id != skip_id {
-          if let Some(addr) = self.sessions.get(id) {
-            let _ = addr.do_send(WSMessage(message.to_owned()));
+          if let Some(info) = self.sessions.get(id) {
+            let _ = info.addr.do_send(WSMessage(message.to_owned()));
           }
         }
       }
@@ -540,8 +561,50 @@ impl ChatServer {
 
     Ok(())
   }
+
+  fn check_rate_limit_register(&mut self, addr: usize) -> Result<(), Error> {
+    self.check_rate_limit_full(addr, RATE_LIMIT_REGISTER_MESSAGES, RATE_LIMIT_REGISTER_PER_SECOND)
+  }
+
+  fn check_rate_limit(&mut self, addr: usize) -> Result<(), Error> {
+    self.check_rate_limit_full(addr, RATE_LIMIT_MESSAGES, RATE_LIMIT_PER_SECOND)
+  }
+
+  fn check_rate_limit_full(&mut self, addr: usize, rate: i32, per: i32) -> Result<(), Error> {
+    if let Some(info) = self.sessions.get(&addr) {
+      if let Some(rate_limit) = self.rate_limits.get_mut(&info.ip) {
+        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(ErrorMessage {
+            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
@@ -555,14 +618,30 @@ impl Actor for ChatServer {
 impl Handler<Connect> for ChatServer {
   type Result = usize;
 
-  fn handle(&mut self, msg: Connect, _: &mut Context<Self>) -> Self::Result {
+  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>();
-    self.sessions.insert(id, msg.addr);
+    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);
@@ -572,6 +651,7 @@ impl Handler<Connect> for ChatServer {
   }
 }
 
+
 /// Handler for Disconnect message.
 impl Handler<Disconnect> for ChatServer {
   type Result = ();
@@ -794,10 +874,12 @@ impl Perform for Register {
   fn op_type(&self) -> UserOperation {
     UserOperation::Register
   }
-  fn perform(&self, _chat: &mut ChatServer, _addr: usize) -> Result<String, Error> {
+  fn perform(&self, chat: &mut ChatServer, addr: usize) -> Result<String, Error> {
 
     let conn = establish_connection();
 
+    chat.check_rate_limit_register(addr)?;
+
     // Make sure passwords match
     if &self.password != &self.password_verify {
       return Err(self.error("Passwords do not match."))?
@@ -884,10 +966,12 @@ impl Perform for CreateCommunity {
     UserOperation::CreateCommunity
   }
 
-  fn perform(&self, _chat: &mut ChatServer, _addr: usize) -> Result<String, Error> {
+  fn perform(&self, chat: &mut ChatServer, addr: usize) -> Result<String, Error> {
 
     let conn = establish_connection();
 
+    chat.check_rate_limit_register(addr)?;
+
     let claims = match Claims::decode(&self.auth) {
       Ok(claims) => claims.claims,
       Err(_e) => {
@@ -1030,10 +1114,12 @@ impl Perform for CreatePost {
     UserOperation::CreatePost
   }
 
-  fn perform(&self, _chat: &mut ChatServer, _addr: usize) -> Result<String, Error> {
+  fn perform(&self, chat: &mut ChatServer, addr: usize) -> Result<String, Error> {
 
     let conn = establish_connection();
 
+    chat.check_rate_limit_register(addr)?;
+
     let claims = match Claims::decode(&self.auth) {
       Ok(claims) => claims.claims,
       Err(_e) => {
@@ -1242,6 +1328,8 @@ impl Perform for CreateComment {
 
     let conn = establish_connection();
 
+    chat.check_rate_limit(addr)?;
+
     let claims = match Claims::decode(&self.auth) {
       Ok(claims) => claims.claims,
       Err(_e) => {
@@ -1390,7 +1478,7 @@ impl Perform for EditComment {
       deleted: self.deleted.to_owned(),
       read: self.read.to_owned(),
       updated: if self.read.is_some() { orig_comment.updated } else {Some(naive_now())}
-      };
+    };
 
     let _updated_comment = match Comment::update(&conn, self.edit_id, &comment_form) {
       Ok(comment) => comment,
@@ -1500,6 +1588,8 @@ impl Perform for CreateCommentLike {
 
     let conn = establish_connection();
 
+    chat.check_rate_limit(addr)?;
+
     let claims = match Claims::decode(&self.auth) {
       Ok(claims) => claims.claims,
       Err(_e) => {
@@ -1628,10 +1718,12 @@ impl Perform for CreatePostLike {
     UserOperation::CreatePostLike
   }
 
-  fn perform(&self, _chat: &mut ChatServer, _addr: usize) -> Result<String, Error> {
+  fn perform(&self, chat: &mut ChatServer, addr: usize) -> Result<String, Error> {
 
     let conn = establish_connection();
 
+    chat.check_rate_limit(addr)?;
+
     let claims = match Claims::decode(&self.auth) {
       Ok(claims) => claims.claims,
       Err(_e) => {
@@ -2695,7 +2787,7 @@ impl Perform for Search {
       },
       SearchType::Comments => {
         comments = CommentView::list(&conn, 
-                                   &sort, 
+                                     &sort, 
                                      None, 
                                      None, 
                                      Some(self.q.to_owned()),
@@ -2717,7 +2809,7 @@ impl Perform for Search {
                                self.page, 
                                self.limit)?;
         comments = CommentView::list(&conn, 
-                                   &sort, 
+                                     &sort, 
                                      None, 
                                      None, 
                                      Some(self.q.to_owned()),