Adding websocket notification system.

- HTTP and APUB clients can now send live updating messages to websocket
  clients
- Rate limiting now affects both HTTP and websockets
- Rate limiting / Websocket logic is now moved into the API Perform
  functions.
- TODO This broke getting current online users, but that will have to
  wait for the perform trait to be made async.
- Fixes #446

server/src/api/comment.rs

@@ -1,9 +1,4 @@
 use super::*;
-use crate::send_email;
-use crate::settings::Settings;
-use diesel::PgConnection;
-use log::error;
-use std::str::FromStr;
 
 #[derive(Serialize, Deserialize)]
 pub struct CreateComment {
@@ -65,7 +60,12 @@ pub struct GetCommentsResponse {
 }
 
 impl Perform<CommentResponse> for Oper<CreateComment> {
-  fn perform(&self, conn: &PgConnection) -> Result<CommentResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<CommentResponse, Error> {
     let data: &CreateComment = &self.data;
 
     let claims = match Claims::decode(&data.auth) {
@@ -77,6 +77,15 @@ impl Perform<CommentResponse> for Oper<CreateComment> {
 
     let hostname = &format!("https://{}", Settings::get().hostname);
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     // Check for a community ban
     let post = Post::read(&conn, data.post_id)?;
     if CommunityUserBanView::get(&conn, user_id, post.community_id).is_ok() {
@@ -223,15 +232,34 @@ impl Perform<CommentResponse> for Oper<CreateComment> {
 
     let comment_view = CommentView::read(&conn, inserted_comment.id, Some(user_id))?;
 
-    Ok(CommentResponse {
+    let mut res = CommentResponse {
       comment: comment_view,
       recipient_ids,
-    })
+    };
+
+    if let Some(ws) = websocket_info {
+      ws.chatserver.do_send(SendComment {
+        op: UserOperation::CreateComment,
+        comment: res.clone(),
+        my_id: ws.id,
+      });
+
+      // strip out the recipient_ids, so that
+      // users don't get double notifs
+      res.recipient_ids = Vec::new();
+    }
+
+    Ok(res)
   }
 }
 
 impl Perform<CommentResponse> for Oper<EditComment> {
-  fn perform(&self, conn: &PgConnection) -> Result<CommentResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<CommentResponse, Error> {
     let data: &EditComment = &self.data;
 
     let claims = match Claims::decode(&data.auth) {
@@ -241,6 +269,15 @@ impl Perform<CommentResponse> for Oper<EditComment> {
 
     let user_id = claims.id;
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     let orig_comment = CommentView::read(&conn, data.edit_id, None)?;
 
     // You are allowed to mark the comment as read even if you're banned.
@@ -353,15 +390,34 @@ impl Perform<CommentResponse> for Oper<EditComment> {
 
     let comment_view = CommentView::read(&conn, data.edit_id, Some(user_id))?;
 
-    Ok(CommentResponse {
+    let mut res = CommentResponse {
       comment: comment_view,
       recipient_ids,
-    })
+    };
+
+    if let Some(ws) = websocket_info {
+      ws.chatserver.do_send(SendComment {
+        op: UserOperation::EditComment,
+        comment: res.clone(),
+        my_id: ws.id,
+      });
+
+      // strip out the recipient_ids, so that
+      // users don't get double notifs
+      res.recipient_ids = Vec::new();
+    }
+
+    Ok(res)
   }
 }
 
 impl Perform<CommentResponse> for Oper<SaveComment> {
-  fn perform(&self, conn: &PgConnection) -> Result<CommentResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    _websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<CommentResponse, Error> {
     let data: &SaveComment = &self.data;
 
     let claims = match Claims::decode(&data.auth) {
@@ -376,6 +432,15 @@ impl Perform<CommentResponse> for Oper<SaveComment> {
       user_id,
     };
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     if data.save {
       match CommentSaved::save(&conn, &comment_saved_form) {
         Ok(comment) => comment,
@@ -398,7 +463,12 @@ impl Perform<CommentResponse> for Oper<SaveComment> {
 }
 
 impl Perform<CommentResponse> for Oper<CreateCommentLike> {
-  fn perform(&self, conn: &PgConnection) -> Result<CommentResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<CommentResponse, Error> {
     let data: &CreateCommentLike = &self.data;
 
     let claims = match Claims::decode(&data.auth) {
@@ -410,6 +480,15 @@ impl Perform<CommentResponse> for Oper<CreateCommentLike> {
 
     let mut recipient_ids = Vec::new();
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     // Don't do a downvote if site has downvotes disabled
     if data.score == -1 {
       let site = SiteView::read(&conn)?;
@@ -467,15 +546,34 @@ impl Perform<CommentResponse> for Oper<CreateCommentLike> {
     // Have to refetch the comment to get the current state
     let liked_comment = CommentView::read(&conn, data.comment_id, Some(user_id))?;
 
-    Ok(CommentResponse {
+    let mut res = CommentResponse {
       comment: liked_comment,
       recipient_ids,
-    })
+    };
+
+    if let Some(ws) = websocket_info {
+      ws.chatserver.do_send(SendComment {
+        op: UserOperation::CreateCommentLike,
+        comment: res.clone(),
+        my_id: ws.id,
+      });
+
+      // strip out the recipient_ids, so that
+      // users don't get double notifs
+      res.recipient_ids = Vec::new();
+    }
+
+    Ok(res)
   }
 }
 
 impl Perform<GetCommentsResponse> for Oper<GetComments> {
-  fn perform(&self, conn: &PgConnection) -> Result<GetCommentsResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<GetCommentsResponse, Error> {
     let data: &GetComments = &self.data;
 
     let user_claims: Option<Claims> = match &data.auth {
@@ -494,6 +592,15 @@ impl Perform<GetCommentsResponse> for Oper<GetComments> {
     let type_ = ListingType::from_str(&data.type_)?;
     let sort = SortType::from_str(&data.sort)?;
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     let comments = match CommentQueryBuilder::create(&conn)
       .listing_type(type_)
       .sort(&sort)
@@ -507,6 +614,20 @@ impl Perform<GetCommentsResponse> for Oper<GetComments> {
       Err(_e) => return Err(APIError::err("couldnt_get_comments").into()),
     };
 
+    if let Some(ws) = websocket_info {
+      // You don't need to join the specific community room, bc this is already handled by
+      // GetCommunity
+      if data.community_id.is_none() {
+        if let Some(id) = ws.id {
+          // 0 is the "all" community
+          ws.chatserver.do_send(JoinCommunityRoom {
+            community_id: 0,
+            id,
+          });
+        }
+      }
+    }
+
     Ok(GetCommentsResponse { comments })
   }
 }

server/src/api/community.rs

@@ -1,6 +1,4 @@
 use super::*;
-use diesel::PgConnection;
-use std::str::FromStr;
 
 #[derive(Serialize, Deserialize)]
 pub struct GetCommunity {
@@ -55,7 +53,7 @@ pub struct BanFromCommunity {
   auth: String,
 }
 
-#[derive(Serialize, Deserialize)]
+#[derive(Serialize, Deserialize, Clone)]
 pub struct BanFromCommunityResponse {
   user: UserView,
   banned: bool,
@@ -69,7 +67,7 @@ pub struct AddModToCommunity {
   auth: String,
 }
 
-#[derive(Serialize, Deserialize)]
+#[derive(Serialize, Deserialize, Clone)]
 pub struct AddModToCommunityResponse {
   moderators: Vec<CommunityModeratorView>,
 }
@@ -114,7 +112,12 @@ pub struct TransferCommunity {
 }
 
 impl Perform<GetCommunityResponse> for Oper<GetCommunity> {
-  fn perform(&self, conn: &PgConnection) -> Result<GetCommunityResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<GetCommunityResponse, Error> {
     let data: &GetCommunity = &self.data;
 
     let user_id: Option<i32> = match &data.auth {
@@ -128,6 +131,15 @@ impl Perform<GetCommunityResponse> for Oper<GetCommunity> {
       None => None,
     };
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     let community_id = match data.id {
       Some(id) => id,
       None => {
@@ -157,18 +169,41 @@ impl Perform<GetCommunityResponse> for Oper<GetCommunity> {
     let creator_user = admins.remove(creator_index);
     admins.insert(0, creator_user);
 
-    // Return the jwt
-    Ok(GetCommunityResponse {
+    let online = if let Some(ws) = websocket_info {
+      if let Some(id) = ws.id {
+        ws.chatserver
+          .do_send(JoinCommunityRoom { community_id, id });
+      }
+
+      // TODO
+      1
+    // let fut = async {
+    //   ws.chatserver.send(GetCommunityUsersOnline {community_id}).await.unwrap()
+    // };
+    // Runtime::new().unwrap().block_on(fut)
+    } else {
+      0
+    };
+
+    let res = GetCommunityResponse {
       community: community_view,
       moderators,
       admins,
-      online: 0,
-    })
+      online,
+    };
+
+    // Return the jwt
+    Ok(res)
   }
 }
 
 impl Perform<CommunityResponse> for Oper<CreateCommunity> {
-  fn perform(&self, conn: &PgConnection) -> Result<CommunityResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    _websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<CommunityResponse, Error> {
     let data: &CreateCommunity = &self.data;
 
     let claims = match Claims::decode(&data.auth) {
@@ -192,6 +227,15 @@ impl Perform<CommunityResponse> for Oper<CreateCommunity> {
 
     let user_id = claims.id;
 
+    if let Some(rl) = &rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_register(&rl.ip, true)?;
+    }
+
+    let conn = pool.get()?;
+
     // Check for a site ban
     if UserView::read(&conn, user_id)?.banned {
       return Err(APIError::err("site_ban").into());
@@ -239,6 +283,13 @@ impl Perform<CommunityResponse> for Oper<CreateCommunity> {
 
     let community_view = CommunityView::read(&conn, inserted_community.id, Some(user_id))?;
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_register(&rl.ip, false)?;
+    }
+
     Ok(CommunityResponse {
       community: community_view,
     })
@@ -246,7 +297,12 @@ impl Perform<CommunityResponse> for Oper<CreateCommunity> {
 }
 
 impl Perform<CommunityResponse> for Oper<EditCommunity> {
-  fn perform(&self, conn: &PgConnection) -> Result<CommunityResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<CommunityResponse, Error> {
     let data: &EditCommunity = &self.data;
 
     if let Err(slurs) = slur_check(&data.name) {
@@ -270,6 +326,15 @@ impl Perform<CommunityResponse> for Oper<EditCommunity> {
 
     let user_id = claims.id;
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     // Check for a site ban
     if UserView::read(&conn, user_id)?.banned {
       return Err(APIError::err("site_ban").into());
@@ -323,14 +388,35 @@ impl Perform<CommunityResponse> for Oper<EditCommunity> {
 
     let community_view = CommunityView::read(&conn, data.edit_id, Some(user_id))?;
 
-    Ok(CommunityResponse {
+    let res = CommunityResponse {
       community: community_view,
-    })
+    };
+
+    if let Some(ws) = websocket_info {
+      // Strip out the user id and subscribed when sending to others
+      let mut res_sent = res.clone();
+      res_sent.community.user_id = None;
+      res_sent.community.subscribed = None;
+
+      ws.chatserver.do_send(SendCommunityRoomMessage {
+        op: UserOperation::EditCommunity,
+        response: res_sent,
+        community_id: data.edit_id,
+        my_id: ws.id,
+      });
+    }
+
+    Ok(res)
   }
 }
 
 impl Perform<ListCommunitiesResponse> for Oper<ListCommunities> {
-  fn perform(&self, conn: &PgConnection) -> Result<ListCommunitiesResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    _websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<ListCommunitiesResponse, Error> {
     let data: &ListCommunities = &self.data;
 
     let user_claims: Option<Claims> = match &data.auth {
@@ -353,6 +439,15 @@ impl Perform<ListCommunitiesResponse> for Oper<ListCommunities> {
 
     let sort = SortType::from_str(&data.sort)?;
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     let communities = CommunityQueryBuilder::create(&conn)
       .sort(&sort)
       .for_user(user_id)
@@ -367,7 +462,12 @@ impl Perform<ListCommunitiesResponse> for Oper<ListCommunities> {
 }
 
 impl Perform<CommunityResponse> for Oper<FollowCommunity> {
-  fn perform(&self, conn: &PgConnection) -> Result<CommunityResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    _websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<CommunityResponse, Error> {
     let data: &FollowCommunity = &self.data;
 
     let claims = match Claims::decode(&data.auth) {
@@ -382,6 +482,15 @@ impl Perform<CommunityResponse> for Oper<FollowCommunity> {
       user_id,
     };
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     if data.follow {
       match CommunityFollower::follow(&conn, &community_follower_form) {
         Ok(user) => user,
@@ -403,7 +512,12 @@ impl Perform<CommunityResponse> for Oper<FollowCommunity> {
 }
 
 impl Perform<GetFollowedCommunitiesResponse> for Oper<GetFollowedCommunities> {
-  fn perform(&self, conn: &PgConnection) -> Result<GetFollowedCommunitiesResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    _websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<GetFollowedCommunitiesResponse, Error> {
     let data: &GetFollowedCommunities = &self.data;
 
     let claims = match Claims::decode(&data.auth) {
@@ -413,6 +527,15 @@ impl Perform<GetFollowedCommunitiesResponse> for Oper<GetFollowedCommunities> {
 
     let user_id = claims.id;
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     let communities: Vec<CommunityFollowerView> =
       match CommunityFollowerView::for_user(&conn, user_id) {
         Ok(communities) => communities,
@@ -425,7 +548,12 @@ impl Perform<GetFollowedCommunitiesResponse> for Oper<GetFollowedCommunities> {
 }
 
 impl Perform<BanFromCommunityResponse> for Oper<BanFromCommunity> {
-  fn perform(&self, conn: &PgConnection) -> Result<BanFromCommunityResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<BanFromCommunityResponse, Error> {
     let data: &BanFromCommunity = &self.data;
 
     let claims = match Claims::decode(&data.auth) {
@@ -440,6 +568,15 @@ impl Perform<BanFromCommunityResponse> for Oper<BanFromCommunity> {
       user_id: data.user_id,
     };
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     if data.ban {
       match CommunityUserBan::ban(&conn, &community_user_ban_form) {
         Ok(user) => user,
@@ -470,15 +607,31 @@ impl Perform<BanFromCommunityResponse> for Oper<BanFromCommunity> {
 
     let user_view = UserView::read(&conn, data.user_id)?;
 
-    Ok(BanFromCommunityResponse {
+    let res = BanFromCommunityResponse {
       user: user_view,
       banned: data.ban,
-    })
+    };
+
+    if let Some(ws) = websocket_info {
+      ws.chatserver.do_send(SendCommunityRoomMessage {
+        op: UserOperation::BanFromCommunity,
+        response: res.clone(),
+        community_id: data.community_id,
+        my_id: ws.id,
+      });
+    }
+
+    Ok(res)
   }
 }
 
 impl Perform<AddModToCommunityResponse> for Oper<AddModToCommunity> {
-  fn perform(&self, conn: &PgConnection) -> Result<AddModToCommunityResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<AddModToCommunityResponse, Error> {
     let data: &AddModToCommunity = &self.data;
 
     let claims = match Claims::decode(&data.auth) {
@@ -493,6 +646,15 @@ impl Perform<AddModToCommunityResponse> for Oper<AddModToCommunity> {
       user_id: data.user_id,
     };
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     if data.added {
       match CommunityModerator::join(&conn, &community_moderator_form) {
         Ok(user) => user,
@@ -516,12 +678,28 @@ impl Perform<AddModToCommunityResponse> for Oper<AddModToCommunity> {
 
     let moderators = CommunityModeratorView::for_community(&conn, data.community_id)?;
 
-    Ok(AddModToCommunityResponse { moderators })
+    let res = AddModToCommunityResponse { moderators };
+
+    if let Some(ws) = websocket_info {
+      ws.chatserver.do_send(SendCommunityRoomMessage {
+        op: UserOperation::AddModToCommunity,
+        response: res.clone(),
+        community_id: data.community_id,
+        my_id: ws.id,
+      });
+    }
+
+    Ok(res)
   }
 }
 
 impl Perform<GetCommunityResponse> for Oper<TransferCommunity> {
-  fn perform(&self, conn: &PgConnection) -> Result<GetCommunityResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    _websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<GetCommunityResponse, Error> {
     let data: &TransferCommunity = &self.data;
 
     let claims = match Claims::decode(&data.auth) {
@@ -531,6 +709,15 @@ impl Perform<GetCommunityResponse> for Oper<TransferCommunity> {
 
     let user_id = claims.id;
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     let read_community = Community::read(&conn, data.community_id)?;
 
     let site_creator_id = Site::read(&conn, 1)?.creator_id;

server/src/api/mod.rs

@@ -18,12 +18,26 @@ use crate::db::user_mention_view::*;
 use crate::db::user_view::*;
 use crate::db::*;
 use crate::{
-  extract_usernames, fetch_iframely_and_pictshare_data, naive_from_unix, naive_now, remove_slurs,
-  slur_check, slurs_vec_to_str,
+  extract_usernames, fetch_iframely_and_pictshare_data, generate_random_string, naive_from_unix,
+  naive_now, remove_slurs, send_email, slur_check, slurs_vec_to_str,
 };
+
+use crate::rate_limit::RateLimitInfo;
+use crate::settings::Settings;
+use crate::websocket::UserOperation;
+use crate::websocket::{
+  server::{
+    JoinCommunityRoom, JoinPostRoom, JoinUserRoom, SendAllMessage, SendComment,
+    SendCommunityRoomMessage, SendPost, SendUserRoomMessage,
+  },
+  WebsocketInfo,
+};
+use diesel::r2d2::{ConnectionManager, Pool};
 use diesel::PgConnection;
 use failure::Error;
+use log::{error, info};
 use serde::{Deserialize, Serialize};
+use std::str::FromStr;
 
 pub mod comment;
 pub mod community;
@@ -56,7 +70,12 @@ impl<T> Oper<T> {
 }
 
 pub trait Perform<T> {
-  fn perform(&self, conn: &PgConnection) -> Result<T, Error>
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<T, Error>
   where
     T: Sized;
 }

server/src/api/post.rs

@@ -1,6 +1,4 @@
 use super::*;
-use diesel::PgConnection;
-use std::str::FromStr;
 
 #[derive(Serialize, Deserialize)]
 pub struct CreatePost {
@@ -80,7 +78,12 @@ pub struct SavePost {
 }
 
 impl Perform<PostResponse> for Oper<CreatePost> {
-  fn perform(&self, conn: &PgConnection) -> Result<PostResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<PostResponse, Error> {
     let data: &CreatePost = &self.data;
 
     let claims = match Claims::decode(&data.auth) {
@@ -100,6 +103,15 @@ impl Perform<PostResponse> for Oper<CreatePost> {
 
     let user_id = claims.id;
 
+    if let Some(rl) = &rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_post(&rl.ip, true)?;
+    }
+
+    let conn = pool.get()?;
+
     // Check for a community ban
     if CommunityUserBanView::get(&conn, user_id, data.community_id).is_ok() {
       return Err(APIError::err("community_ban").into());
@@ -164,12 +176,34 @@ impl Perform<PostResponse> for Oper<CreatePost> {
       Err(_e) => return Err(APIError::err("couldnt_find_post").into()),
     };
 
-    Ok(PostResponse { post: post_view })
+    if let Some(rl) = &rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_post(&rl.ip, false)?;
+    }
+
+    let res = PostResponse { post: post_view };
+
+    if let Some(ws) = websocket_info {
+      ws.chatserver.do_send(SendPost {
+        op: UserOperation::CreatePost,
+        post: res.clone(),
+        my_id: ws.id,
+      });
+    }
+
+    Ok(res)
   }
 }
 
 impl Perform<GetPostResponse> for Oper<GetPost> {
-  fn perform(&self, conn: &PgConnection) -> Result<GetPostResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<GetPostResponse, Error> {
     let data: &GetPost = &self.data;
 
     let user_id: Option<i32> = match &data.auth {
@@ -183,6 +217,15 @@ impl Perform<GetPostResponse> for Oper<GetPost> {
       None => None,
     };
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     let post_view = match PostView::read(&conn, data.id, user_id) {
       Ok(post) => post,
       Err(_e) => return Err(APIError::err("couldnt_find_post").into()),
@@ -204,6 +247,24 @@ impl Perform<GetPostResponse> for Oper<GetPost> {
     let creator_user = admins.remove(creator_index);
     admins.insert(0, creator_user);
 
+    let online = if let Some(ws) = websocket_info {
+      if let Some(id) = ws.id {
+        ws.chatserver.do_send(JoinPostRoom {
+          post_id: data.id,
+          id,
+        });
+      }
+
+      // TODO
+      1
+    // let fut = async {
+    //   ws.chatserver.send(GetPostUsersOnline {post_id: data.id}).await.unwrap()
+    // };
+    // Runtime::new().unwrap().block_on(fut)
+    } else {
+      0
+    };
+
     // Return the jwt
     Ok(GetPostResponse {
       post: post_view,
@@ -211,13 +272,18 @@ impl Perform<GetPostResponse> for Oper<GetPost> {
       community,
       moderators,
       admins,
-      online: 0,
+      online,
     })
   }
 }
 
 impl Perform<GetPostsResponse> for Oper<GetPosts> {
-  fn perform(&self, conn: &PgConnection) -> Result<GetPostsResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<GetPostsResponse, Error> {
     let data: &GetPosts = &self.data;
 
     let user_claims: Option<Claims> = match &data.auth {
@@ -241,6 +307,15 @@ impl Perform<GetPostsResponse> for Oper<GetPosts> {
     let type_ = ListingType::from_str(&data.type_)?;
     let sort = SortType::from_str(&data.sort)?;
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     let posts = match PostQueryBuilder::create(&conn)
       .listing_type(type_)
       .sort(&sort)
@@ -255,12 +330,31 @@ impl Perform<GetPostsResponse> for Oper<GetPosts> {
       Err(_e) => return Err(APIError::err("couldnt_get_posts").into()),
     };
 
+    if let Some(ws) = websocket_info {
+      // You don't need to join the specific community room, bc this is already handled by
+      // GetCommunity
+      if data.community_id.is_none() {
+        if let Some(id) = ws.id {
+          // 0 is the "all" community
+          ws.chatserver.do_send(JoinCommunityRoom {
+            community_id: 0,
+            id,
+          });
+        }
+      }
+    }
+
     Ok(GetPostsResponse { posts })
   }
 }
 
 impl Perform<PostResponse> for Oper<CreatePostLike> {
-  fn perform(&self, conn: &PgConnection) -> Result<PostResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<PostResponse, Error> {
     let data: &CreatePostLike = &self.data;
 
     let claims = match Claims::decode(&data.auth) {
@@ -270,6 +364,15 @@ impl Perform<PostResponse> for Oper<CreatePostLike> {
 
     let user_id = claims.id;
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     // Don't do a downvote if site has downvotes disabled
     if data.score == -1 {
       let site = SiteView::read(&conn)?;
@@ -312,13 +415,27 @@ impl Perform<PostResponse> for Oper<CreatePostLike> {
       Err(_e) => return Err(APIError::err("couldnt_find_post").into()),
     };
 
-    // just output the score
-    Ok(PostResponse { post: post_view })
+    let res = PostResponse { post: post_view };
+
+    if let Some(ws) = websocket_info {
+      ws.chatserver.do_send(SendPost {
+        op: UserOperation::CreatePostLike,
+        post: res.clone(),
+        my_id: ws.id,
+      });
+    }
+
+    Ok(res)
   }
 }
 
 impl Perform<PostResponse> for Oper<EditPost> {
-  fn perform(&self, conn: &PgConnection) -> Result<PostResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<PostResponse, Error> {
     let data: &EditPost = &self.data;
 
     if let Err(slurs) = slur_check(&data.name) {
@@ -338,6 +455,15 @@ impl Perform<PostResponse> for Oper<EditPost> {
 
     let user_id = claims.id;
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     // Verify its the creator or a mod or admin
     let mut editors: Vec<i32> = vec![data.creator_id];
     editors.append(
@@ -427,12 +553,27 @@ impl Perform<PostResponse> for Oper<EditPost> {
 
     let post_view = PostView::read(&conn, data.edit_id, Some(user_id))?;
 
-    Ok(PostResponse { post: post_view })
+    let res = PostResponse { post: post_view };
+
+    if let Some(ws) = websocket_info {
+      ws.chatserver.do_send(SendPost {
+        op: UserOperation::EditPost,
+        post: res.clone(),
+        my_id: ws.id,
+      });
+    }
+
+    Ok(res)
   }
 }
 
 impl Perform<PostResponse> for Oper<SavePost> {
-  fn perform(&self, conn: &PgConnection) -> Result<PostResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    _websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<PostResponse, Error> {
     let data: &SavePost = &self.data;
 
     let claims = match Claims::decode(&data.auth) {
@@ -447,6 +588,15 @@ impl Perform<PostResponse> for Oper<SavePost> {
       user_id,
     };
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     if data.save {
       match PostSaved::save(&conn, &post_saved_form) {
         Ok(post) => post,

server/src/api/site.rs

@@ -1,10 +1,5 @@
+use super::user::Register;
 use super::*;
-use crate::api::user::Register;
-use crate::api::{Oper, Perform};
-use crate::settings::Settings;
-use diesel::PgConnection;
-use log::info;
-use std::str::FromStr;
 
 #[derive(Serialize, Deserialize)]
 pub struct ListCategories {}
@@ -78,7 +73,7 @@ pub struct EditSite {
 #[derive(Serialize, Deserialize)]
 pub struct GetSite {}
 
-#[derive(Serialize, Deserialize)]
+#[derive(Serialize, Deserialize, Clone)]
 pub struct SiteResponse {
   site: SiteView,
 }
@@ -114,9 +109,23 @@ pub struct SaveSiteConfig {
 }
 
 impl Perform<ListCategoriesResponse> for Oper<ListCategories> {
-  fn perform(&self, conn: &PgConnection) -> Result<ListCategoriesResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    _websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<ListCategoriesResponse, Error> {
     let _data: &ListCategories = &self.data;
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     let categories: Vec<Category> = Category::list_all(&conn)?;
 
     // Return the jwt
@@ -125,9 +134,23 @@ impl Perform<ListCategoriesResponse> for Oper<ListCategories> {
 }
 
 impl Perform<GetModlogResponse> for Oper<GetModlog> {
-  fn perform(&self, conn: &PgConnection) -> Result<GetModlogResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    _websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<GetModlogResponse, Error> {
     let data: &GetModlog = &self.data;
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     let removed_posts = ModRemovePostView::list(
       &conn,
       data.community_id,
@@ -198,7 +221,12 @@ impl Perform<GetModlogResponse> for Oper<GetModlog> {
 }
 
 impl Perform<SiteResponse> for Oper<CreateSite> {
-  fn perform(&self, conn: &PgConnection) -> Result<SiteResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    _websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<SiteResponse, Error> {
     let data: &CreateSite = &self.data;
 
     let claims = match Claims::decode(&data.auth) {
@@ -218,6 +246,15 @@ impl Perform<SiteResponse> for Oper<CreateSite> {
 
     let user_id = claims.id;
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     // Make sure user is an admin
     if !UserView::read(&conn, user_id)?.admin {
       return Err(APIError::err("not_an_admin").into());
@@ -245,7 +282,12 @@ impl Perform<SiteResponse> for Oper<CreateSite> {
 }
 
 impl Perform<SiteResponse> for Oper<EditSite> {
-  fn perform(&self, conn: &PgConnection) -> Result<SiteResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<SiteResponse, Error> {
     let data: &EditSite = &self.data;
 
     let claims = match Claims::decode(&data.auth) {
@@ -265,6 +307,15 @@ impl Perform<SiteResponse> for Oper<EditSite> {
 
     let user_id = claims.id;
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     // Make sure user is an admin
     if !UserView::read(&conn, user_id)?.admin {
       return Err(APIError::err("not_an_admin").into());
@@ -289,14 +340,39 @@ impl Perform<SiteResponse> for Oper<EditSite> {
 
     let site_view = SiteView::read(&conn)?;
 
-    Ok(SiteResponse { site: site_view })
+    let res = SiteResponse { site: site_view };
+
+    if let Some(ws) = websocket_info {
+      ws.chatserver.do_send(SendAllMessage {
+        op: UserOperation::EditSite,
+        response: res.clone(),
+        my_id: ws.id,
+      });
+    }
+
+    Ok(res)
   }
 }
 
 impl Perform<GetSiteResponse> for Oper<GetSite> {
-  fn perform(&self, conn: &PgConnection) -> Result<GetSiteResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<GetSiteResponse, Error> {
     let _data: &GetSite = &self.data;
 
+    if let Some(rl) = &rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
+    // TODO refactor this a little
     let site = Site::read(&conn, 1);
     let site_view = if site.is_ok() {
       Some(SiteView::read(&conn)?)
@@ -309,7 +385,11 @@ impl Perform<GetSiteResponse> for Oper<GetSite> {
         admin: true,
         show_nsfw: true,
       };
-      let login_response = Oper::new(register).perform(&conn)?;
+      let login_response = Oper::new(register).perform(
+        pool.clone(),
+        websocket_info.clone(),
+        rate_limit_info.clone(),
+      )?;
       info!("Admin {} created", setup.admin_username);
 
       let create_site = CreateSite {
@@ -320,7 +400,7 @@ impl Perform<GetSiteResponse> for Oper<GetSite> {
         enable_nsfw: false,
         auth: login_response.jwt,
       };
-      Oper::new(create_site).perform(&conn)?;
+      Oper::new(create_site).perform(pool, websocket_info.clone(), rate_limit_info)?;
       info!("Site {} created", setup.site_name);
       Some(SiteView::read(&conn)?)
     } else {
@@ -337,17 +417,33 @@ impl Perform<GetSiteResponse> for Oper<GetSite> {
 
     let banned = UserView::banned(&conn)?;
 
+    let online = if let Some(_ws) = websocket_info {
+      // TODO
+      1
+    // let fut = async {
+    //   ws.chatserver.send(GetUsersOnline).await.unwrap()
+    // };
+    // Runtime::new().unwrap().block_on(fut)
+    } else {
+      0
+    };
+
     Ok(GetSiteResponse {
       site: site_view,
       admins,
       banned,
-      online: 0,
+      online,
     })
   }
 }
 
 impl Perform<SearchResponse> for Oper<Search> {
-  fn perform(&self, conn: &PgConnection) -> Result<SearchResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    _websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<SearchResponse, Error> {
     let data: &Search = &self.data;
 
     let user_id: Option<i32> = match &data.auth {
@@ -371,6 +467,15 @@ impl Perform<SearchResponse> for Oper<Search> {
 
     // TODO no clean / non-nsfw searching rn
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     match type_ {
       SearchType::Posts => {
         posts = PostQueryBuilder::create(&conn)
@@ -465,7 +570,12 @@ impl Perform<SearchResponse> for Oper<Search> {
 }
 
 impl Perform<GetSiteResponse> for Oper<TransferSite> {
-  fn perform(&self, conn: &PgConnection) -> Result<GetSiteResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    _websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<GetSiteResponse, Error> {
     let data: &TransferSite = &self.data;
 
     let claims = match Claims::decode(&data.auth) {
@@ -475,6 +585,15 @@ impl Perform<GetSiteResponse> for Oper<TransferSite> {
 
     let user_id = claims.id;
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     let read_site = Site::read(&conn, 1)?;
 
     // Make sure user is the creator
@@ -528,7 +647,12 @@ impl Perform<GetSiteResponse> for Oper<TransferSite> {
 }
 
 impl Perform<GetSiteConfigResponse> for Oper<GetSiteConfig> {
-  fn perform(&self, conn: &PgConnection) -> Result<GetSiteConfigResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    _websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<GetSiteConfigResponse, Error> {
     let data: &GetSiteConfig = &self.data;
 
     let claims = match Claims::decode(&data.auth) {
@@ -538,6 +662,15 @@ impl Perform<GetSiteConfigResponse> for Oper<GetSiteConfig> {
 
     let user_id = claims.id;
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     // Only let admins read this
     let admins = UserView::admins(&conn)?;
     let admin_ids: Vec<i32> = admins.into_iter().map(|m| m.id).collect();
@@ -553,7 +686,12 @@ impl Perform<GetSiteConfigResponse> for Oper<GetSiteConfig> {
 }
 
 impl Perform<GetSiteConfigResponse> for Oper<SaveSiteConfig> {
-  fn perform(&self, conn: &PgConnection) -> Result<GetSiteConfigResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    _websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<GetSiteConfigResponse, Error> {
     let data: &SaveSiteConfig = &self.data;
 
     let claims = match Claims::decode(&data.auth) {
@@ -563,6 +701,15 @@ impl Perform<GetSiteConfigResponse> for Oper<SaveSiteConfig> {
 
     let user_id = claims.id;
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     // Only let admins read this
     let admins = UserView::admins(&conn)?;
     let admin_ids: Vec<i32> = admins.into_iter().map(|m| m.id).collect();

server/src/api/user.rs

@@ -1,10 +1,5 @@
 use super::*;
-use crate::settings::Settings;
-use crate::{generate_random_string, send_email};
 use bcrypt::verify;
-use diesel::PgConnection;
-use log::error;
-use std::str::FromStr;
 
 #[derive(Serialize, Deserialize, Debug)]
 pub struct Login {
@@ -89,7 +84,7 @@ pub struct AddAdmin {
   auth: String,
 }
 
-#[derive(Serialize, Deserialize)]
+#[derive(Serialize, Deserialize, Clone)]
 pub struct AddAdminResponse {
   admins: Vec<UserView>,
 }
@@ -103,7 +98,7 @@ pub struct BanUser {
   auth: String,
 }
 
-#[derive(Serialize, Deserialize)]
+#[derive(Serialize, Deserialize, Clone)]
 pub struct BanUserResponse {
   user: UserView,
   banned: bool,
@@ -205,9 +200,23 @@ pub struct UserJoinResponse {
 }
 
 impl Perform<LoginResponse> for Oper<Login> {
-  fn perform(&self, conn: &PgConnection) -> Result<LoginResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    _websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<LoginResponse, Error> {
     let data: &Login = &self.data;
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     // Fetch that username / email
     let user: User_ = match User_::find_by_email_or_username(&conn, &data.username_or_email) {
       Ok(user) => user,
@@ -226,9 +235,23 @@ impl Perform<LoginResponse> for Oper<Login> {
 }
 
 impl Perform<LoginResponse> for Oper<Register> {
-  fn perform(&self, conn: &PgConnection) -> Result<LoginResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    _websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<LoginResponse, Error> {
     let data: &Register = &self.data;
 
+    if let Some(rl) = &rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_register(&rl.ip, true)?;
+    }
+
+    let conn = pool.get()?;
+
     // Make sure site has open registration
     if let Ok(site) = SiteView::read(&conn) {
       if !site.open_registration {
@@ -332,6 +355,13 @@ impl Perform<LoginResponse> for Oper<Register> {
         };
     }
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_register(&rl.ip, false)?;
+    }
+
     // Return the jwt
     Ok(LoginResponse {
       jwt: inserted_user.jwt(),
@@ -340,7 +370,12 @@ impl Perform<LoginResponse> for Oper<Register> {
 }
 
 impl Perform<LoginResponse> for Oper<SaveUserSettings> {
-  fn perform(&self, conn: &PgConnection) -> Result<LoginResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    _websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<LoginResponse, Error> {
     let data: &SaveUserSettings = &self.data;
 
     let claims = match Claims::decode(&data.auth) {
@@ -350,6 +385,15 @@ impl Perform<LoginResponse> for Oper<SaveUserSettings> {
 
     let user_id = claims.id;
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     let read_user = User_::read(&conn, user_id)?;
 
     let email = match &data.email {
@@ -428,9 +472,23 @@ impl Perform<LoginResponse> for Oper<SaveUserSettings> {
 }
 
 impl Perform<GetUserDetailsResponse> for Oper<GetUserDetails> {
-  fn perform(&self, conn: &PgConnection) -> Result<GetUserDetailsResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    _websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<GetUserDetailsResponse, Error> {
     let data: &GetUserDetails = &self.data;
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     let user_claims: Option<Claims> = match &data.auth {
       Some(auth) => match Claims::decode(&auth) {
         Ok(claims) => Some(claims.claims),
@@ -525,7 +583,12 @@ impl Perform<GetUserDetailsResponse> for Oper<GetUserDetails> {
 }
 
 impl Perform<AddAdminResponse> for Oper<AddAdmin> {
-  fn perform(&self, conn: &PgConnection) -> Result<AddAdminResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<AddAdminResponse, Error> {
     let data: &AddAdmin = &self.data;
 
     let claims = match Claims::decode(&data.auth) {
@@ -535,6 +598,15 @@ impl Perform<AddAdminResponse> for Oper<AddAdmin> {
 
     let user_id = claims.id;
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     // Make sure user is an admin
     if !UserView::read(&conn, user_id)?.admin {
       return Err(APIError::err("not_an_admin").into());
@@ -583,12 +655,27 @@ impl Perform<AddAdminResponse> for Oper<AddAdmin> {
     let creator_user = admins.remove(creator_index);
     admins.insert(0, creator_user);
 
-    Ok(AddAdminResponse { admins })
+    let res = AddAdminResponse { admins };
+
+    if let Some(ws) = websocket_info {
+      ws.chatserver.do_send(SendAllMessage {
+        op: UserOperation::AddAdmin,
+        response: res.clone(),
+        my_id: ws.id,
+      });
+    }
+
+    Ok(res)
   }
 }
 
 impl Perform<BanUserResponse> for Oper<BanUser> {
-  fn perform(&self, conn: &PgConnection) -> Result<BanUserResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<BanUserResponse, Error> {
     let data: &BanUser = &self.data;
 
     let claims = match Claims::decode(&data.auth) {
@@ -598,6 +685,15 @@ impl Perform<BanUserResponse> for Oper<BanUser> {
 
     let user_id = claims.id;
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     // Make sure user is an admin
     if !UserView::read(&conn, user_id)?.admin {
       return Err(APIError::err("not_an_admin").into());
@@ -649,15 +745,30 @@ impl Perform<BanUserResponse> for Oper<BanUser> {
 
     let user_view = UserView::read(&conn, data.user_id)?;
 
-    Ok(BanUserResponse {
+    let res = BanUserResponse {
       user: user_view,
       banned: data.ban,
-    })
+    };
+
+    if let Some(ws) = websocket_info {
+      ws.chatserver.do_send(SendAllMessage {
+        op: UserOperation::BanUser,
+        response: res.clone(),
+        my_id: ws.id,
+      });
+    }
+
+    Ok(res)
   }
 }
 
 impl Perform<GetRepliesResponse> for Oper<GetReplies> {
-  fn perform(&self, conn: &PgConnection) -> Result<GetRepliesResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    _websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<GetRepliesResponse, Error> {
     let data: &GetReplies = &self.data;
 
     let claims = match Claims::decode(&data.auth) {
@@ -669,6 +780,15 @@ impl Perform<GetRepliesResponse> for Oper<GetReplies> {
 
     let sort = SortType::from_str(&data.sort)?;
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     let replies = ReplyQueryBuilder::create(&conn, user_id)
       .sort(&sort)
       .unread_only(data.unread_only)
@@ -681,7 +801,12 @@ impl Perform<GetRepliesResponse> for Oper<GetReplies> {
 }
 
 impl Perform<GetUserMentionsResponse> for Oper<GetUserMentions> {
-  fn perform(&self, conn: &PgConnection) -> Result<GetUserMentionsResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    _websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<GetUserMentionsResponse, Error> {
     let data: &GetUserMentions = &self.data;
 
     let claims = match Claims::decode(&data.auth) {
@@ -693,6 +818,15 @@ impl Perform<GetUserMentionsResponse> for Oper<GetUserMentions> {
 
     let sort = SortType::from_str(&data.sort)?;
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     let mentions = UserMentionQueryBuilder::create(&conn, user_id)
       .sort(&sort)
       .unread_only(data.unread_only)
@@ -705,7 +839,12 @@ impl Perform<GetUserMentionsResponse> for Oper<GetUserMentions> {
 }
 
 impl Perform<UserMentionResponse> for Oper<EditUserMention> {
-  fn perform(&self, conn: &PgConnection) -> Result<UserMentionResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    _websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<UserMentionResponse, Error> {
     let data: &EditUserMention = &self.data;
 
     let claims = match Claims::decode(&data.auth) {
@@ -715,6 +854,15 @@ impl Perform<UserMentionResponse> for Oper<EditUserMention> {
 
     let user_id = claims.id;
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     let user_mention = UserMention::read(&conn, data.user_mention_id)?;
 
     let user_mention_form = UserMentionForm {
@@ -738,7 +886,12 @@ impl Perform<UserMentionResponse> for Oper<EditUserMention> {
 }
 
 impl Perform<GetRepliesResponse> for Oper<MarkAllAsRead> {
-  fn perform(&self, conn: &PgConnection) -> Result<GetRepliesResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    _websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<GetRepliesResponse, Error> {
     let data: &MarkAllAsRead = &self.data;
 
     let claims = match Claims::decode(&data.auth) {
@@ -748,6 +901,15 @@ impl Perform<GetRepliesResponse> for Oper<MarkAllAsRead> {
 
     let user_id = claims.id;
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     let replies = ReplyQueryBuilder::create(&conn, user_id)
       .unread_only(true)
       .page(1)
@@ -822,7 +984,12 @@ impl Perform<GetRepliesResponse> for Oper<MarkAllAsRead> {
 }
 
 impl Perform<LoginResponse> for Oper<DeleteAccount> {
-  fn perform(&self, conn: &PgConnection) -> Result<LoginResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    _websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<LoginResponse, Error> {
     let data: &DeleteAccount = &self.data;
 
     let claims = match Claims::decode(&data.auth) {
@@ -832,6 +999,15 @@ impl Perform<LoginResponse> for Oper<DeleteAccount> {
 
     let user_id = claims.id;
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     let user: User_ = User_::read(&conn, user_id)?;
 
     // Verify the password
@@ -903,9 +1079,23 @@ impl Perform<LoginResponse> for Oper<DeleteAccount> {
 }
 
 impl Perform<PasswordResetResponse> for Oper<PasswordReset> {
-  fn perform(&self, conn: &PgConnection) -> Result<PasswordResetResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    _websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<PasswordResetResponse, Error> {
     let data: &PasswordReset = &self.data;
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     // Fetch that email
     let user: User_ = match User_::find_by_email(&conn, &data.email) {
       Ok(user) => user,
@@ -934,9 +1124,23 @@ impl Perform<PasswordResetResponse> for Oper<PasswordReset> {
 }
 
 impl Perform<LoginResponse> for Oper<PasswordChange> {
-  fn perform(&self, conn: &PgConnection) -> Result<LoginResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    _websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<LoginResponse, Error> {
     let data: &PasswordChange = &self.data;
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     // Fetch the user_id from the token
     let user_id = PasswordResetRequest::read_from_token(&conn, &data.token)?.user_id;
 
@@ -959,7 +1163,12 @@ impl Perform<LoginResponse> for Oper<PasswordChange> {
 }
 
 impl Perform<PrivateMessageResponse> for Oper<CreatePrivateMessage> {
-  fn perform(&self, conn: &PgConnection) -> Result<PrivateMessageResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<PrivateMessageResponse, Error> {
     let data: &CreatePrivateMessage = &self.data;
 
     let claims = match Claims::decode(&data.auth) {
@@ -971,6 +1180,15 @@ impl Perform<PrivateMessageResponse> for Oper<CreatePrivateMessage> {
 
     let hostname = &format!("https://{}", Settings::get().hostname);
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     // Check for a site ban
     if UserView::read(&conn, user_id)?.banned {
       return Err(APIError::err("site_ban").into());
@@ -1016,12 +1234,28 @@ impl Perform<PrivateMessageResponse> for Oper<CreatePrivateMessage> {
 
     let message = PrivateMessageView::read(&conn, inserted_private_message.id)?;
 
-    Ok(PrivateMessageResponse { message })
+    let res = PrivateMessageResponse { message };
+
+    if let Some(ws) = websocket_info {
+      ws.chatserver.do_send(SendUserRoomMessage {
+        op: UserOperation::CreatePrivateMessage,
+        response: res.clone(),
+        recipient_id: recipient_user.id,
+        my_id: ws.id,
+      });
+    }
+
+    Ok(res)
   }
 }
 
 impl Perform<PrivateMessageResponse> for Oper<EditPrivateMessage> {
-  fn perform(&self, conn: &PgConnection) -> Result<PrivateMessageResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    _websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<PrivateMessageResponse, Error> {
     let data: &EditPrivateMessage = &self.data;
 
     let claims = match Claims::decode(&data.auth) {
@@ -1031,6 +1265,15 @@ impl Perform<PrivateMessageResponse> for Oper<EditPrivateMessage> {
 
     let user_id = claims.id;
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     let orig_private_message = PrivateMessage::read(&conn, data.edit_id)?;
 
     // Check for a site ban
@@ -1076,7 +1319,12 @@ impl Perform<PrivateMessageResponse> for Oper<EditPrivateMessage> {
 }
 
 impl Perform<PrivateMessagesResponse> for Oper<GetPrivateMessages> {
-  fn perform(&self, conn: &PgConnection) -> Result<PrivateMessagesResponse, Error> {
+  fn perform(
+    &self,
+    pool: Pool<ConnectionManager<PgConnection>>,
+    _websocket_info: Option<WebsocketInfo>,
+    rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<PrivateMessagesResponse, Error> {
     let data: &GetPrivateMessages = &self.data;
 
     let claims = match Claims::decode(&data.auth) {
@@ -1086,6 +1334,15 @@ impl Perform<PrivateMessagesResponse> for Oper<GetPrivateMessages> {
 
     let user_id = claims.id;
 
+    if let Some(rl) = rate_limit_info {
+      rl.rate_limiter
+        .lock()
+        .unwrap()
+        .check_rate_limit_message(&rl.ip, false)?;
+    }
+
+    let conn = pool.get()?;
+
     let messages = PrivateMessageQueryBuilder::create(&conn, user_id)
       .page(data.page)
       .limit(data.limit)
@@ -1097,7 +1354,12 @@ impl Perform<PrivateMessagesResponse> for Oper<GetPrivateMessages> {
 }
 
 impl Perform<UserJoinResponse> for Oper<UserJoin> {
-  fn perform(&self, _conn: &PgConnection) -> Result<UserJoinResponse, Error> {
+  fn perform(
+    &self,
+    _pool: Pool<ConnectionManager<PgConnection>>,
+    websocket_info: Option<WebsocketInfo>,
+    _rate_limit_info: Option<RateLimitInfo>,
+  ) -> Result<UserJoinResponse, Error> {
     let data: &UserJoin = &self.data;
 
     let claims = match Claims::decode(&data.auth) {
@@ -1106,6 +1368,13 @@ impl Perform<UserJoinResponse> for Oper<UserJoin> {
     };
 
     let user_id = claims.id;
+
+    if let Some(ws) = websocket_info {
+      if let Some(id) = ws.id {
+        ws.chatserver.do_send(JoinUserRoom { user_id, id });
+      }
+    }
+
     Ok(UserJoinResponse { user_id })
   }
 }

server/src/lib.rs

@@ -27,13 +27,14 @@ pub extern crate strum;
 pub mod api;
 pub mod apub;
 pub mod db;
+pub mod rate_limit;
 pub mod routes;
 pub mod schema;
 pub mod settings;
 pub mod version;
 pub mod websocket;
 
-use crate::settings::Settings;
+use actix_web::HttpRequest;
 use chrono::{DateTime, NaiveDateTime, Utc};
 use isahc::prelude::*;
 use lettre::smtp::authentication::{Credentials, Mechanism};
@@ -48,6 +49,14 @@ use rand::{thread_rng, Rng};
 use regex::{Regex, RegexBuilder};
 use serde::Deserialize;
 
+use crate::settings::Settings;
+
+pub type ConnectionId = usize;
+pub type PostId = i32;
+pub type CommunityId = i32;
+pub type UserId = i32;
+pub type IPAddr = String;
+
 pub fn to_datetime_utc(ndt: NaiveDateTime) -> DateTime<Utc> {
   DateTime::<Utc>::from_utc(ndt, Utc)
 }
@@ -224,6 +233,17 @@ pub fn markdown_to_html(text: &str) -> String {
   comrak::markdown_to_html(text, &comrak::ComrakOptions::default())
 }
 
+pub fn get_ip(req: &HttpRequest) -> String {
+  req
+    .connection_info()
+    .remote()
+    .unwrap_or("127.0.0.1:12345")
+    .split(':')
+    .next()
+    .unwrap_or("127.0.0.1")
+    .to_string()
+}
+
 #[cfg(test)]
 mod tests {
   use crate::{extract_usernames, is_email_regex, remove_slurs, slur_check, slurs_vec_to_str};

server/src/main.rs

@@ -6,10 +6,16 @@ use actix::prelude::*;
 use actix_web::*;
 use diesel::r2d2::{ConnectionManager, Pool};
 use diesel::PgConnection;
-use lemmy_server::routes::{api, federation, feeds, index, nodeinfo, webfinger, websocket};
-use lemmy_server::settings::Settings;
-use lemmy_server::websocket::server::*;
-use std::io;
+use lemmy_server::{
+  rate_limit::rate_limiter::RateLimiter,
+  routes::{api, federation, feeds, index, nodeinfo, webfinger, websocket},
+  settings::Settings,
+  websocket::server::*,
+};
+use std::{
+  io,
+  sync::{Arc, Mutex},
+};
 
 embed_migrations!();
 
@@ -29,8 +35,11 @@ async fn main() -> io::Result<()> {
   let conn = pool.get().unwrap();
   embedded_migrations::run(&conn).unwrap();
 
+  // Set up the rate limiter
+  let rate_limiter = Arc::new(Mutex::new(RateLimiter::default()));
+
   // Set up websocket server
-  let server = ChatServer::startup(pool.clone()).start();
+  let server = ChatServer::startup(pool.clone(), rate_limiter.clone()).start();
 
   println!(
     "Starting http server at {}:{}",
@@ -44,6 +53,7 @@ async fn main() -> io::Result<()> {
       .wrap(middleware::Logger::default())
       .data(pool.clone())
       .data(server.clone())
+      .data(rate_limiter.clone())
       // The routes
       .configure(api::config)
       .configure(federation::config)

server/src/rate_limit/mod.rs

@@ -0,0 +1,18 @@
+pub mod rate_limiter;
+
+use super::{IPAddr, Settings};
+use crate::api::APIError;
+use failure::Error;
+use log::warn;
+use rate_limiter::RateLimiter;
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::sync::Mutex;
+use std::time::SystemTime;
+use strum::IntoEnumIterator;
+
+#[derive(Debug, Clone)]
+pub struct RateLimitInfo {
+  pub rate_limiter: Arc<Mutex<RateLimiter>>,
+  pub ip: IPAddr,
+}

server/src/rate_limit/rate_limiter.rs

@@ -0,0 +1,131 @@
+use super::*;
+
+#[derive(Debug, Clone)]
+pub struct RateLimitBucket {
+  last_checked: SystemTime,
+  allowance: f64,
+}
+
+#[derive(Eq, PartialEq, Hash, Debug, EnumIter, Copy, Clone)]
+pub enum RateLimitType {
+  Message,
+  Register,
+  Post,
+}
+
+/// Rate limiting based on rate type and IP addr
+#[derive(Debug, Clone)]
+pub struct RateLimiter {
+  pub buckets: HashMap<RateLimitType, HashMap<IPAddr, RateLimitBucket>>,
+}
+
+impl Default for RateLimiter {
+  fn default() -> Self {
+    Self {
+      buckets: HashMap::new(),
+    }
+  }
+}
+
+impl RateLimiter {
+  fn insert_ip(&mut self, ip: &str) {
+    for rate_limit_type in RateLimitType::iter() {
+      if self.buckets.get(&rate_limit_type).is_none() {
+        self.buckets.insert(rate_limit_type, HashMap::new());
+      }
+
+      if let Some(bucket) = self.buckets.get_mut(&rate_limit_type) {
+        if bucket.get(ip).is_none() {
+          bucket.insert(
+            ip.to_string(),
+            RateLimitBucket {
+              last_checked: SystemTime::now(),
+              allowance: -2f64,
+            },
+          );
+        }
+      }
+    }
+  }
+
+  pub fn check_rate_limit_register(&mut self, ip: &str, check_only: bool) -> Result<(), Error> {
+    self.check_rate_limit_full(
+      RateLimitType::Register,
+      ip,
+      Settings::get().rate_limit.register,
+      Settings::get().rate_limit.register_per_second,
+      check_only,
+    )
+  }
+
+  pub fn check_rate_limit_post(&mut self, ip: &str, check_only: bool) -> Result<(), Error> {
+    self.check_rate_limit_full(
+      RateLimitType::Post,
+      ip,
+      Settings::get().rate_limit.post,
+      Settings::get().rate_limit.post_per_second,
+      check_only,
+    )
+  }
+
+  pub fn check_rate_limit_message(&mut self, ip: &str, check_only: bool) -> Result<(), Error> {
+    self.check_rate_limit_full(
+      RateLimitType::Message,
+      ip,
+      Settings::get().rate_limit.message,
+      Settings::get().rate_limit.message_per_second,
+      check_only,
+    )
+  }
+
+  #[allow(clippy::float_cmp)]
+  fn check_rate_limit_full(
+    &mut self,
+    type_: RateLimitType,
+    ip: &str,
+    rate: i32,
+    per: i32,
+    check_only: bool,
+  ) -> Result<(), Error> {
+    self.insert_ip(ip);
+    if let Some(bucket) = self.buckets.get_mut(&type_) {
+      if let Some(rate_limit) = bucket.get_mut(ip) {
+        let current = SystemTime::now();
+        let time_passed = current.duration_since(rate_limit.last_checked)?.as_secs() as f64;
+
+        // The initial value
+        if rate_limit.allowance == -2f64 {
+          rate_limit.allowance = rate as f64;
+        };
+
+        rate_limit.last_checked = current;
+        rate_limit.allowance += time_passed * (rate as f64 / per as f64);
+        if !check_only && rate_limit.allowance > rate as f64 {
+          rate_limit.allowance = rate as f64;
+        }
+
+        if rate_limit.allowance < 1.0 {
+          warn!(
+            "Rate limited IP: {}, time_passed: {}, allowance: {}",
+            ip, time_passed, rate_limit.allowance
+          );
+          Err(
+            APIError {
+              message: format!("Too many requests. {} per {} seconds", rate, per),
+            }
+            .into(),
+          )
+        } else {
+          if !check_only {
+            rate_limit.allowance -= 1.0;
+          }
+          Ok(())
+        }
+      } else {
+        Ok(())
+      }
+    } else {
+      Ok(())
+    }
+  }
+}

server/src/routes/api.rs

@@ -1,16 +1,9 @@
+use super::*;
 use crate::api::comment::*;
 use crate::api::community::*;
 use crate::api::post::*;
 use crate::api::site::*;
 use crate::api::user::*;
-use crate::api::{Oper, Perform};
-use actix_web::{web, HttpResponse};
-use diesel::r2d2::{ConnectionManager, Pool};
-use diesel::PgConnection;
-use failure::Error;
-use serde::Serialize;
-
-type DbParam = web::Data<Pool<ConnectionManager<PgConnection>>>;
 
 #[rustfmt::skip]
 pub fn config(cfg: &mut web::ServiceConfig) {
@@ -66,40 +59,64 @@ pub fn config(cfg: &mut web::ServiceConfig) {
     .route("/api/v1/user/save_user_settings", web::put().to(route_post::<SaveUserSettings, LoginResponse>));
 }
 
-fn perform<Request, Response>(data: Request, db: DbParam) -> Result<HttpResponse, Error>
+fn perform<Request, Response>(
+  data: Request,
+  db: DbPoolParam,
+  rate_limit_param: RateLimitParam,
+  chat_server: ChatServerParam,
+  req: HttpRequest,
+) -> Result<HttpResponse, Error>
 where
   Response: Serialize,
   Oper<Request>: Perform<Response>,
 {
-  let conn = match db.get() {
-    Ok(c) => c,
-    Err(e) => return Err(format_err!("{}", e)),
+  let ws_info = WebsocketInfo {
+    chatserver: chat_server.get_ref().to_owned(),
+    id: None,
+  };
+
+  let rate_limit_info = RateLimitInfo {
+    rate_limiter: rate_limit_param.get_ref().to_owned(),
+    ip: get_ip(&req),
   };
+
   let oper: Oper<Request> = Oper::new(data);
-  let response = oper.perform(&conn);
-  Ok(HttpResponse::Ok().json(response?))
+
+  let res = oper.perform(
+    db.get_ref().to_owned(),
+    Some(ws_info),
+    Some(rate_limit_info),
+  );
+
+  Ok(HttpResponse::Ok().json(res?))
 }
 
 async fn route_get<Data, Response>(
   data: web::Query<Data>,
-  db: DbParam,
+  db: DbPoolParam,
+  rate_limit_param: RateLimitParam,
+  chat_server: ChatServerParam,
+  req: HttpRequest,
 ) -> Result<HttpResponse, Error>
 where
   Data: Serialize,
   Response: Serialize,
   Oper<Data>: Perform<Response>,
 {
-  perform::<Data, Response>(data.0, db)
+  perform::<Data, Response>(data.0, db, rate_limit_param, chat_server, req)
 }
 
 async fn route_post<Data, Response>(
   data: web::Json<Data>,
-  db: DbParam,
+  db: DbPoolParam,
+  rate_limit_param: RateLimitParam,
+  chat_server: ChatServerParam,
+  req: HttpRequest,
 ) -> Result<HttpResponse, Error>
 where
   Data: Serialize,
   Response: Serialize,
   Oper<Data>: Perform<Response>,
 {
-  perform::<Data, Response>(data.0, db)
+  perform::<Data, Response>(data.0, db, rate_limit_param, chat_server, req)
 }

server/src/routes/federation.rs

@@ -1,5 +1,5 @@
+use super::*;
 use crate::apub;
-use actix_web::web;
 
 pub fn config(cfg: &mut web::ServiceConfig) {
   cfg

server/src/routes/feeds.rs

@@ -6,16 +6,6 @@ use crate::db::site_view::SiteView;
 use crate::db::user::{Claims, User_};
 use crate::db::user_mention_view::{UserMentionQueryBuilder, UserMentionView};
 use crate::db::{ListingType, SortType};
-use crate::{markdown_to_html, Settings};
-use actix_web::{web, HttpResponse, Result};
-use chrono::{DateTime, NaiveDateTime, Utc};
-use diesel::r2d2::{ConnectionManager, Pool};
-use diesel::PgConnection;
-use failure::Error;
-use rss::{CategoryBuilder, ChannelBuilder, GuidBuilder, Item, ItemBuilder};
-use serde::Deserialize;
-use std::str::FromStr;
-use strum::ParseError;
 
 #[derive(Deserialize)]
 pub struct Params {

server/src/routes/index.rs

@@ -1,6 +1,4 @@
-use crate::settings::Settings;
-use actix_files::NamedFile;
-use actix_web::web;
+use super::*;
 
 pub fn config(cfg: &mut web::ServiceConfig) {
   cfg

server/src/routes/mod.rs

@@ -1,3 +1,32 @@
+use crate::api::{Oper, Perform};
+use crate::db::site_view::SiteView;
+use crate::rate_limit::{rate_limiter::RateLimiter, RateLimitInfo};
+use crate::websocket::{server::ChatServer, WebsocketInfo};
+use crate::{get_ip, markdown_to_html, version, Settings};
+use actix::prelude::*;
+use actix_files::NamedFile;
+use actix_web::{body::Body, web::Query, *};
+use actix_web_actors::ws;
+use chrono::{DateTime, NaiveDateTime, Utc};
+use diesel::{
+  r2d2::{ConnectionManager, Pool},
+  PgConnection,
+};
+use failure::Error;
+use log::{error, info};
+use regex::Regex;
+use rss::{CategoryBuilder, ChannelBuilder, GuidBuilder, Item, ItemBuilder};
+use serde::{Deserialize, Serialize};
+use serde_json::json;
+use std::str::FromStr;
+use std::sync::{Arc, Mutex};
+use std::time::{Duration, Instant};
+use strum::ParseError;
+
+pub type DbPoolParam = web::Data<Pool<ConnectionManager<PgConnection>>>;
+pub type RateLimitParam = web::Data<Arc<Mutex<RateLimiter>>>;
+pub type ChatServerParam = web::Data<Addr<ChatServer>>;
+
 pub mod api;
 pub mod federation;
 pub mod feeds;

server/src/routes/nodeinfo.rs

@@ -1,12 +1,4 @@
-use crate::db::site_view::SiteView;
-use crate::version;
-use crate::Settings;
-use actix_web::body::Body;
-use actix_web::web;
-use actix_web::HttpResponse;
-use diesel::r2d2::{ConnectionManager, Pool};
-use diesel::PgConnection;
-use serde::Serialize;
+use super::*;
 
 pub fn config(cfg: &mut web::ServiceConfig) {
   cfg

server/src/routes/webfinger.rs

@@ -1,13 +1,5 @@
+use super::*;
 use crate::db::community::Community;
-use crate::Settings;
-use actix_web::web;
-use actix_web::web::Query;
-use actix_web::HttpResponse;
-use diesel::r2d2::{ConnectionManager, Pool};
-use diesel::PgConnection;
-use regex::Regex;
-use serde::Deserialize;
-use serde_json::json;
 
 #[derive(Deserialize)]
 pub struct Params {

server/src/routes/websocket.rs

@@ -1,10 +1,6 @@
+use super::*;
 use crate::websocket::server::*;
-use actix::prelude::*;
-use actix_web::web;
-use actix_web::*;
-use actix_web_actors::ws;
-use log::{error, info};
-use std::time::{Duration, Instant};
+use actix_web::{Error, Result};
 
 pub fn config(cfg: &mut web::ServiceConfig) {
   cfg.service(web::resource("/api/v1/ws").to(chat_route));
@@ -21,20 +17,12 @@ async fn chat_route(
   stream: web::Payload,
   chat_server: web::Data<Addr<ChatServer>>,
 ) -> Result<HttpResponse, Error> {
-  // TODO not sure if the blocking should be here or not
   ws::start(
     WSSession {
       cs_addr: chat_server.get_ref().to_owned(),
       id: 0,
       hb: Instant::now(),
-      ip: req
-        .connection_info()
-        .remote()
-        .unwrap_or("127.0.0.1:12345")
-        .split(':')
-        .next()
-        .unwrap_or("127.0.0.1")
-        .to_string(),
+      ip: get_ip(&req),
     },
     &req,
     stream,

server/src/websocket/mod.rs

@@ -1,6 +1,21 @@
 pub mod server;
 
-#[derive(EnumString, ToString, Debug)]
+use crate::ConnectionId;
+use actix::prelude::*;
+use diesel::r2d2::{ConnectionManager, Pool};
+use diesel::PgConnection;
+use failure::Error;
+use log::{error, info};
+use rand::{rngs::ThreadRng, Rng};
+use serde::{Deserialize, Serialize};
+use serde_json::Value;
+use server::ChatServer;
+use std::collections::{HashMap, HashSet};
+use std::str::FromStr;
+use std::sync::Arc;
+use std::sync::Mutex;
+
+#[derive(EnumString, ToString, Debug, Clone)]
 pub enum UserOperation {
   Login,
   Register,
@@ -49,3 +64,9 @@ pub enum UserOperation {
   GetSiteConfig,
   SaveSiteConfig,
 }
+
+#[derive(Clone)]
+pub struct WebsocketInfo {
+  pub chatserver: Addr<ChatServer>,
+  pub id: Option<ConnectionId>,
+}