3 serialize_websocket_message,
10 use anyhow::Context as acontext;
11 use background_jobs::QueueHandle;
13 r2d2::{ConnectionManager, Pool},
16 use lemmy_api_common::{comment::*, post::*};
17 use lemmy_db_schema::{CommunityId, LocalUserId, PostId};
20 rate_limit::RateLimit,
26 use rand::rngs::ThreadRng;
29 use serde_json::Value;
31 collections::{HashMap, HashSet},
35 use tokio::macros::support::Pin;
37 type MessageHandlerType = fn(
38 context: LemmyContext,
42 ) -> Pin<Box<dyn Future<Output = Result<String, LemmyError>> + '_>>;
44 type MessageHandlerCrudType = fn(
45 context: LemmyContext,
47 op: UserOperationCrud,
49 ) -> Pin<Box<dyn Future<Output = Result<String, LemmyError>> + '_>>;
51 /// `ChatServer` manages chat rooms and responsible for coordinating chat
53 pub struct ChatServer {
54 /// A map from generated random ID to session addr
55 pub sessions: HashMap<ConnectionId, SessionInfo>,
57 /// A map from post_id to set of connectionIDs
58 pub post_rooms: HashMap<PostId, HashSet<ConnectionId>>,
60 /// A map from community to set of connectionIDs
61 pub community_rooms: HashMap<CommunityId, HashSet<ConnectionId>>,
63 pub mod_rooms: HashMap<CommunityId, HashSet<ConnectionId>>,
65 /// A map from user id to its connection ID for joined users. Remember a user can have multiple
66 /// sessions (IE clients)
67 pub(super) user_rooms: HashMap<LocalUserId, HashSet<ConnectionId>>,
69 pub(super) rng: ThreadRng,
72 pub(super) pool: Pool<ConnectionManager<PgConnection>>,
74 /// Rate limiting based on rate type and IP addr
75 pub(super) rate_limiter: RateLimit,
77 /// A list of the current captchas
78 pub(super) captchas: Vec<CaptchaItem>,
80 message_handler: MessageHandlerType,
81 message_handler_crud: MessageHandlerCrudType,
86 activity_queue: QueueHandle,
89 pub struct SessionInfo {
90 pub addr: Recipient<WsMessage>,
94 /// `ChatServer` is an actor. It maintains list of connection client session.
95 /// And manages available rooms. Peers send messages to other peers in same
96 /// room through `ChatServer`.
99 pool: Pool<ConnectionManager<PgConnection>>,
100 rate_limiter: RateLimit,
101 message_handler: MessageHandlerType,
102 message_handler_crud: MessageHandlerCrudType,
104 activity_queue: QueueHandle,
107 sessions: HashMap::new(),
108 post_rooms: HashMap::new(),
109 community_rooms: HashMap::new(),
110 mod_rooms: HashMap::new(),
111 user_rooms: HashMap::new(),
112 rng: rand::thread_rng(),
115 captchas: Vec::new(),
117 message_handler_crud,
123 pub fn join_community_room(
125 community_id: CommunityId,
127 ) -> Result<(), LemmyError> {
128 // remove session from all rooms
129 for sessions in self.community_rooms.values_mut() {
130 sessions.remove(&id);
133 // Also leave all post rooms
134 // This avoids double messages
135 for sessions in self.post_rooms.values_mut() {
136 sessions.remove(&id);
139 // If the room doesn't exist yet
140 if self.community_rooms.get_mut(&community_id).is_none() {
141 self.community_rooms.insert(community_id, HashSet::new());
146 .get_mut(&community_id)
147 .context(location_info!())?
152 pub fn join_mod_room(
154 community_id: CommunityId,
156 ) -> Result<(), LemmyError> {
157 // remove session from all rooms
158 for sessions in self.mod_rooms.values_mut() {
159 sessions.remove(&id);
162 // If the room doesn't exist yet
163 if self.mod_rooms.get_mut(&community_id).is_none() {
164 self.mod_rooms.insert(community_id, HashSet::new());
169 .get_mut(&community_id)
170 .context(location_info!())?
175 pub fn join_post_room(&mut self, post_id: PostId, id: ConnectionId) -> Result<(), LemmyError> {
176 // remove session from all rooms
177 for sessions in self.post_rooms.values_mut() {
178 sessions.remove(&id);
181 // Also leave all communities
182 // This avoids double messages
183 // TODO found a bug, whereby community messages like
184 // delete and remove aren't sent, because
185 // you left the community room
186 for sessions in self.community_rooms.values_mut() {
187 sessions.remove(&id);
190 // If the room doesn't exist yet
191 if self.post_rooms.get_mut(&post_id).is_none() {
192 self.post_rooms.insert(post_id, HashSet::new());
198 .context(location_info!())?
204 pub fn join_user_room(
206 user_id: LocalUserId,
208 ) -> Result<(), LemmyError> {
209 // remove session from all rooms
210 for sessions in self.user_rooms.values_mut() {
211 sessions.remove(&id);
214 // If the room doesn't exist yet
215 if self.user_rooms.get_mut(&user_id).is_none() {
216 self.user_rooms.insert(user_id, HashSet::new());
222 .context(location_info!())?
228 fn send_post_room_message<OP, Response>(
233 websocket_id: Option<ConnectionId>,
234 ) -> Result<(), LemmyError>
236 OP: OperationType + ToString,
239 let res_str = &serialize_websocket_message(op, response)?;
240 if let Some(sessions) = self.post_rooms.get(&post_id) {
242 if let Some(my_id) = websocket_id {
247 self.sendit(res_str, *id);
253 pub fn send_community_room_message<OP, Response>(
257 community_id: CommunityId,
258 websocket_id: Option<ConnectionId>,
259 ) -> Result<(), LemmyError>
261 OP: OperationType + ToString,
264 let res_str = &serialize_websocket_message(op, response)?;
265 if let Some(sessions) = self.community_rooms.get(&community_id) {
267 if let Some(my_id) = websocket_id {
272 self.sendit(res_str, *id);
278 pub fn send_mod_room_message<OP, Response>(
282 community_id: CommunityId,
283 websocket_id: Option<ConnectionId>,
284 ) -> Result<(), LemmyError>
286 OP: OperationType + ToString,
289 let res_str = &serialize_websocket_message(op, response)?;
290 if let Some(sessions) = self.mod_rooms.get(&community_id) {
292 if let Some(my_id) = websocket_id {
297 self.sendit(res_str, *id);
303 pub fn send_all_message<OP, Response>(
307 websocket_id: Option<ConnectionId>,
308 ) -> Result<(), LemmyError>
310 OP: OperationType + ToString,
313 let res_str = &serialize_websocket_message(op, response)?;
314 for id in self.sessions.keys() {
315 if let Some(my_id) = websocket_id {
320 self.sendit(res_str, *id);
325 pub fn send_user_room_message<OP, Response>(
329 recipient_id: LocalUserId,
330 websocket_id: Option<ConnectionId>,
331 ) -> Result<(), LemmyError>
333 OP: OperationType + ToString,
336 let res_str = &serialize_websocket_message(op, response)?;
337 if let Some(sessions) = self.user_rooms.get(&recipient_id) {
339 if let Some(my_id) = websocket_id {
344 self.sendit(res_str, *id);
350 pub fn send_comment<OP>(
353 comment: &CommentResponse,
354 websocket_id: Option<ConnectionId>,
355 ) -> Result<(), LemmyError>
357 OP: OperationType + ToString,
359 let mut comment_reply_sent = comment.clone();
361 // Strip out my specific user info
362 comment_reply_sent.comment_view.my_vote = None;
364 // Send it to the post room
365 let mut comment_post_sent = comment_reply_sent.clone();
366 // Remove the recipients here to separate mentions / user messages from post or community comments
367 comment_post_sent.recipient_ids = Vec::new();
368 self.send_post_room_message(
371 comment_post_sent.comment_view.post.id,
375 // Send it to the community too
376 self.send_community_room_message(
382 self.send_community_room_message(
385 comment.comment_view.community.id,
389 // Send it to the recipient(s) including the mentioned users
390 for recipient_id in &comment_reply_sent.recipient_ids {
391 self.send_user_room_message(
402 pub fn send_post<OP>(
405 post_res: &PostResponse,
406 websocket_id: Option<ConnectionId>,
407 ) -> Result<(), LemmyError>
409 OP: OperationType + ToString,
411 let community_id = post_res.post_view.community.id;
413 // Don't send my data with it
414 let mut post_sent = post_res.clone();
415 post_sent.post_view.my_vote = None;
417 // Send it to /c/all and that community
418 self.send_community_room_message(user_operation, &post_sent, CommunityId(0), websocket_id)?;
419 self.send_community_room_message(user_operation, &post_sent, community_id, websocket_id)?;
421 // Send it to the post room
422 self.send_post_room_message(
425 post_res.post_view.post.id,
432 fn sendit(&self, message: &str, id: ConnectionId) {
433 if let Some(info) = self.sessions.get(&id) {
434 let _ = info.addr.do_send(WsMessage(message.to_owned()));
438 pub(super) fn parse_json_message(
440 msg: StandardMessage,
441 ctx: &mut Context<Self>,
442 ) -> impl Future<Output = Result<String, LemmyError>> {
443 let rate_limiter = self.rate_limiter.clone();
445 let ip: IpAddr = match self.sessions.get(&msg.id) {
446 Some(info) => info.ip.to_owned(),
447 None => IpAddr("blank_ip".to_string()),
450 let context = LemmyContext {
451 pool: self.pool.clone(),
452 chat_server: ctx.address(),
453 client: self.client.to_owned(),
454 activity_queue: self.activity_queue.to_owned(),
456 let message_handler_crud = self.message_handler_crud;
457 let message_handler = self.message_handler;
459 let json: Value = serde_json::from_str(&msg.msg)?;
460 let data = &json["data"].to_string();
461 let op = &json["op"].as_str().ok_or(ApiError {
462 message: "Unknown op type".to_string(),
465 if let Ok(user_operation_crud) = UserOperationCrud::from_str(op) {
466 let fut = (message_handler_crud)(context, msg.id, user_operation_crud.clone(), data);
467 match user_operation_crud {
468 UserOperationCrud::Register => rate_limiter.register().wrap(ip, fut).await,
469 UserOperationCrud::CreatePost => rate_limiter.post().wrap(ip, fut).await,
470 UserOperationCrud::CreateCommunity => rate_limiter.register().wrap(ip, fut).await,
471 _ => rate_limiter.message().wrap(ip, fut).await,
474 let user_operation = UserOperation::from_str(op)?;
475 let fut = (message_handler)(context, msg.id, user_operation.clone(), data);
476 rate_limiter.message().wrap(ip, fut).await