use std::{ sync::Arc, time::{Duration, Instant}, }; use http_body_util::{combinators::BoxBody, StreamBody}; use hyper::{ body::{Bytes, Frame}, header, StatusCode, }; use jmap_proto::{error::request::RequestError, types::type_state::TypeState}; use utils::map::bitmap::Bitmap; use crate::{auth::AclToken, JMAP, LONG_SLUMBER}; use super::{http::ToHttpResponse, HttpRequest, HttpResponse, StateChangeResponse}; struct Ping { interval: Duration, last_ping: Instant, payload: Bytes, } impl JMAP { pub async fn handle_event_source( &self, req: HttpRequest, acl_token: Arc, ) -> HttpResponse { // Parse query let mut ping = 0; let mut types = Bitmap::default(); let mut close_after_state = false; for (key, value) in form_urlencoded::parse(req.uri().query().unwrap_or_default().as_bytes()) { match key.as_ref() { "types" => { for type_state in value.split(',') { if type_state == "*" { types = Bitmap::all(); break; } else if let Ok(type_state) = TypeState::try_from(type_state) { types.insert(type_state); } else { return RequestError::invalid_parameters().into_http_response(); } } } "closeafter" => match value.as_ref() { "state" => { close_after_state = true; } "no" => {} _ => return RequestError::invalid_parameters().into_http_response(), }, "ping" => match value.parse::() { Ok(value) => { ping = value; } Err(_) => return RequestError::invalid_parameters().into_http_response(), }, _ => {} } } let mut ping = if ping > 0 { #[cfg(not(feature = "test_mode"))] let interval = std::cmp::max(ping, 30) * 1000; #[cfg(feature = "test_mode")] let interval = ping * 1000; Ping { interval: Duration::from_millis(interval as u64), last_ping: Instant::now() - Duration::from_millis(interval as u64), payload: Bytes::from(format!( "event: ping\ndata: {{\"interval\": {}}}\n\n", interval )), } .into() } else { None }; let mut response = StateChangeResponse::new(); let throttle = self.config.event_source_throttle; // Register with state manager let mut change_rx = if let Some(change_rx) = self .subscribe_state_manager(acl_token.primary_id(), acl_token.primary_id(), types) .await { change_rx } else { return RequestError::internal_server_error().into_http_response(); }; hyper::Response::builder() .status(StatusCode::OK) .header(header::CONTENT_TYPE, "text/event-stream") .header(header::CACHE_CONTROL, "no-store") .body(BoxBody::new(StreamBody::new(async_stream::stream! { let mut last_message = Instant::now() - throttle; let mut timeout = ping.as_ref().map(|p| p.interval).unwrap_or(LONG_SLUMBER); loop { match tokio::time::timeout(timeout, change_rx.recv()).await { Ok(Some(state_change)) => { for (type_state, change_id) in state_change.types { response .changed .get_mut_or_insert(state_change.account_id.into()) .set(type_state, change_id.into()); } } Ok(None) => { tracing::debug!("Broadcast channel was closed."); break; } Err(_) => (), } timeout = if !response.changed.is_empty() { let elapsed = last_message.elapsed(); if elapsed >= throttle { last_message = Instant::now(); yield Ok(Frame::data(Bytes::from(format!( "event: state\ndata: {}\n\n", serde_json::to_string(&response).unwrap() )))); if close_after_state { break; } response.changed.clear(); ping.as_ref().map(|p| p.interval).unwrap_or(LONG_SLUMBER) } else { throttle - elapsed } } else if let Some(ping) = &mut ping { let elapsed = ping.last_ping.elapsed(); if elapsed >= ping.interval { ping.last_ping = Instant::now(); yield Ok(Frame::data(ping.payload.clone())); ping.interval } else { ping.interval - elapsed } } else { LONG_SLUMBER }; } }))) .unwrap() } }