/* * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC * * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL */ use crate::api::{ToRequestError, notifications_into_push_objects, request::RequestHandler}; use common::{Server, auth::AccessToken, ipc::PushNotification}; use futures_util::{SinkExt, StreamExt}; use http_proto::HttpSessionData; use hyper::upgrade::Upgraded; use hyper_util::rt::TokioIo; use jmap_proto::{ error::request::RequestError, request::websocket::{ WebSocketMessage, WebSocketPushObject, WebSocketRequestError, WebSocketResponse, }, }; use std::future::Future; use std::time::Instant; use tokio_tungstenite::WebSocketStream; use trc::JmapEvent; use tungstenite::Message; use types::type_state::{DataType, StateChange}; use utils::map::bitmap::Bitmap; pub trait WebSocketHandler: Sync + Send { fn handle_websocket_stream( &self, stream: WebSocketStream>, access_token: AccessToken, session: HttpSessionData, ) -> impl Future + Send; } impl WebSocketHandler for Server { async fn handle_websocket_stream( &self, mut stream: WebSocketStream>, access_token: AccessToken, session: HttpSessionData, ) { trc::event!( Jmap(JmapEvent::WebsocketStart), SpanId = session.session_id, AccountId = access_token.account_id(), ); // Set timeouts let throttle = self.core.jmap.web_socket_throttle; let timeout = self.core.jmap.web_socket_timeout; let heartbeat = self.core.jmap.web_socket_heartbeat; let mut last_request = Instant::now(); let mut last_changes_sent = Instant::now() - throttle; let mut last_heartbeat = Instant::now() - heartbeat; let mut next_event = heartbeat; // Register with push manager let mut push_rx = match self .subscribe_push_manager(&access_token, Bitmap::all()) .await { Ok(push_rx) => push_rx, Err(err) => { trc::error!( err.details("Failed to subscribe to push manager") .span_id(session.session_id) ); let _ = stream .send(Message::Text( WebSocketRequestError::from(RequestError::internal_server_error()) .to_json() .into(), )) .await; return; } }; let mut notifications = Vec::new(); let mut change_types: Bitmap = Bitmap::new(); loop { tokio::select! { event = tokio::time::timeout(next_event, stream.next()) => { match event { Ok(Some(Ok(event))) => { match event { Message::Text(text) => { let response = match WebSocketMessage::parse( text.as_bytes(), self.core.jmap.request_max_calls, self.core.jmap.request_max_size, ) { Ok(WebSocketMessage::Request(request)) => { let response = self .handle_jmap_request( request.request, &access_token, &session, ) .await; WebSocketResponse::from_response(response, request.id) .to_json() } Ok(WebSocketMessage::PushEnable(push_enable)) => { change_types = if !push_enable.data_types.is_empty() { push_enable.data_types.into() } else { Bitmap::all() }; continue; } Ok(WebSocketMessage::PushDisable) => { change_types = Bitmap::new(); continue; } Err(err) => { let response = WebSocketRequestError::from(err.to_request_error()).to_json(); trc::error!(err.details("Failed to parse WebSocket message").span_id(session.session_id)); response }, }; if let Err(err) = stream.send(Message::Text(response.into())).await { trc::event!(Jmap(JmapEvent::WebsocketError), Details = "Failed to send text message", SpanId = session.session_id, Reason = err.to_string() ); } } Message::Ping(bytes) => { if let Err(err) = stream.send(Message::Pong(bytes)).await { trc::event!(Jmap(JmapEvent::WebsocketError), Details = "Failed to send pong message", SpanId = session.session_id, Reason = err.to_string() ); } } Message::Close(frame) => { let _ = stream.close(frame).await; break; } _ => (), } last_request = Instant::now(); last_heartbeat = Instant::now(); } Ok(Some(Err(err))) => { trc::event!(Jmap(JmapEvent::WebsocketError), Details = "Websocket error", SpanId = session.session_id, Reason = err.to_string() ); break; } Ok(None) => break, Err(_) => { // Verify timeout if last_request.elapsed() > timeout { trc::event!( Jmap(JmapEvent::WebsocketStop), SpanId = session.session_id, Reason = "Idle client" ); break; } } } } push_notification = push_rx.recv() => { if let Some(push_notification) = push_notification { match push_notification { PushNotification::StateChange(state_change) => { let mut types = state_change.types; types.intersection(&change_types); if !types.is_empty() { notifications.push(PushNotification::StateChange( StateChange { account_id: state_change.account_id, types, change_id: state_change.change_id, } )); } }, PushNotification::EmailPush(email_push) => { let state_change = email_push.to_state_change(); let mut types = state_change.types; types.intersection(&change_types); if !types.is_empty() { notifications.push(PushNotification::StateChange( StateChange { account_id: state_change.account_id, types, change_id: state_change.change_id, } )); } }, PushNotification::CalendarAlert(calendar_alert) => { if change_types.contains(DataType::CalendarAlert) { notifications.push(PushNotification::CalendarAlert( calendar_alert )); } }, } } else { trc::event!( Jmap(JmapEvent::WebsocketStop), SpanId = session.session_id, Reason = "State manager channel closed" ); break; } } } if !notifications.is_empty() { // Send any queued changes let elapsed = last_changes_sent.elapsed(); if elapsed >= throttle { for push in notifications_into_push_objects(std::mem::take(&mut notifications)) { let payload = WebSocketPushObject { push, push_state: None, }; if let Err(err) = stream.send(Message::Text(payload.to_json().into())).await { trc::event!( Jmap(JmapEvent::WebsocketError), Details = "Failed to send state change message.", SpanId = session.session_id, Reason = err.to_string() ); } } last_changes_sent = Instant::now(); last_heartbeat = Instant::now(); next_event = heartbeat; } else { next_event = throttle - elapsed; } } else if last_heartbeat.elapsed() > heartbeat { if let Err(err) = stream.send(Message::Ping(Vec::::new().into())).await { trc::event!( Jmap(JmapEvent::WebsocketError), Details = "Failed to send ping message.", SpanId = session.session_id, Reason = err.to_string() ); break; } last_heartbeat = Instant::now(); next_event = heartbeat; } } } }