crates/websocket/src/handlers.rs

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