src/websocket/handlers.rs

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