Refactor websocket to split it into multiple files

server/src/api/comment.rs

@@ -11,7 +11,7 @@ use crate::{
   apub::{ApubLikeableType, ApubObjectType},
   blocking,
   websocket::{
-    server::{JoinCommunityRoom, SendComment},
+    messages::{JoinCommunityRoom, SendComment},
     UserOperation,
   },
   ConnectionId,

server/src/api/community.rs

@@ -4,7 +4,7 @@ use crate::{
   apub::ActorType,
   blocking,
   websocket::{
-    server::{GetCommunityUsersOnline, JoinCommunityRoom, SendCommunityRoomMessage},
+    messages::{GetCommunityUsersOnline, JoinCommunityRoom, SendCommunityRoomMessage},
     UserOperation,
   },
   ConnectionId,

server/src/api/post.rs

@@ -13,7 +13,7 @@ use crate::{
   blocking,
   fetch_iframely_and_pictrs_data,
   websocket::{
-    server::{GetPostUsersOnline, JoinCommunityRoom, JoinPostRoom, SendPost},
+    messages::{GetPostUsersOnline, JoinCommunityRoom, JoinPostRoom, SendPost},
     UserOperation,
   },
   ConnectionId,

server/src/api/site.rs

@@ -13,7 +13,7 @@ use crate::{
   blocking,
   version,
   websocket::{
-    server::{GetUsersOnline, SendAllMessage},
+    messages::{GetUsersOnline, SendAllMessage},
     UserOperation,
   },
   ConnectionId,

server/src/api/user.rs

@@ -12,7 +12,7 @@ use crate::{
   blocking,
   captcha_espeak_wav_base64,
   websocket::{
-    server::{CaptchaItem, CheckCaptcha, JoinUserRoom, SendAllMessage, SendUserRoomMessage},
+    messages::{CaptchaItem, CheckCaptcha, JoinUserRoom, SendAllMessage, SendUserRoomMessage},
     UserOperation,
   },
   ConnectionId,

server/src/apub/inbox/activities/create.rs

@@ -15,7 +15,7 @@ use crate::{
   },
   blocking,
   websocket::{
-    server::{SendComment, SendPost},
+    messages::{SendComment, SendPost},
     UserOperation,
   },
   LemmyContext,

server/src/apub/inbox/activities/delete.rs

@@ -14,7 +14,7 @@ use crate::{
   },
   blocking,
   websocket::{
-    server::{SendComment, SendCommunityRoomMessage, SendPost},
+    messages::{SendComment, SendCommunityRoomMessage, SendPost},
     UserOperation,
   },
   LemmyContext,

server/src/apub/inbox/activities/dislike.rs

@@ -12,7 +12,7 @@ use crate::{
   },
   blocking,
   websocket::{
-    server::{SendComment, SendPost},
+    messages::{SendComment, SendPost},
     UserOperation,
   },
   LemmyContext,

server/src/apub/inbox/activities/like.rs

@@ -12,7 +12,7 @@ use crate::{
   },
   blocking,
   websocket::{
-    server::{SendComment, SendPost},
+    messages::{SendComment, SendPost},
     UserOperation,
   },
   LemmyContext,

server/src/apub/inbox/activities/remove.rs

@@ -15,7 +15,7 @@ use crate::{
   },
   blocking,
   websocket::{
-    server::{SendComment, SendCommunityRoomMessage, SendPost},
+    messages::{SendComment, SendCommunityRoomMessage, SendPost},
     UserOperation,
   },
   LemmyContext,

server/src/apub/inbox/activities/undo.rs

@@ -14,7 +14,7 @@ use crate::{
   },
   blocking,
   websocket::{
-    server::{SendComment, SendCommunityRoomMessage, SendPost},
+    messages::{SendComment, SendCommunityRoomMessage, SendPost},
     UserOperation,
   },
   LemmyContext,

server/src/apub/inbox/activities/update.rs

@@ -16,7 +16,7 @@ use crate::{
   },
   blocking,
   websocket::{
-    server::{SendComment, SendPost},
+    messages::{SendComment, SendPost},
     UserOperation,
   },
   LemmyContext,

server/src/apub/inbox/user_inbox.rs

@@ -8,7 +8,7 @@ use crate::{
     FromApub,
   },
   blocking,
-  websocket::{server::SendUserRoomMessage, UserOperation},
+  websocket::{messages::SendUserRoomMessage, UserOperation},
   LemmyContext,
   LemmyError,
 };

server/src/lib.rs

@@ -30,11 +30,9 @@ pub mod routes;
 pub mod version;
 pub mod websocket;
 
-use crate::{
+use crate::request::{retry, RecvError};
-  request::{retry, RecvError},
-  websocket::server::ChatServer,
-};
 
+use crate::websocket::chat_server::ChatServer;
 use actix::Addr;
 use actix_web::dev::ConnectionInfo;
 use anyhow::anyhow;

server/src/main.rs

@@ -24,7 +24,7 @@ use lemmy_server::{
   code_migrations::run_advanced_migrations,
   rate_limit::{rate_limiter::RateLimiter, RateLimit},
   routes::*,
-  websocket::server::*,
+  websocket::chat_server::ChatServer,
   LemmyContext,
   LemmyError,
 };

server/src/routes/websocket.rs

@@ -1,6 +1,9 @@
 use crate::{
   get_ip,
-  websocket::server::{ChatServer, *},
+  websocket::{
+    chat_server::ChatServer,
+    messages::{Connect, Disconnect, StandardMessage, WSMessage},
+  },
   LemmyContext,
 };
 use actix::prelude::*;

server/src/websocket/chat_server.rs

@@ -1,12 +1,12 @@
-//! `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 super::*;
 use crate::{
   api::{comment::*, community::*, post::*, site::*, user::*, *},
   rate_limit::RateLimit,
-  websocket::UserOperation,
+  websocket::{
+    handlers::{do_user_operation, to_json_string, Args},
+    messages::*,
+    UserOperation,
+  },
   CommunityId,
   ConnectionId,
   IPAddr,
@@ -15,145 +15,11 @@ use crate::{
   PostId,
   UserId,
 };
-use actix_web::web;
 use anyhow::Context as acontext;
 use background_jobs::QueueHandle;
-use lemmy_db::naive_now;
 use lemmy_utils::location_info;
 use reqwest::Client;
 
-/// Chat server sends this messages to session
-#[derive(Message)]
-#[rtype(result = "()")]
-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: IPAddr,
-}
-
-/// Session is disconnected
-#[derive(Message)]
-#[rtype(result = "()")]
-pub struct Disconnect {
-  pub id: ConnectionId,
-  pub ip: IPAddr,
-}
-
-/// The messages sent to websocket clients
-#[derive(Serialize, Deserialize, Message)]
-#[rtype(result = "Result<String, std::convert::Infallible>")]
-pub struct StandardMessage {
-  /// Id of the client session
-  pub id: ConnectionId,
-  /// Peer message
-  pub msg: String,
-}
-
-#[derive(Message)]
-#[rtype(result = "()")]
-pub struct SendAllMessage<Response> {
-  pub op: UserOperation,
-  pub response: Response,
-  pub websocket_id: Option<ConnectionId>,
-}
-
-#[derive(Message)]
-#[rtype(result = "()")]
-pub struct SendUserRoomMessage<Response> {
-  pub op: UserOperation,
-  pub response: Response,
-  pub recipient_id: UserId,
-  pub websocket_id: Option<ConnectionId>,
-}
-
-#[derive(Message)]
-#[rtype(result = "()")]
-pub struct SendCommunityRoomMessage<Response> {
-  pub op: UserOperation,
-  pub response: Response,
-  pub community_id: CommunityId,
-  pub websocket_id: Option<ConnectionId>,
-}
-
-#[derive(Message)]
-#[rtype(result = "()")]
-pub struct SendPost {
-  pub op: UserOperation,
-  pub post: PostResponse,
-  pub websocket_id: Option<ConnectionId>,
-}
-
-#[derive(Message)]
-#[rtype(result = "()")]
-pub struct SendComment {
-  pub op: UserOperation,
-  pub comment: CommentResponse,
-  pub websocket_id: Option<ConnectionId>,
-}
-
-#[derive(Message)]
-#[rtype(result = "()")]
-pub struct JoinUserRoom {
-  pub user_id: UserId,
-  pub id: ConnectionId,
-}
-
-#[derive(Message)]
-#[rtype(result = "()")]
-pub struct JoinCommunityRoom {
-  pub community_id: CommunityId,
-  pub id: ConnectionId,
-}
-
-#[derive(Message)]
-#[rtype(result = "()")]
-pub struct JoinPostRoom {
-  pub post_id: PostId,
-  pub id: ConnectionId,
-}
-
-#[derive(Message)]
-#[rtype(usize)]
-pub struct GetUsersOnline;
-
-#[derive(Message)]
-#[rtype(usize)]
-pub struct GetPostUsersOnline {
-  pub post_id: PostId,
-}
-
-#[derive(Message)]
-#[rtype(usize)]
-pub struct GetCommunityUsersOnline {
-  pub community_id: CommunityId,
-}
-
-pub struct SessionInfo {
-  pub addr: Recipient<WSMessage>,
-  pub ip: IPAddr,
-}
-
-#[derive(Message, Debug)]
-#[rtype(result = "()")]
-pub struct CaptchaItem {
-  pub uuid: String,
-  pub answer: String,
-  pub expires: chrono::NaiveDateTime,
-}
-
-#[derive(Message)]
-#[rtype(bool)]
-pub struct CheckCaptcha {
-  pub uuid: String,
-  pub answer: String,
-}
-
 /// `ChatServer` manages chat rooms and responsible for coordinating chat
 /// session.
 pub struct ChatServer {
@@ -168,18 +34,18 @@ pub struct ChatServer {
 
   /// A map from user id to its connection ID for joined users. Remember a user can have multiple
   /// sessions (IE clients)
-  user_rooms: HashMap<UserId, HashSet<ConnectionId>>,
+  pub(super) user_rooms: HashMap<UserId, HashSet<ConnectionId>>,
 
-  rng: ThreadRng,
+  pub(super) rng: ThreadRng,
 
   /// The DB Pool
-  pool: Pool<ConnectionManager<PgConnection>>,
+  pub(super) pool: Pool<ConnectionManager<PgConnection>>,
 
   /// Rate limiting based on rate type and IP addr
-  rate_limiter: RateLimit,
+  pub(super) rate_limiter: RateLimit,
 
   /// A list of the current captchas
-  captchas: Vec<CaptchaItem>,
+  pub(super) captchas: Vec<CaptchaItem>,
 
   /// An HTTP Client
   client: Client,
@@ -187,6 +53,14 @@ pub struct ChatServer {
   activity_queue: QueueHandle,
 }
 
+pub struct SessionInfo {
+  pub addr: Recipient<WSMessage>,
+  pub ip: IPAddr,
+}
+
+/// `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`.
 impl ChatServer {
   pub fn startup(
     pool: Pool<ConnectionManager<PgConnection>>,
@@ -451,7 +325,7 @@ impl ChatServer {
     }
   }
 
-  fn parse_json_message(
+  pub(super) fn parse_json_message(
     &mut self,
     msg: StandardMessage,
     ctx: &mut Context<Self>,
@@ -576,275 +450,3 @@ impl ChatServer {
     }
   }
 }
-
-struct Args<'a> {
-  context: LemmyContext,
-  rate_limiter: RateLimit,
-  id: ConnectionId,
-  ip: IPAddr,
-  op: UserOperation,
-  data: &'a str,
-}
-
-async fn do_user_operation<'a, 'b, Data>(args: Args<'b>) -> Result<String, LemmyError>
-where
-  for<'de> Data: Deserialize<'de> + 'a,
-  Data: Perform,
-{
-  let Args {
-    context,
-    rate_limiter,
-    id,
-    ip,
-    op,
-    data,
-  } = args;
-
-  let data = data.to_string();
-  let op2 = op.clone();
-
-  let fut = async move {
-    let parsed_data: Data = serde_json::from_str(&data)?;
-    let res = parsed_data
-      .perform(&web::Data::new(context), Some(id))
-      .await?;
-    to_json_string(&op, &res)
-  };
-
-  match op2 {
-    UserOperation::Register => rate_limiter.register().wrap(ip, fut).await,
-    UserOperation::CreatePost => rate_limiter.post().wrap(ip, fut).await,
-    UserOperation::CreateCommunity => rate_limiter.register().wrap(ip, fut).await,
-    _ => rate_limiter.message().wrap(ip, fut).await,
-  }
-}
-
-/// 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 {
-    // register session with random id
-    let id = self.rng.gen::<usize>();
-    info!("{} joined", &msg.ip);
-
-    self.sessions.insert(
-      id,
-      SessionInfo {
-        addr: msg.addr,
-        ip: msg.ip,
-      },
-    );
-
-    id
-  }
-}
-
-/// Handler for Disconnect message.
-impl Handler<Disconnect> for ChatServer {
-  type Result = ();
-
-  fn handle(&mut self, msg: Disconnect, _: &mut Context<Self>) {
-    // Remove connections from sessions and all 3 scopes
-    if self.sessions.remove(&msg.id).is_some() {
-      for sessions in self.user_rooms.values_mut() {
-        sessions.remove(&msg.id);
-      }
-
-      for sessions in self.post_rooms.values_mut() {
-        sessions.remove(&msg.id);
-      }
-
-      for sessions in self.community_rooms.values_mut() {
-        sessions.remove(&msg.id);
-      }
-    }
-  }
-}
-
-/// Handler for Message message.
-impl Handler<StandardMessage> for ChatServer {
-  type Result = ResponseFuture<Result<String, std::convert::Infallible>>;
-
-  fn handle(&mut self, msg: StandardMessage, ctx: &mut Context<Self>) -> Self::Result {
-    let fut = self.parse_json_message(msg, ctx);
-    Box::pin(async move {
-      match fut.await {
-        Ok(m) => {
-          // info!("Message Sent: {}", m);
-          Ok(m)
-        }
-        Err(e) => {
-          error!("Error during message handling {}", e);
-          Ok(e.to_string())
-        }
-      }
-    })
-  }
-}
-
-impl<Response> Handler<SendAllMessage<Response>> for ChatServer
-where
-  Response: Serialize,
-{
-  type Result = ();
-
-  fn handle(&mut self, msg: SendAllMessage<Response>, _: &mut Context<Self>) {
-    self
-      .send_all_message(&msg.op, &msg.response, msg.websocket_id)
-      .ok();
-  }
-}
-
-impl<Response> Handler<SendUserRoomMessage<Response>> for ChatServer
-where
-  Response: Serialize,
-{
-  type Result = ();
-
-  fn handle(&mut self, msg: SendUserRoomMessage<Response>, _: &mut Context<Self>) {
-    self
-      .send_user_room_message(&msg.op, &msg.response, msg.recipient_id, msg.websocket_id)
-      .ok();
-  }
-}
-
-impl<Response> Handler<SendCommunityRoomMessage<Response>> for ChatServer
-where
-  Response: Serialize,
-{
-  type Result = ();
-
-  fn handle(&mut self, msg: SendCommunityRoomMessage<Response>, _: &mut Context<Self>) {
-    self
-      .send_community_room_message(&msg.op, &msg.response, msg.community_id, msg.websocket_id)
-      .ok();
-  }
-}
-
-impl Handler<SendPost> for ChatServer {
-  type Result = ();
-
-  fn handle(&mut self, msg: SendPost, _: &mut Context<Self>) {
-    self.send_post(&msg.op, &msg.post, msg.websocket_id).ok();
-  }
-}
-
-impl Handler<SendComment> for ChatServer {
-  type Result = ();
-
-  fn handle(&mut self, msg: SendComment, _: &mut Context<Self>) {
-    self
-      .send_comment(&msg.op, &msg.comment, msg.websocket_id)
-      .ok();
-  }
-}
-
-impl Handler<JoinUserRoom> for ChatServer {
-  type Result = ();
-
-  fn handle(&mut self, msg: JoinUserRoom, _: &mut Context<Self>) {
-    self.join_user_room(msg.user_id, msg.id).ok();
-  }
-}
-
-impl Handler<JoinCommunityRoom> for ChatServer {
-  type Result = ();
-
-  fn handle(&mut self, msg: JoinCommunityRoom, _: &mut Context<Self>) {
-    self.join_community_room(msg.community_id, msg.id).ok();
-  }
-}
-
-impl Handler<JoinPostRoom> for ChatServer {
-  type Result = ();
-
-  fn handle(&mut self, msg: JoinPostRoom, _: &mut Context<Self>) {
-    self.join_post_room(msg.post_id, msg.id).ok();
-  }
-}
-
-impl Handler<GetUsersOnline> for ChatServer {
-  type Result = usize;
-
-  fn handle(&mut self, _msg: GetUsersOnline, _: &mut Context<Self>) -> Self::Result {
-    self.sessions.len()
-  }
-}
-
-impl Handler<GetPostUsersOnline> for ChatServer {
-  type Result = usize;
-
-  fn handle(&mut self, msg: GetPostUsersOnline, _: &mut Context<Self>) -> Self::Result {
-    if let Some(users) = self.post_rooms.get(&msg.post_id) {
-      users.len()
-    } else {
-      0
-    }
-  }
-}
-
-impl Handler<GetCommunityUsersOnline> for ChatServer {
-  type Result = usize;
-
-  fn handle(&mut self, msg: GetCommunityUsersOnline, _: &mut Context<Self>) -> Self::Result {
-    if let Some(users) = self.community_rooms.get(&msg.community_id) {
-      users.len()
-    } else {
-      0
-    }
-  }
-}
-
-#[derive(Serialize)]
-struct WebsocketResponse<T> {
-  op: String,
-  data: T,
-}
-
-fn to_json_string<Response>(op: &UserOperation, data: &Response) -> Result<String, LemmyError>
-where
-  Response: Serialize,
-{
-  let response = WebsocketResponse {
-    op: op.to_string(),
-    data,
-  };
-  Ok(serde_json::to_string(&response)?)
-}
-
-impl Handler<CaptchaItem> for ChatServer {
-  type Result = ();
-
-  fn handle(&mut self, msg: CaptchaItem, _: &mut Context<Self>) {
-    self.captchas.push(msg);
-  }
-}
-
-impl Handler<CheckCaptcha> for ChatServer {
-  type Result = bool;
-
-  fn handle(&mut self, msg: CheckCaptcha, _: &mut Context<Self>) -> Self::Result {
-    // Remove all the ones that are past the expire time
-    self.captchas.retain(|x| x.expires.gt(&naive_now()));
-
-    let check = self
-      .captchas
-      .iter()
-      .any(|r| r.uuid == msg.uuid && r.answer == msg.answer);
-
-    // Remove this uuid so it can't be re-checked (Checks only work once)
-    self.captchas.retain(|x| x.uuid != msg.uuid);
-
-    check
-  }
-}

server/src/websocket/handlers.rs

@@ -0,0 +1,291 @@
+use super::*;
+use crate::{
+  api::Perform,
+  rate_limit::RateLimit,
+  websocket::{
+    chat_server::{ChatServer, SessionInfo},
+    messages::*,
+    UserOperation,
+  },
+  ConnectionId,
+  IPAddr,
+  LemmyContext,
+  LemmyError,
+};
+use actix_web::web;
+use lemmy_db::naive_now;
+
+pub(super) struct Args<'a> {
+  pub(super) context: LemmyContext,
+  pub(super) rate_limiter: RateLimit,
+  pub(super) id: ConnectionId,
+  pub(super) ip: IPAddr,
+  pub(super) op: UserOperation,
+  pub(super) data: &'a str,
+}
+
+pub(super) async fn do_user_operation<'a, 'b, Data>(args: Args<'b>) -> Result<String, LemmyError>
+where
+  for<'de> Data: Deserialize<'de> + 'a,
+  Data: Perform,
+{
+  let Args {
+    context,
+    rate_limiter,
+    id,
+    ip,
+    op,
+    data,
+  } = args;
+
+  let data = data.to_string();
+  let op2 = op.clone();
+
+  let fut = async move {
+    let parsed_data: Data = serde_json::from_str(&data)?;
+    let res = parsed_data
+      .perform(&web::Data::new(context), Some(id))
+      .await?;
+    to_json_string(&op, &res)
+  };
+
+  match op2 {
+    UserOperation::Register => rate_limiter.register().wrap(ip, fut).await,
+    UserOperation::CreatePost => rate_limiter.post().wrap(ip, fut).await,
+    UserOperation::CreateCommunity => rate_limiter.register().wrap(ip, fut).await,
+    _ => rate_limiter.message().wrap(ip, fut).await,
+  }
+}
+
+/// 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 {
+    // register session with random id
+    let id = self.rng.gen::<usize>();
+    info!("{} joined", &msg.ip);
+
+    self.sessions.insert(
+      id,
+      SessionInfo {
+        addr: msg.addr,
+        ip: msg.ip,
+      },
+    );
+
+    id
+  }
+}
+
+/// Handler for Disconnect message.
+impl Handler<Disconnect> for ChatServer {
+  type Result = ();
+
+  fn handle(&mut self, msg: Disconnect, _: &mut Context<Self>) {
+    // Remove connections from sessions and all 3 scopes
+    if self.sessions.remove(&msg.id).is_some() {
+      for sessions in self.user_rooms.values_mut() {
+        sessions.remove(&msg.id);
+      }
+
+      for sessions in self.post_rooms.values_mut() {
+        sessions.remove(&msg.id);
+      }
+
+      for sessions in self.community_rooms.values_mut() {
+        sessions.remove(&msg.id);
+      }
+    }
+  }
+}
+
+/// Handler for Message message.
+impl Handler<StandardMessage> for ChatServer {
+  type Result = ResponseFuture<Result<String, std::convert::Infallible>>;
+
+  fn handle(&mut self, msg: StandardMessage, ctx: &mut Context<Self>) -> Self::Result {
+    let fut = self.parse_json_message(msg, ctx);
+    Box::pin(async move {
+      match fut.await {
+        Ok(m) => {
+          // info!("Message Sent: {}", m);
+          Ok(m)
+        }
+        Err(e) => {
+          error!("Error during message handling {}", e);
+          Ok(e.to_string())
+        }
+      }
+    })
+  }
+}
+
+impl<Response> Handler<SendAllMessage<Response>> for ChatServer
+where
+  Response: Serialize,
+{
+  type Result = ();
+
+  fn handle(&mut self, msg: SendAllMessage<Response>, _: &mut Context<Self>) {
+    self
+      .send_all_message(&msg.op, &msg.response, msg.websocket_id)
+      .ok();
+  }
+}
+
+impl<Response> Handler<SendUserRoomMessage<Response>> for ChatServer
+where
+  Response: Serialize,
+{
+  type Result = ();
+
+  fn handle(&mut self, msg: SendUserRoomMessage<Response>, _: &mut Context<Self>) {
+    self
+      .send_user_room_message(&msg.op, &msg.response, msg.recipient_id, msg.websocket_id)
+      .ok();
+  }
+}
+
+impl<Response> Handler<SendCommunityRoomMessage<Response>> for ChatServer
+where
+  Response: Serialize,
+{
+  type Result = ();
+
+  fn handle(&mut self, msg: SendCommunityRoomMessage<Response>, _: &mut Context<Self>) {
+    self
+      .send_community_room_message(&msg.op, &msg.response, msg.community_id, msg.websocket_id)
+      .ok();
+  }
+}
+
+impl Handler<SendPost> for ChatServer {
+  type Result = ();
+
+  fn handle(&mut self, msg: SendPost, _: &mut Context<Self>) {
+    self.send_post(&msg.op, &msg.post, msg.websocket_id).ok();
+  }
+}
+
+impl Handler<SendComment> for ChatServer {
+  type Result = ();
+
+  fn handle(&mut self, msg: SendComment, _: &mut Context<Self>) {
+    self
+      .send_comment(&msg.op, &msg.comment, msg.websocket_id)
+      .ok();
+  }
+}
+
+impl Handler<JoinUserRoom> for ChatServer {
+  type Result = ();
+
+  fn handle(&mut self, msg: JoinUserRoom, _: &mut Context<Self>) {
+    self.join_user_room(msg.user_id, msg.id).ok();
+  }
+}
+
+impl Handler<JoinCommunityRoom> for ChatServer {
+  type Result = ();
+
+  fn handle(&mut self, msg: JoinCommunityRoom, _: &mut Context<Self>) {
+    self.join_community_room(msg.community_id, msg.id).ok();
+  }
+}
+
+impl Handler<JoinPostRoom> for ChatServer {
+  type Result = ();
+
+  fn handle(&mut self, msg: JoinPostRoom, _: &mut Context<Self>) {
+    self.join_post_room(msg.post_id, msg.id).ok();
+  }
+}
+
+impl Handler<GetUsersOnline> for ChatServer {
+  type Result = usize;
+
+  fn handle(&mut self, _msg: GetUsersOnline, _: &mut Context<Self>) -> Self::Result {
+    self.sessions.len()
+  }
+}
+
+impl Handler<GetPostUsersOnline> for ChatServer {
+  type Result = usize;
+
+  fn handle(&mut self, msg: GetPostUsersOnline, _: &mut Context<Self>) -> Self::Result {
+    if let Some(users) = self.post_rooms.get(&msg.post_id) {
+      users.len()
+    } else {
+      0
+    }
+  }
+}
+
+impl Handler<GetCommunityUsersOnline> for ChatServer {
+  type Result = usize;
+
+  fn handle(&mut self, msg: GetCommunityUsersOnline, _: &mut Context<Self>) -> Self::Result {
+    if let Some(users) = self.community_rooms.get(&msg.community_id) {
+      users.len()
+    } else {
+      0
+    }
+  }
+}
+
+#[derive(Serialize)]
+struct WebsocketResponse<T> {
+  op: String,
+  data: T,
+}
+
+pub(super) fn to_json_string<Response>(
+  op: &UserOperation,
+  data: &Response,
+) -> Result<String, LemmyError>
+where
+  Response: Serialize,
+{
+  let response = WebsocketResponse {
+    op: op.to_string(),
+    data,
+  };
+  Ok(serde_json::to_string(&response)?)
+}
+
+impl Handler<CaptchaItem> for ChatServer {
+  type Result = ();
+
+  fn handle(&mut self, msg: CaptchaItem, _: &mut Context<Self>) {
+    self.captchas.push(msg);
+  }
+}
+
+impl Handler<CheckCaptcha> for ChatServer {
+  type Result = bool;
+
+  fn handle(&mut self, msg: CheckCaptcha, _: &mut Context<Self>) -> Self::Result {
+    // Remove all the ones that are past the expire time
+    self.captchas.retain(|x| x.expires.gt(&naive_now()));
+
+    let check = self
+      .captchas
+      .iter()
+      .any(|r| r.uuid == msg.uuid && r.answer == msg.answer);
+
+    // Remove this uuid so it can't be re-checked (Checks only work once)
+    self.captchas.retain(|x| x.uuid != msg.uuid);
+
+    check
+  }
+}

server/src/websocket/messages.rs

@@ -0,0 +1,137 @@
+use super::*;
+use crate::{
+  api::{comment::*, post::*},
+  websocket::UserOperation,
+  CommunityId,
+  ConnectionId,
+  IPAddr,
+  PostId,
+  UserId,
+};
+
+/// Chat server sends this messages to session
+#[derive(Message)]
+#[rtype(result = "()")]
+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: IPAddr,
+}
+
+/// Session is disconnected
+#[derive(Message)]
+#[rtype(result = "()")]
+pub struct Disconnect {
+  pub id: ConnectionId,
+  pub ip: IPAddr,
+}
+
+/// The messages sent to websocket clients
+#[derive(Serialize, Deserialize, Message)]
+#[rtype(result = "Result<String, std::convert::Infallible>")]
+pub struct StandardMessage {
+  /// Id of the client session
+  pub id: ConnectionId,
+  /// Peer message
+  pub msg: String,
+}
+
+#[derive(Message)]
+#[rtype(result = "()")]
+pub struct SendAllMessage<Response> {
+  pub op: UserOperation,
+  pub response: Response,
+  pub websocket_id: Option<ConnectionId>,
+}
+
+#[derive(Message)]
+#[rtype(result = "()")]
+pub struct SendUserRoomMessage<Response> {
+  pub op: UserOperation,
+  pub response: Response,
+  pub recipient_id: UserId,
+  pub websocket_id: Option<ConnectionId>,
+}
+
+#[derive(Message)]
+#[rtype(result = "()")]
+pub struct SendCommunityRoomMessage<Response> {
+  pub op: UserOperation,
+  pub response: Response,
+  pub community_id: CommunityId,
+  pub websocket_id: Option<ConnectionId>,
+}
+
+#[derive(Message)]
+#[rtype(result = "()")]
+pub struct SendPost {
+  pub op: UserOperation,
+  pub post: PostResponse,
+  pub websocket_id: Option<ConnectionId>,
+}
+
+#[derive(Message)]
+#[rtype(result = "()")]
+pub struct SendComment {
+  pub op: UserOperation,
+  pub comment: CommentResponse,
+  pub websocket_id: Option<ConnectionId>,
+}
+
+#[derive(Message)]
+#[rtype(result = "()")]
+pub struct JoinUserRoom {
+  pub user_id: UserId,
+  pub id: ConnectionId,
+}
+
+#[derive(Message)]
+#[rtype(result = "()")]
+pub struct JoinCommunityRoom {
+  pub community_id: CommunityId,
+  pub id: ConnectionId,
+}
+
+#[derive(Message)]
+#[rtype(result = "()")]
+pub struct JoinPostRoom {
+  pub post_id: PostId,
+  pub id: ConnectionId,
+}
+
+#[derive(Message)]
+#[rtype(usize)]
+pub struct GetUsersOnline;
+
+#[derive(Message)]
+#[rtype(usize)]
+pub struct GetPostUsersOnline {
+  pub post_id: PostId,
+}
+
+#[derive(Message)]
+#[rtype(usize)]
+pub struct GetCommunityUsersOnline {
+  pub community_id: CommunityId,
+}
+
+#[derive(Message, Debug)]
+#[rtype(result = "()")]
+pub struct CaptchaItem {
+  pub uuid: String,
+  pub answer: String,
+  pub expires: chrono::NaiveDateTime,
+}
+
+#[derive(Message)]
+#[rtype(bool)]
+pub struct CheckCaptcha {
+  pub uuid: String,
+  pub answer: String,
+}

server/src/websocket/mod.rs

@@ -1,4 +1,6 @@
-pub mod server;
+pub mod chat_server;
+pub mod handlers;
+pub mod messages;
 
 use actix::prelude::*;
 use diesel::{