event_stream_types/
lib.rs

1// Copyright (C) The Strand-Braid Authors
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4//! types for http event streams
5use bytes::Bytes;
6use futures::StreamExt;
7use http::{StatusCode, header::ACCEPT, request::Parts};
8use http_body::Frame;
9use std::{
10    collections::HashMap,
11    convert::Infallible,
12    pin::Pin,
13    sync::{Arc, RwLock},
14};
15use strand_bui_backend_session_types::ConnectionKey;
16use tokio::sync::mpsc::Sender;
17use tokio_stream::wrappers::ReceiverStream;
18
19pub type EventChunkSender = Sender<Frame<Bytes>>;
20type EventReceiver = ReceiverStream<Frame<Bytes>>;
21
22/// The type of possible connect event, either connect or disconnect.
23#[derive(Debug)]
24pub enum ConnectionEventType {
25    /// A connection event with sink for event stream messages to the connected client.
26    Connect(EventChunkSender),
27    /// A disconnection event.
28    Disconnect,
29}
30
31/// State associated with connection or disconnection.
32#[derive(Debug)]
33pub struct ConnectionEvent {
34    /// The type of connection for this event.
35    pub typ: ConnectionEventType,
36    /// Identifier for the connection (one per tab).
37    pub connection_key: ConnectionKey,
38}
39
40// header extractor for "Accept: text/event-stream" --------------------------
41
42pub struct AcceptsEventStream;
43
44impl<S> axum::extract::FromRequestParts<S> for AcceptsEventStream
45where
46    S: Send + Sync,
47{
48    type Rejection = (StatusCode, &'static str);
49    async fn from_request_parts(p: &mut Parts, _: &S) -> Result<Self, Self::Rejection> {
50        const ES: &[u8] = b"text/event-stream";
51        if p.headers.get_all(ACCEPT).iter().any(|v| v.as_bytes() == ES) {
52            Ok(AcceptsEventStream)
53        } else {
54            Err((
55                StatusCode::BAD_REQUEST,
56                "Bad request: It is required that you have an \
57                HTTP Header \"Accept: text/event-stream\"",
58            ))
59        }
60    }
61}
62
63// TolerantJson extractor --------------------------
64
65/// This is much like `axum::Json` but does not fail if the request does not set
66/// the 'Content-Type' header.
67///
68/// This is purely for backwards-compatibility and can be removed sometime.
69pub struct TolerantJson<T>(pub T);
70
71impl<T, S> axum::extract::FromRequest<S> for TolerantJson<T>
72where
73    T: serde::de::DeserializeOwned,
74    S: Send + Sync,
75{
76    type Rejection = axum::extract::rejection::JsonRejection;
77
78    async fn from_request(
79        mut req: axum::extract::Request,
80        state: &S,
81    ) -> Result<Self, Self::Rejection> {
82        if !json_content_type(req.headers()) {
83            tracing::error!("request should indicate \"Content-Type: application/json\"");
84            req.headers_mut().insert(
85                http::header::CONTENT_TYPE,
86                http::HeaderValue::from_static("application/json"),
87            );
88        }
89        match axum::Json::from_request(req, state).await {
90            Ok(payload) => Ok(TolerantJson(payload.0)),
91            Err(e) => Err(e),
92        }
93    }
94}
95
96// events body ---------------------------
97
98pub struct EventsBody {
99    events: EventReceiver,
100}
101
102impl EventsBody {
103    fn new(events: EventReceiver) -> Self {
104        Self { events }
105    }
106}
107
108impl http_body::Body for EventsBody {
109    type Data = Bytes;
110    type Error = Infallible;
111
112    fn poll_frame(
113        mut self: Pin<&mut Self>,
114        cx: &mut std::task::Context<'_>,
115    ) -> std::task::Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
116        match self.events.poll_next_unpin(cx) {
117            std::task::Poll::Ready(Some(frame)) => std::task::Poll::Ready(Some(Ok(frame))),
118            std::task::Poll::Ready(None) => std::task::Poll::Ready(None),
119            std::task::Poll::Pending => std::task::Poll::Pending,
120        }
121    }
122}
123
124impl axum::response::IntoResponse for EventsBody {
125    fn into_response(self) -> axum::response::Response {
126        let mut response = axum::response::Response::new(axum::body::Body::new(self));
127        response.headers_mut().insert(
128            "content-type",
129            http::header::HeaderValue::from_static("text/event-stream"),
130        );
131        response
132    }
133}
134
135// -----
136
137#[derive(Debug, Clone, Eq, PartialEq, Hash)]
138pub struct ConnectionSessionKey {
139    session_key: uuid::Uuid,
140    connection_key: std::net::SocketAddr,
141}
142
143impl ConnectionSessionKey {
144    pub fn new(session_key: uuid::Uuid, connection_key: std::net::SocketAddr) -> Self {
145        Self {
146            session_key,
147            connection_key,
148        }
149    }
150}
151
152/// broadcasts events to many listeners.
153///
154/// This is generic over the key type.
155#[derive(Debug, Clone)]
156pub struct EventBroadcaster<KEY> {
157    txers: Arc<RwLock<HashMap<KEY, EventChunkSender>>>,
158}
159
160impl<KEY> Default for EventBroadcaster<KEY> {
161    fn default() -> Self {
162        Self {
163            txers: Default::default(),
164        }
165    }
166}
167
168impl<KEY> EventBroadcaster<KEY>
169where
170    KEY: std::cmp::Eq + std::hash::Hash,
171{
172    /// Add a new connection indexed by a key.
173    ///
174    /// This returns an [EventsBody].
175    pub fn new_connection(&self, key: KEY) -> (EventChunkSender, EventsBody) {
176        let (tx, rx) = tokio::sync::mpsc::channel(10);
177        let mut txers = self.txers.write().unwrap();
178        txers.insert(key, tx.clone());
179        let rx = tokio_stream::wrappers::ReceiverStream::new(rx);
180        let body = EventsBody::new(rx);
181
182        (tx, body)
183    }
184    /// Transmit bytes as frame
185    ///
186    /// This will drop connections which have errored.
187    pub async fn broadcast_frame(&self, frame_string: String) {
188        let txers: Vec<_> = {
189            // Keep lock in this scope.
190            // Move all listeners out of shared map.
191            self.txers.write().unwrap().drain().collect()
192        };
193
194        // now we have released the lock and can await without holding the lock.
195        let mut keep_event_listeners = Vec::with_capacity(txers.len());
196        for (key, tx) in txers.into_iter() {
197            match tx.send(Frame::data(frame_string.clone().into())).await {
198                Ok(()) => {
199                    keep_event_listeners.push((key, tx));
200                }
201                Err(tokio::sync::mpsc::error::SendError(_frame)) => {
202                    // The receiver was dropped because the connection closed.
203                    tracing::debug!("send error");
204                }
205            }
206        }
207
208        {
209            // Keep lock in this scope.
210            // Move all listeners back into shared map.
211            let mut event_listeners = self.txers.write().unwrap();
212            for (key, value) in keep_event_listeners.into_iter() {
213                event_listeners.insert(key, value);
214            }
215        };
216    }
217}
218
219// ----
220
221// This does not really belong here...
222
223fn json_content_type(headers: &http::HeaderMap) -> bool {
224    let content_type = if let Some(content_type) = headers.get(http::header::CONTENT_TYPE) {
225        content_type
226    } else {
227        return false;
228    };
229
230    let content_type = if let Ok(content_type) = content_type.to_str() {
231        content_type
232    } else {
233        return false;
234    };
235
236    let mime = if let Ok(mime) = content_type.parse::<mime::Mime>() {
237        mime
238    } else {
239        return false;
240    };
241
242    mime.type_() == "application"
243        && (mime.subtype() == "json" || mime.suffix().is_some_and(|name| name == "json"))
244}
event_stream_types/lib.rs

event_stream_types/
lib.rs
