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},
34 use tokio::macros::support::Pin;
36 type MessageHandlerType = fn(
37 context: LemmyContext,
41 ) -> Pin<Box<dyn Future<Output = Result<String, LemmyError>> + '_>>;
43 type MessageHandlerCrudType = fn(
44 context: LemmyContext,
46 op: UserOperationCrud,
48 ) -> Pin<Box<dyn Future<Output = Result<String, LemmyError>> + '_>>;
50 /// `ChatServer` manages chat rooms and responsible for coordinating chat
52 pub struct ChatServer {
53 /// A map from generated random ID to session addr
54 pub sessions: HashMap<ConnectionId, SessionInfo>,
56 /// A map from post_id to set of connectionIDs
57 pub post_rooms: HashMap<PostId, HashSet<ConnectionId>>,
59 /// A map from community to set of connectionIDs
60 pub community_rooms: HashMap<CommunityId, HashSet<ConnectionId>>,
62 pub mod_rooms: HashMap<CommunityId, HashSet<ConnectionId>>,
64 /// A map from user id to its connection ID for joined users. Remember a user can have multiple
65 /// sessions (IE clients)
66 pub(super) user_rooms: HashMap<LocalUserId, HashSet<ConnectionId>>,
68 pub(super) rng: ThreadRng,
71 pub(super) pool: Pool<ConnectionManager<PgConnection>>,
73 /// Rate limiting based on rate type and IP addr
74 pub(super) rate_limiter: RateLimit,
76 /// A list of the current captchas
77 pub(super) captchas: Vec<CaptchaItem>,
79 message_handler: MessageHandlerType,
80 message_handler_crud: MessageHandlerCrudType,
85 activity_queue: QueueHandle,
88 pub struct SessionInfo {
89 pub addr: Recipient<WsMessage>,
93 /// `ChatServer` is an actor. It maintains list of connection client session.
94 /// And manages available rooms. Peers send messages to other peers in same
95 /// room through `ChatServer`.
98 pool: Pool<ConnectionManager<PgConnection>>,
99 rate_limiter: RateLimit,
100 message_handler: MessageHandlerType,
101 message_handler_crud: MessageHandlerCrudType,
103 activity_queue: QueueHandle,
106 sessions: HashMap::new(),
107 post_rooms: HashMap::new(),
108 community_rooms: HashMap::new(),
109 mod_rooms: HashMap::new(),
110 user_rooms: HashMap::new(),
111 rng: rand::thread_rng(),
114 captchas: Vec::new(),
116 message_handler_crud,
122 pub fn join_community_room(
124 community_id: CommunityId,
126 ) -> Result<(), LemmyError> {
127 // remove session from all rooms
128 for sessions in self.community_rooms.values_mut() {
129 sessions.remove(&id);
132 // Also leave all post rooms
133 // This avoids double messages
134 for sessions in self.post_rooms.values_mut() {
135 sessions.remove(&id);
138 // If the room doesn't exist yet
139 if self.community_rooms.get_mut(&community_id).is_none() {
140 self.community_rooms.insert(community_id, HashSet::new());
145 .get_mut(&community_id)
146 .context(location_info!())?
151 pub fn join_mod_room(
153 community_id: CommunityId,
155 ) -> Result<(), LemmyError> {
156 // remove session from all rooms
157 for sessions in self.mod_rooms.values_mut() {
158 sessions.remove(&id);
161 // If the room doesn't exist yet
162 if self.mod_rooms.get_mut(&community_id).is_none() {
163 self.mod_rooms.insert(community_id, HashSet::new());
168 .get_mut(&community_id)
169 .context(location_info!())?
174 pub fn join_post_room(&mut self, post_id: PostId, id: ConnectionId) -> Result<(), LemmyError> {
175 // remove session from all rooms
176 for sessions in self.post_rooms.values_mut() {
177 sessions.remove(&id);
180 // Also leave all communities
181 // This avoids double messages
182 // TODO found a bug, whereby community messages like
183 // delete and remove aren't sent, because
184 // you left the community room
185 for sessions in self.community_rooms.values_mut() {
186 sessions.remove(&id);
189 // If the room doesn't exist yet
190 if self.post_rooms.get_mut(&post_id).is_none() {
191 self.post_rooms.insert(post_id, HashSet::new());
197 .context(location_info!())?
203 pub fn join_user_room(
205 user_id: LocalUserId,
207 ) -> Result<(), LemmyError> {
208 // remove session from all rooms
209 for sessions in self.user_rooms.values_mut() {
210 sessions.remove(&id);
213 // If the room doesn't exist yet
214 if self.user_rooms.get_mut(&user_id).is_none() {
215 self.user_rooms.insert(user_id, HashSet::new());
221 .context(location_info!())?
227 fn send_post_room_message<OP, Response>(
232 websocket_id: Option<ConnectionId>,
233 ) -> Result<(), LemmyError>
235 OP: OperationType + ToString,
238 let res_str = &serialize_websocket_message(op, response)?;
239 if let Some(sessions) = self.post_rooms.get(&post_id) {
241 if let Some(my_id) = websocket_id {
246 self.sendit(res_str, *id);
252 pub fn send_community_room_message<OP, Response>(
256 community_id: CommunityId,
257 websocket_id: Option<ConnectionId>,
258 ) -> Result<(), LemmyError>
260 OP: OperationType + ToString,
263 let res_str = &serialize_websocket_message(op, response)?;
264 if let Some(sessions) = self.community_rooms.get(&community_id) {
266 if let Some(my_id) = websocket_id {
271 self.sendit(res_str, *id);
277 pub fn send_mod_room_message<OP, Response>(
281 community_id: CommunityId,
282 websocket_id: Option<ConnectionId>,
283 ) -> Result<(), LemmyError>
285 OP: OperationType + ToString,
288 let res_str = &serialize_websocket_message(op, response)?;
289 if let Some(sessions) = self.mod_rooms.get(&community_id) {
291 if let Some(my_id) = websocket_id {
296 self.sendit(res_str, *id);
302 pub fn send_all_message<OP, Response>(
306 websocket_id: Option<ConnectionId>,
307 ) -> Result<(), LemmyError>
309 OP: OperationType + ToString,
312 let res_str = &serialize_websocket_message(op, response)?;
313 for id in self.sessions.keys() {
314 if let Some(my_id) = websocket_id {
319 self.sendit(res_str, *id);
324 pub fn send_user_room_message<OP, Response>(
328 recipient_id: LocalUserId,
329 websocket_id: Option<ConnectionId>,
330 ) -> Result<(), LemmyError>
332 OP: OperationType + ToString,
335 let res_str = &serialize_websocket_message(op, response)?;
336 if let Some(sessions) = self.user_rooms.get(&recipient_id) {
338 if let Some(my_id) = websocket_id {
343 self.sendit(res_str, *id);
349 pub fn send_comment<OP>(
352 comment: &CommentResponse,
353 websocket_id: Option<ConnectionId>,
354 ) -> Result<(), LemmyError>
356 OP: OperationType + ToString,
358 let mut comment_reply_sent = comment.clone();
360 // Strip out my specific user info
361 comment_reply_sent.comment_view.my_vote = None;
363 // Send it to the post room
364 let mut comment_post_sent = comment_reply_sent.clone();
365 // Remove the recipients here to separate mentions / user messages from post or community comments
366 comment_post_sent.recipient_ids = Vec::new();
367 self.send_post_room_message(
370 comment_post_sent.comment_view.post.id,
374 // Send it to the community too
375 self.send_community_room_message(
381 self.send_community_room_message(
384 comment.comment_view.community.id,
388 // Send it to the recipient(s) including the mentioned users
389 for recipient_id in &comment_reply_sent.recipient_ids {
390 self.send_user_room_message(
401 pub fn send_post<OP>(
404 post_res: &PostResponse,
405 websocket_id: Option<ConnectionId>,
406 ) -> Result<(), LemmyError>
408 OP: OperationType + ToString,
410 let community_id = post_res.post_view.community.id;
412 // Don't send my data with it
413 let mut post_sent = post_res.clone();
414 post_sent.post_view.my_vote = None;
416 // Send it to /c/all and that community
417 self.send_community_room_message(user_operation, &post_sent, CommunityId(0), websocket_id)?;
418 self.send_community_room_message(user_operation, &post_sent, community_id, websocket_id)?;
420 // Send it to the post room
421 self.send_post_room_message(
424 post_res.post_view.post.id,
431 fn sendit(&self, message: &str, id: ConnectionId) {
432 if let Some(info) = self.sessions.get(&id) {
433 let _ = info.addr.do_send(WsMessage(message.to_owned()));
437 pub(super) fn parse_json_message(
439 msg: StandardMessage,
440 ctx: &mut Context<Self>,
441 ) -> impl Future<Output = Result<String, LemmyError>> {
442 let rate_limiter = self.rate_limiter.clone();
444 let ip: IpAddr = match self.sessions.get(&msg.id) {
445 Some(info) => info.ip.to_owned(),
446 None => IpAddr("blank_ip".to_string()),
449 let context = LemmyContext {
450 pool: self.pool.clone(),
451 chat_server: ctx.address(),
452 client: self.client.to_owned(),
453 activity_queue: self.activity_queue.to_owned(),
455 let message_handler_crud = self.message_handler_crud;
456 let message_handler = self.message_handler;
458 let json: Value = serde_json::from_str(&msg.msg)?;
459 let data = &json["data"].to_string();
460 let op = &json["op"].as_str().ok_or(ApiError {
461 message: "Unknown op type".to_string(),
464 if let Ok(user_operation_crud) = UserOperationCrud::from_str(&op) {
465 let fut = (message_handler_crud)(context, msg.id, user_operation_crud.clone(), data);
466 match user_operation_crud {
467 UserOperationCrud::Register => rate_limiter.register().wrap(ip, fut).await,
468 UserOperationCrud::CreatePost => rate_limiter.post().wrap(ip, fut).await,
469 UserOperationCrud::CreateCommunity => rate_limiter.register().wrap(ip, fut).await,
470 _ => rate_limiter.message().wrap(ip, fut).await,
473 let user_operation = UserOperation::from_str(&op)?;
474 let fut = (message_handler)(context, msg.id, user_operation.clone(), data);
475 rate_limiter.message().wrap(ip, fut).await