PushSubscription and EventSource tests passing.

crates/jmap/src/api/event_source.rs

@@ -0,0 +1,158 @@
+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<AclToken>,
+    ) -> 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::<u32>() {
+                    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()
+    }
+}