use crate::{ chat_server::{ChatServer, SessionInfo}, messages::{ CaptchaItem, CheckCaptcha, Connect, Disconnect, GetCommunityUsersOnline, GetPostUsersOnline, GetUsersOnline, JoinCommunityRoom, JoinModRoom, JoinPostRoom, JoinUserRoom, SendAllMessage, SendComment, SendCommunityRoomMessage, SendModRoomMessage, SendPost, SendUserRoomMessage, StandardMessage, }, OperationType, }; use actix::{Actor, Context, Handler, ResponseFuture}; use lemmy_db_schema::utils::naive_now; use lemmy_utils::ConnectionId; use opentelemetry::trace::TraceContextExt; use rand::Rng; use serde::Serialize; use tracing::{error, info}; use tracing_opentelemetry::OpenTelemetrySpanExt; /// Make actor from `ChatServer` impl Actor for ChatServer { /// We are going to use simple Context, we just need ability to communicate /// with other actors. type Context = Context; } /// Handler for Connect message. /// /// Register new session and assign unique id to this session impl Handler for ChatServer { type Result = ConnectionId; fn handle(&mut self, msg: Connect, _ctx: &mut Context) -> Self::Result { // register session with random id let id = self.rng.gen::(); info!("{} joined", &msg.ip); self.sessions.insert( id, SessionInfo { addr: msg.addr, ip: msg.ip, }, ); id } } /// Handler for Disconnect message. impl Handler for ChatServer { type Result = (); fn handle(&mut self, msg: Disconnect, _: &mut Context) { // Remove connections from sessions and all 3 scopes if self.sessions.remove(&msg.id).is_some() { for sessions in self.user_rooms.values_mut() { sessions.remove(&msg.id); } for sessions in self.post_rooms.values_mut() { sessions.remove(&msg.id); } for sessions in self.community_rooms.values_mut() { sessions.remove(&msg.id); } } } } fn root_span() -> tracing::Span { let span = tracing::info_span!( parent: None, "Websocket Request", trace_id = tracing::field::Empty, ); { let trace_id = span.context().span().span_context().trace_id().to_string(); span.record("trace_id", &tracing::field::display(trace_id)); } span } /// Handler for Message message. impl Handler for ChatServer { type Result = ResponseFuture>; fn handle(&mut self, msg: StandardMessage, ctx: &mut Context) -> Self::Result { use tracing::Instrument; let fut = self.parse_json_message(msg, ctx); let span = root_span(); Box::pin( async move { match fut.await { Ok(m) => { // info!("Message Sent: {}", m); Ok(m) } Err(e) => { error!("Error during message handling {}", e); Ok( e.to_json() .unwrap_or_else(|_| String::from(r#"{"error":"failed to serialize json"}"#)), ) } } } .instrument(span), ) } } impl Handler> for ChatServer where OP: OperationType + ToString, Response: Serialize, { type Result = (); fn handle(&mut self, msg: SendAllMessage, _: &mut Context) { self .send_all_message(&msg.op, &msg.response, msg.websocket_id) .ok(); } } impl Handler> for ChatServer where OP: OperationType + ToString, Response: Serialize, { type Result = (); fn handle(&mut self, msg: SendUserRoomMessage, _: &mut Context) { self .send_user_room_message( &msg.op, &msg.response, msg.local_recipient_id, msg.websocket_id, ) .ok(); } } impl Handler> for ChatServer where OP: OperationType + ToString, Response: Serialize, { type Result = (); fn handle(&mut self, msg: SendCommunityRoomMessage, _: &mut Context) { self .send_community_room_message(&msg.op, &msg.response, msg.community_id, msg.websocket_id) .ok(); } } impl Handler> for ChatServer where Response: Serialize, { type Result = (); fn handle(&mut self, msg: SendModRoomMessage, _: &mut Context) { self .send_mod_room_message(&msg.op, &msg.response, msg.community_id, msg.websocket_id) .ok(); } } impl Handler> for ChatServer where OP: OperationType + ToString, { type Result = (); fn handle(&mut self, msg: SendPost, _: &mut Context) { self.send_post(&msg.op, &msg.post, msg.websocket_id).ok(); } } impl Handler> for ChatServer where OP: OperationType + ToString, { type Result = (); fn handle(&mut self, msg: SendComment, _: &mut Context) { self .send_comment(&msg.op, &msg.comment, msg.websocket_id) .ok(); } } impl Handler for ChatServer { type Result = (); fn handle(&mut self, msg: JoinUserRoom, _: &mut Context) { self.join_user_room(msg.local_user_id, msg.id).ok(); } } impl Handler for ChatServer { type Result = (); fn handle(&mut self, msg: JoinCommunityRoom, _: &mut Context) { self.join_community_room(msg.community_id, msg.id).ok(); } } impl Handler for ChatServer { type Result = (); fn handle(&mut self, msg: JoinModRoom, _: &mut Context) { self.join_mod_room(msg.community_id, msg.id).ok(); } } impl Handler for ChatServer { type Result = (); fn handle(&mut self, msg: JoinPostRoom, _: &mut Context) { self.join_post_room(msg.post_id, msg.id).ok(); } } impl Handler for ChatServer { type Result = usize; fn handle(&mut self, _msg: GetUsersOnline, _: &mut Context) -> Self::Result { self.sessions.len() } } impl Handler for ChatServer { type Result = usize; fn handle(&mut self, msg: GetPostUsersOnline, _: &mut Context) -> Self::Result { if let Some(users) = self.post_rooms.get(&msg.post_id) { users.len() } else { 0 } } } impl Handler for ChatServer { type Result = usize; fn handle(&mut self, msg: GetCommunityUsersOnline, _: &mut Context) -> Self::Result { if let Some(users) = self.community_rooms.get(&msg.community_id) { users.len() } else { 0 } } } impl Handler for ChatServer { type Result = (); fn handle(&mut self, msg: CaptchaItem, _: &mut Context) { self.captchas.push(msg); } } impl Handler for ChatServer { type Result = bool; fn handle(&mut self, msg: CheckCaptcha, _: &mut Context) -> Self::Result { // Remove all the ones that are past the expire time self.captchas.retain(|x| x.expires.gt(&naive_now())); let check = self .captchas .iter() .any(|r| r.uuid == msg.uuid && r.answer.to_lowercase() == msg.answer.to_lowercase()); // Remove this uuid so it can't be re-checked (Checks only work once) self.captchas.retain(|x| x.uuid != msg.uuid); check } }