1 use crate::{messages::*, serialize_websocket_message, LemmyContext, UserOperation};
3 use anyhow::Context as acontext;
4 use background_jobs::QueueHandle;
6 r2d2::{ConnectionManager, Pool},
9 use lemmy_api_structs::{comment::*, post::*};
12 rate_limit::RateLimit,
21 use rand::rngs::ThreadRng;
24 use serde_json::Value;
26 collections::{HashMap, HashSet},
29 use tokio::macros::support::Pin;
31 type MessageHandlerType = fn(
32 context: LemmyContext,
36 ) -> Pin<Box<dyn Future<Output = Result<String, LemmyError>> + '_>>;
38 /// `ChatServer` manages chat rooms and responsible for coordinating chat
40 pub struct ChatServer {
41 /// A map from generated random ID to session addr
42 pub sessions: HashMap<ConnectionId, SessionInfo>,
44 /// A map from post_id to set of connectionIDs
45 pub post_rooms: HashMap<PostId, HashSet<ConnectionId>>,
47 /// A map from community to set of connectionIDs
48 pub community_rooms: HashMap<CommunityId, HashSet<ConnectionId>>,
50 pub mod_rooms: HashMap<CommunityId, HashSet<ConnectionId>>,
52 /// A map from user id to its connection ID for joined users. Remember a user can have multiple
53 /// sessions (IE clients)
54 pub(super) user_rooms: HashMap<UserId, HashSet<ConnectionId>>,
56 pub(super) rng: ThreadRng,
59 pub(super) pool: Pool<ConnectionManager<PgConnection>>,
61 /// Rate limiting based on rate type and IP addr
62 pub(super) rate_limiter: RateLimit,
64 /// A list of the current captchas
65 pub(super) captchas: Vec<CaptchaItem>,
67 message_handler: MessageHandlerType,
72 activity_queue: QueueHandle,
75 pub struct SessionInfo {
76 pub addr: Recipient<WsMessage>,
80 /// `ChatServer` is an actor. It maintains list of connection client session.
81 /// And manages available rooms. Peers send messages to other peers in same
82 /// room through `ChatServer`.
85 pool: Pool<ConnectionManager<PgConnection>>,
86 rate_limiter: RateLimit,
87 message_handler: MessageHandlerType,
89 activity_queue: QueueHandle,
92 sessions: HashMap::new(),
93 post_rooms: HashMap::new(),
94 community_rooms: HashMap::new(),
95 mod_rooms: HashMap::new(),
96 user_rooms: HashMap::new(),
97 rng: rand::thread_rng(),
100 captchas: Vec::new(),
107 pub fn join_community_room(
109 community_id: CommunityId,
111 ) -> Result<(), LemmyError> {
112 // remove session from all rooms
113 for sessions in self.community_rooms.values_mut() {
114 sessions.remove(&id);
117 // Also leave all post rooms
118 // This avoids double messages
119 for sessions in self.post_rooms.values_mut() {
120 sessions.remove(&id);
123 // If the room doesn't exist yet
124 if self.community_rooms.get_mut(&community_id).is_none() {
125 self.community_rooms.insert(community_id, HashSet::new());
130 .get_mut(&community_id)
131 .context(location_info!())?
136 pub fn join_mod_room(
138 community_id: CommunityId,
140 ) -> Result<(), LemmyError> {
141 // remove session from all rooms
142 for sessions in self.mod_rooms.values_mut() {
143 sessions.remove(&id);
146 // If the room doesn't exist yet
147 if self.mod_rooms.get_mut(&community_id).is_none() {
148 self.mod_rooms.insert(community_id, HashSet::new());
153 .get_mut(&community_id)
154 .context(location_info!())?
159 pub fn join_post_room(&mut self, post_id: PostId, id: ConnectionId) -> Result<(), LemmyError> {
160 // remove session from all rooms
161 for sessions in self.post_rooms.values_mut() {
162 sessions.remove(&id);
165 // Also leave all communities
166 // This avoids double messages
167 // TODO found a bug, whereby community messages like
168 // delete and remove aren't sent, because
169 // you left the community room
170 for sessions in self.community_rooms.values_mut() {
171 sessions.remove(&id);
174 // If the room doesn't exist yet
175 if self.post_rooms.get_mut(&post_id).is_none() {
176 self.post_rooms.insert(post_id, HashSet::new());
182 .context(location_info!())?
188 pub fn join_user_room(&mut self, user_id: UserId, id: ConnectionId) -> Result<(), LemmyError> {
189 // remove session from all rooms
190 for sessions in self.user_rooms.values_mut() {
191 sessions.remove(&id);
194 // If the room doesn't exist yet
195 if self.user_rooms.get_mut(&user_id).is_none() {
196 self.user_rooms.insert(user_id, HashSet::new());
202 .context(location_info!())?
208 fn send_post_room_message<Response>(
213 websocket_id: Option<ConnectionId>,
214 ) -> Result<(), LemmyError>
218 let res_str = &serialize_websocket_message(op, response)?;
219 if let Some(sessions) = self.post_rooms.get(&post_id) {
221 if let Some(my_id) = websocket_id {
226 self.sendit(res_str, *id);
232 pub fn send_community_room_message<Response>(
236 community_id: CommunityId,
237 websocket_id: Option<ConnectionId>,
238 ) -> Result<(), LemmyError>
242 let res_str = &serialize_websocket_message(op, response)?;
243 if let Some(sessions) = self.community_rooms.get(&community_id) {
245 if let Some(my_id) = websocket_id {
250 self.sendit(res_str, *id);
256 pub fn send_mod_room_message<Response>(
260 community_id: CommunityId,
261 websocket_id: Option<ConnectionId>,
262 ) -> Result<(), LemmyError>
266 let res_str = &serialize_websocket_message(op, response)?;
267 if let Some(sessions) = self.mod_rooms.get(&community_id) {
269 if let Some(my_id) = websocket_id {
274 self.sendit(res_str, *id);
280 pub fn send_all_message<Response>(
284 websocket_id: Option<ConnectionId>,
285 ) -> Result<(), LemmyError>
289 let res_str = &serialize_websocket_message(op, response)?;
290 for id in self.sessions.keys() {
291 if let Some(my_id) = websocket_id {
296 self.sendit(res_str, *id);
301 pub fn send_user_room_message<Response>(
305 recipient_id: UserId,
306 websocket_id: Option<ConnectionId>,
307 ) -> Result<(), LemmyError>
311 let res_str = &serialize_websocket_message(op, response)?;
312 if let Some(sessions) = self.user_rooms.get(&recipient_id) {
314 if let Some(my_id) = websocket_id {
319 self.sendit(res_str, *id);
327 user_operation: &UserOperation,
328 comment: &CommentResponse,
329 websocket_id: Option<ConnectionId>,
330 ) -> Result<(), LemmyError> {
331 let mut comment_reply_sent = comment.clone();
333 // Strip out my specific user info
334 comment_reply_sent.comment_view.my_vote = None;
336 // Send it to the post room
337 let mut comment_post_sent = comment_reply_sent.clone();
338 // Remove the recipients here to separate mentions / user messages from post or community comments
339 comment_post_sent.recipient_ids = Vec::new();
340 self.send_post_room_message(
343 comment_post_sent.comment_view.post.id,
347 // Send it to the community too
348 self.send_community_room_message(user_operation, &comment_post_sent, 0, websocket_id)?;
349 self.send_community_room_message(
352 comment.comment_view.community.id,
356 // Send it to the recipient(s) including the mentioned users
357 for recipient_id in &comment_reply_sent.recipient_ids {
358 self.send_user_room_message(
371 user_operation: &UserOperation,
372 post_res: &PostResponse,
373 websocket_id: Option<ConnectionId>,
374 ) -> Result<(), LemmyError> {
375 let community_id = post_res.post_view.community.id;
377 // Don't send my data with it
378 let mut post_sent = post_res.clone();
379 post_sent.post_view.my_vote = None;
381 // Send it to /c/all and that community
382 self.send_community_room_message(user_operation, &post_sent, 0, websocket_id)?;
383 self.send_community_room_message(user_operation, &post_sent, community_id, websocket_id)?;
385 // Send it to the post room
386 self.send_post_room_message(
389 post_res.post_view.post.id,
396 fn sendit(&self, message: &str, id: ConnectionId) {
397 if let Some(info) = self.sessions.get(&id) {
398 let _ = info.addr.do_send(WsMessage(message.to_owned()));
402 pub(super) fn parse_json_message(
404 msg: StandardMessage,
405 ctx: &mut Context<Self>,
406 ) -> impl Future<Output = Result<String, LemmyError>> {
407 let rate_limiter = self.rate_limiter.clone();
409 let ip: IpAddr = match self.sessions.get(&msg.id) {
410 Some(info) => info.ip.to_owned(),
411 None => "blank_ip".to_string(),
414 let context = LemmyContext {
415 pool: self.pool.clone(),
416 chat_server: ctx.address(),
417 client: self.client.to_owned(),
418 activity_queue: self.activity_queue.to_owned(),
420 let message_handler = self.message_handler;
422 let json: Value = serde_json::from_str(&msg.msg)?;
423 let data = &json["data"].to_string();
424 let op = &json["op"].as_str().ok_or(ApiError {
425 message: "Unknown op type".to_string(),
428 let user_operation = UserOperation::from_str(&op)?;
429 let fut = (message_handler)(context, msg.id, user_operation.clone(), data);
430 match user_operation {
431 UserOperation::Register => rate_limiter.register().wrap(ip, fut).await,
432 UserOperation::CreatePost => rate_limiter.post().wrap(ip, fut).await,
433 UserOperation::CreateCommunity => rate_limiter.register().wrap(ip, fut).await,
434 _ => rate_limiter.message().wrap(ip, fut).await,