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::*};
10 use lemmy_db_schema::{CommunityId, LocalUserId, PostId};
13 rate_limit::RateLimit,
19 use rand::rngs::ThreadRng;
22 use serde_json::Value;
24 collections::{HashMap, HashSet},
27 use tokio::macros::support::Pin;
29 type MessageHandlerType = fn(
30 context: LemmyContext,
34 ) -> Pin<Box<dyn Future<Output = Result<String, LemmyError>> + '_>>;
36 /// `ChatServer` manages chat rooms and responsible for coordinating chat
38 pub struct ChatServer {
39 /// A map from generated random ID to session addr
40 pub sessions: HashMap<ConnectionId, SessionInfo>,
42 /// A map from post_id to set of connectionIDs
43 pub post_rooms: HashMap<PostId, HashSet<ConnectionId>>,
45 /// A map from community to set of connectionIDs
46 pub community_rooms: HashMap<CommunityId, HashSet<ConnectionId>>,
48 pub mod_rooms: HashMap<CommunityId, HashSet<ConnectionId>>,
50 /// A map from user id to its connection ID for joined users. Remember a user can have multiple
51 /// sessions (IE clients)
52 pub(super) user_rooms: HashMap<LocalUserId, HashSet<ConnectionId>>,
54 pub(super) rng: ThreadRng,
57 pub(super) pool: Pool<ConnectionManager<PgConnection>>,
59 /// Rate limiting based on rate type and IP addr
60 pub(super) rate_limiter: RateLimit,
62 /// A list of the current captchas
63 pub(super) captchas: Vec<CaptchaItem>,
65 message_handler: MessageHandlerType,
70 activity_queue: QueueHandle,
73 pub struct SessionInfo {
74 pub addr: Recipient<WsMessage>,
78 /// `ChatServer` is an actor. It maintains list of connection client session.
79 /// And manages available rooms. Peers send messages to other peers in same
80 /// room through `ChatServer`.
83 pool: Pool<ConnectionManager<PgConnection>>,
84 rate_limiter: RateLimit,
85 message_handler: MessageHandlerType,
87 activity_queue: QueueHandle,
90 sessions: HashMap::new(),
91 post_rooms: HashMap::new(),
92 community_rooms: HashMap::new(),
93 mod_rooms: HashMap::new(),
94 user_rooms: HashMap::new(),
95 rng: rand::thread_rng(),
105 pub fn join_community_room(
107 community_id: CommunityId,
109 ) -> Result<(), LemmyError> {
110 // remove session from all rooms
111 for sessions in self.community_rooms.values_mut() {
112 sessions.remove(&id);
115 // Also leave all post rooms
116 // This avoids double messages
117 for sessions in self.post_rooms.values_mut() {
118 sessions.remove(&id);
121 // If the room doesn't exist yet
122 if self.community_rooms.get_mut(&community_id).is_none() {
123 self.community_rooms.insert(community_id, HashSet::new());
128 .get_mut(&community_id)
129 .context(location_info!())?
134 pub fn join_mod_room(
136 community_id: CommunityId,
138 ) -> Result<(), LemmyError> {
139 // remove session from all rooms
140 for sessions in self.mod_rooms.values_mut() {
141 sessions.remove(&id);
144 // If the room doesn't exist yet
145 if self.mod_rooms.get_mut(&community_id).is_none() {
146 self.mod_rooms.insert(community_id, HashSet::new());
151 .get_mut(&community_id)
152 .context(location_info!())?
157 pub fn join_post_room(&mut self, post_id: PostId, id: ConnectionId) -> Result<(), LemmyError> {
158 // remove session from all rooms
159 for sessions in self.post_rooms.values_mut() {
160 sessions.remove(&id);
163 // Also leave all communities
164 // This avoids double messages
165 // TODO found a bug, whereby community messages like
166 // delete and remove aren't sent, because
167 // you left the community room
168 for sessions in self.community_rooms.values_mut() {
169 sessions.remove(&id);
172 // If the room doesn't exist yet
173 if self.post_rooms.get_mut(&post_id).is_none() {
174 self.post_rooms.insert(post_id, HashSet::new());
180 .context(location_info!())?
186 pub fn join_user_room(
188 user_id: LocalUserId,
190 ) -> Result<(), LemmyError> {
191 // remove session from all rooms
192 for sessions in self.user_rooms.values_mut() {
193 sessions.remove(&id);
196 // If the room doesn't exist yet
197 if self.user_rooms.get_mut(&user_id).is_none() {
198 self.user_rooms.insert(user_id, HashSet::new());
204 .context(location_info!())?
210 fn send_post_room_message<Response>(
215 websocket_id: Option<ConnectionId>,
216 ) -> Result<(), LemmyError>
220 let res_str = &serialize_websocket_message(op, response)?;
221 if let Some(sessions) = self.post_rooms.get(&post_id) {
223 if let Some(my_id) = websocket_id {
228 self.sendit(res_str, *id);
234 pub fn send_community_room_message<Response>(
238 community_id: CommunityId,
239 websocket_id: Option<ConnectionId>,
240 ) -> Result<(), LemmyError>
244 let res_str = &serialize_websocket_message(op, response)?;
245 if let Some(sessions) = self.community_rooms.get(&community_id) {
247 if let Some(my_id) = websocket_id {
252 self.sendit(res_str, *id);
258 pub fn send_mod_room_message<Response>(
262 community_id: CommunityId,
263 websocket_id: Option<ConnectionId>,
264 ) -> Result<(), LemmyError>
268 let res_str = &serialize_websocket_message(op, response)?;
269 if let Some(sessions) = self.mod_rooms.get(&community_id) {
271 if let Some(my_id) = websocket_id {
276 self.sendit(res_str, *id);
282 pub fn send_all_message<Response>(
286 websocket_id: Option<ConnectionId>,
287 ) -> Result<(), LemmyError>
291 let res_str = &serialize_websocket_message(op, response)?;
292 for id in self.sessions.keys() {
293 if let Some(my_id) = websocket_id {
298 self.sendit(res_str, *id);
303 pub fn send_user_room_message<Response>(
307 recipient_id: LocalUserId,
308 websocket_id: Option<ConnectionId>,
309 ) -> Result<(), LemmyError>
313 let res_str = &serialize_websocket_message(op, response)?;
314 if let Some(sessions) = self.user_rooms.get(&recipient_id) {
316 if let Some(my_id) = websocket_id {
321 self.sendit(res_str, *id);
329 user_operation: &UserOperation,
330 comment: &CommentResponse,
331 websocket_id: Option<ConnectionId>,
332 ) -> Result<(), LemmyError> {
333 let mut comment_reply_sent = comment.clone();
335 // Strip out my specific user info
336 comment_reply_sent.comment_view.my_vote = None;
338 // Send it to the post room
339 let mut comment_post_sent = comment_reply_sent.clone();
340 // Remove the recipients here to separate mentions / user messages from post or community comments
341 comment_post_sent.recipient_ids = Vec::new();
342 self.send_post_room_message(
345 comment_post_sent.comment_view.post.id,
349 // Send it to the community too
350 self.send_community_room_message(
356 self.send_community_room_message(
359 comment.comment_view.community.id,
363 // Send it to the recipient(s) including the mentioned users
364 for recipient_id in &comment_reply_sent.recipient_ids {
365 self.send_user_room_message(
378 user_operation: &UserOperation,
379 post_res: &PostResponse,
380 websocket_id: Option<ConnectionId>,
381 ) -> Result<(), LemmyError> {
382 let community_id = post_res.post_view.community.id;
384 // Don't send my data with it
385 let mut post_sent = post_res.clone();
386 post_sent.post_view.my_vote = None;
388 // Send it to /c/all and that community
389 self.send_community_room_message(user_operation, &post_sent, CommunityId(0), websocket_id)?;
390 self.send_community_room_message(user_operation, &post_sent, community_id, websocket_id)?;
392 // Send it to the post room
393 self.send_post_room_message(
396 post_res.post_view.post.id,
403 fn sendit(&self, message: &str, id: ConnectionId) {
404 if let Some(info) = self.sessions.get(&id) {
405 let _ = info.addr.do_send(WsMessage(message.to_owned()));
409 pub(super) fn parse_json_message(
411 msg: StandardMessage,
412 ctx: &mut Context<Self>,
413 ) -> impl Future<Output = Result<String, LemmyError>> {
414 let rate_limiter = self.rate_limiter.clone();
416 let ip: IpAddr = match self.sessions.get(&msg.id) {
417 Some(info) => info.ip.to_owned(),
418 None => IpAddr("blank_ip".to_string()),
421 let context = LemmyContext {
422 pool: self.pool.clone(),
423 chat_server: ctx.address(),
424 client: self.client.to_owned(),
425 activity_queue: self.activity_queue.to_owned(),
427 let message_handler = self.message_handler;
429 let json: Value = serde_json::from_str(&msg.msg)?;
430 let data = &json["data"].to_string();
431 let op = &json["op"].as_str().ok_or(ApiError {
432 message: "Unknown op type".to_string(),
435 let user_operation = UserOperation::from_str(&op)?;
436 let fut = (message_handler)(context, msg.id, user_operation.clone(), data);
437 match user_operation {
438 UserOperation::Register => rate_limiter.register().wrap(ip, fut).await,
439 UserOperation::CreatePost => rate_limiter.post().wrap(ip, fut).await,
440 UserOperation::CreateCommunity => rate_limiter.register().wrap(ip, fut).await,
441 _ => rate_limiter.message().wrap(ip, fut).await,