crates/websocket/src/handlers.rs

   1 use crate::{
   2   chat_server::{ChatServer, SessionInfo},
   3   messages::*,
   4   OperationType,
   5 };
   6 use actix::{Actor, Context, Handler, ResponseFuture};
   7 use lemmy_db_schema::naive_now;
   8 use lemmy_utils::ConnectionId;
   9 use rand::Rng;
  10 use serde::Serialize;
  11 use tracing::{error, info};
  12
  13 /// Make actor from `ChatServer`
  14 impl Actor for ChatServer {
  15   /// We are going to use simple Context, we just need ability to communicate
  16   /// with other actors.
  17   type Context = Context<Self>;
  18 }
  19
  20 /// Handler for Connect message.
  21 ///
  22 /// Register new session and assign unique id to this session
  23 impl Handler<Connect> for ChatServer {
  24   type Result = ConnectionId;
  25
  26   fn handle(&mut self, msg: Connect, _ctx: &mut Context<Self>) -> Self::Result {
  27     // register session with random id
  28     let id = self.rng.gen::<usize>();
  29     info!("{} joined", &msg.ip);
  30
  31     self.sessions.insert(
  32       id,
  33       SessionInfo {
  34         addr: msg.addr,
  35         ip: msg.ip,
  36       },
  37     );
  38
  39     id
  40   }
  41 }
  42
  43 /// Handler for Disconnect message.
  44 impl Handler<Disconnect> for ChatServer {
  45   type Result = ();
  46
  47   fn handle(&mut self, msg: Disconnect, _: &mut Context<Self>) {
  48     // Remove connections from sessions and all 3 scopes
  49     if self.sessions.remove(&msg.id).is_some() {
  50       for sessions in self.user_rooms.values_mut() {
  51         sessions.remove(&msg.id);
  52       }
  53
  54       for sessions in self.post_rooms.values_mut() {
  55         sessions.remove(&msg.id);
  56       }
  57
  58       for sessions in self.community_rooms.values_mut() {
  59         sessions.remove(&msg.id);
  60       }
  61     }
  62   }
  63 }
  64
  65 /// Handler for Message message.
  66 impl Handler<StandardMessage> for ChatServer {
  67   type Result = ResponseFuture<Result<String, std::convert::Infallible>>;
  68
  69   fn handle(&mut self, msg: StandardMessage, ctx: &mut Context<Self>) -> Self::Result {
  70     let fut = self.parse_json_message(msg, ctx);
  71     Box::pin(async move {
  72       match fut.await {
  73         Ok(m) => {
  74           // info!("Message Sent: {}", m);
  75           Ok(m)
  76         }
  77         Err(e) => {
  78           error!("Error during message handling {}", e);
  79           Ok(
  80             e.to_json()
  81               .unwrap_or_else(|_| String::from(r#"{"error":"failed to serialize json"}"#)),
  82           )
  83         }
  84       }
  85     })
  86   }
  87 }
  88
  89 impl<OP, Response> Handler<SendAllMessage<OP, Response>> for ChatServer
  90 where
  91   OP: OperationType + ToString,
  92   Response: Serialize,
  93 {
  94   type Result = ();
  95
  96   fn handle(&mut self, msg: SendAllMessage<OP, Response>, _: &mut Context<Self>) {
  97     self
  98       .send_all_message(&msg.op, &msg.response, msg.websocket_id)
  99       .ok();
 100   }
 101 }
 102
 103 impl<OP, Response> Handler<SendUserRoomMessage<OP, Response>> for ChatServer
 104 where
 105   OP: OperationType + ToString,
 106   Response: Serialize,
 107 {
 108   type Result = ();
 109
 110   fn handle(&mut self, msg: SendUserRoomMessage<OP, Response>, _: &mut Context<Self>) {
 111     self
 112       .send_user_room_message(
 113         &msg.op,
 114         &msg.response,
 115         msg.local_recipient_id,
 116         msg.websocket_id,
 117       )
 118       .ok();
 119   }
 120 }
 121
 122 impl<OP, Response> Handler<SendCommunityRoomMessage<OP, Response>> for ChatServer
 123 where
 124   OP: OperationType + ToString,
 125   Response: Serialize,
 126 {
 127   type Result = ();
 128
 129   fn handle(&mut self, msg: SendCommunityRoomMessage<OP, Response>, _: &mut Context<Self>) {
 130     self
 131       .send_community_room_message(&msg.op, &msg.response, msg.community_id, msg.websocket_id)
 132       .ok();
 133   }
 134 }
 135
 136 impl<Response> Handler<SendModRoomMessage<Response>> for ChatServer
 137 where
 138   Response: Serialize,
 139 {
 140   type Result = ();
 141
 142   fn handle(&mut self, msg: SendModRoomMessage<Response>, _: &mut Context<Self>) {
 143     self
 144       .send_mod_room_message(&msg.op, &msg.response, msg.community_id, msg.websocket_id)
 145       .ok();
 146   }
 147 }
 148
 149 impl<OP> Handler<SendPost<OP>> for ChatServer
 150 where
 151   OP: OperationType + ToString,
 152 {
 153   type Result = ();
 154
 155   fn handle(&mut self, msg: SendPost<OP>, _: &mut Context<Self>) {
 156     self.send_post(&msg.op, &msg.post, msg.websocket_id).ok();
 157   }
 158 }
 159
 160 impl<OP> Handler<SendComment<OP>> for ChatServer
 161 where
 162   OP: OperationType + ToString,
 163 {
 164   type Result = ();
 165
 166   fn handle(&mut self, msg: SendComment<OP>, _: &mut Context<Self>) {
 167     self
 168       .send_comment(&msg.op, &msg.comment, msg.websocket_id)
 169       .ok();
 170   }
 171 }
 172
 173 impl Handler<JoinUserRoom> for ChatServer {
 174   type Result = ();
 175
 176   fn handle(&mut self, msg: JoinUserRoom, _: &mut Context<Self>) {
 177     self.join_user_room(msg.local_user_id, msg.id).ok();
 178   }
 179 }
 180
 181 impl Handler<JoinCommunityRoom> for ChatServer {
 182   type Result = ();
 183
 184   fn handle(&mut self, msg: JoinCommunityRoom, _: &mut Context<Self>) {
 185     self.join_community_room(msg.community_id, msg.id).ok();
 186   }
 187 }
 188
 189 impl Handler<JoinModRoom> for ChatServer {
 190   type Result = ();
 191
 192   fn handle(&mut self, msg: JoinModRoom, _: &mut Context<Self>) {
 193     self.join_mod_room(msg.community_id, msg.id).ok();
 194   }
 195 }
 196
 197 impl Handler<JoinPostRoom> for ChatServer {
 198   type Result = ();
 199
 200   fn handle(&mut self, msg: JoinPostRoom, _: &mut Context<Self>) {
 201     self.join_post_room(msg.post_id, msg.id).ok();
 202   }
 203 }
 204
 205 impl Handler<GetUsersOnline> for ChatServer {
 206   type Result = usize;
 207
 208   fn handle(&mut self, _msg: GetUsersOnline, _: &mut Context<Self>) -> Self::Result {
 209     self.sessions.len()
 210   }
 211 }
 212
 213 impl Handler<GetPostUsersOnline> for ChatServer {
 214   type Result = usize;
 215
 216   fn handle(&mut self, msg: GetPostUsersOnline, _: &mut Context<Self>) -> Self::Result {
 217     if let Some(users) = self.post_rooms.get(&msg.post_id) {
 218       users.len()
 219     } else {
 220       0
 221     }
 222   }
 223 }
 224
 225 impl Handler<GetCommunityUsersOnline> for ChatServer {
 226   type Result = usize;
 227
 228   fn handle(&mut self, msg: GetCommunityUsersOnline, _: &mut Context<Self>) -> Self::Result {
 229     if let Some(users) = self.community_rooms.get(&msg.community_id) {
 230       users.len()
 231     } else {
 232       0
 233     }
 234   }
 235 }
 236
 237 impl Handler<CaptchaItem> for ChatServer {
 238   type Result = ();
 239
 240   fn handle(&mut self, msg: CaptchaItem, _: &mut Context<Self>) {
 241     self.captchas.push(msg);
 242   }
 243 }
 244
 245 impl Handler<CheckCaptcha> for ChatServer {
 246   type Result = bool;
 247
 248   fn handle(&mut self, msg: CheckCaptcha, _: &mut Context<Self>) -> Self::Result {
 249     // Remove all the ones that are past the expire time
 250     self.captchas.retain(|x| x.expires.gt(&naive_now()));
 251
 252     let check = self
 253       .captchas
 254       .iter()
 255       .any(|r| r.uuid == msg.uuid && r.answer.to_lowercase() == msg.answer.to_lowercase());
 256
 257     // Remove this uuid so it can't be re-checked (Checks only work once)
 258     self.captchas.retain(|x| x.uuid != msg.uuid);
 259
 260     check
 261   }
 262 }