jcoffey-dev
7dae9b29fd
Upstream commit: 474dd0229cb20cf513036619781ed97bd8073c3f Enterprise-only files removed or emptied: 63 Enterprise-only snippets removed: 117 in 50 files Dangling module declarations removed: 5 Cargo edits turning enterprise off: 14 Verification: clean Enterprise feature gates left for rebuilt features: 19 in 18 files Produced by tools/fork/strip.py. The full report is in docs/fork/strip-reports/ on main.
270 lines
12 KiB
Rust
270 lines
12 KiB
Rust
/*
|
|
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
|
|
*
|
|
* 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<TokioIo<Upgraded>>,
|
|
access_token: AccessToken,
|
|
session: HttpSessionData,
|
|
) -> impl Future<Output = ()> + Send;
|
|
}
|
|
|
|
impl WebSocketHandler for Server {
|
|
async fn handle_websocket_stream(
|
|
&self,
|
|
mut stream: WebSocketStream<TokioIo<Upgraded>>,
|
|
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<DataType> = 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::<u8>::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;
|
|
}
|
|
}
|
|
}
|
|
}
|