Skip to main content

rest/
sse.rs

1use actix_web::{
2    http::header::{CACHE_CONTROL, CONTENT_TYPE},
3    web::Bytes,
4    HttpResponse,
5};
6use futures::{Stream, StreamExt};
7use std::{convert::Infallible, time::Duration};
8
9/// One event in a `text/event-stream` response.
10#[derive(Debug, Clone, Default, PartialEq, Eq)]
11pub struct SseEvent {
12    event: Option<String>,
13    data: Option<String>,
14    id: Option<String>,
15    retry: Option<Duration>,
16    comment: Option<String>,
17}
18
19impl SseEvent {
20    pub fn data(data: impl Into<String>) -> Self {
21        Self {
22            data: Some(data.into()),
23            ..Self::default()
24        }
25    }
26
27    /// Creates a comment frame, which is useful as a connection heartbeat.
28    pub fn comment(comment: impl Into<String>) -> Self {
29        Self {
30            comment: Some(comment.into()),
31            ..Self::default()
32        }
33    }
34
35    pub fn with_event(mut self, event: impl Into<String>) -> Self {
36        self.event = Some(sanitize_single_line(event.into()));
37        self
38    }
39
40    pub fn with_id(mut self, id: impl Into<String>) -> Self {
41        self.id = Some(sanitize_single_line(id.into()).replace('\0', ""));
42        self
43    }
44
45    pub fn with_retry(mut self, retry: Duration) -> Self {
46        self.retry = Some(retry);
47        self
48    }
49
50    /// Encodes the event according to the EventSource wire format.
51    pub fn encode(&self) -> Bytes {
52        let mut encoded = String::new();
53        if let Some(comment) = &self.comment {
54            for line in lines(comment) {
55                encoded.push_str(": ");
56                encoded.push_str(line);
57                encoded.push('\n');
58            }
59        }
60        if let Some(event) = &self.event {
61            encoded.push_str("event: ");
62            encoded.push_str(event);
63            encoded.push('\n');
64        }
65        if let Some(id) = &self.id {
66            encoded.push_str("id: ");
67            encoded.push_str(id);
68            encoded.push('\n');
69        }
70        if let Some(retry) = self.retry {
71            encoded.push_str("retry: ");
72            encoded.push_str(&retry.as_millis().to_string());
73            encoded.push('\n');
74        }
75        if let Some(data) = &self.data {
76            for line in lines(data) {
77                encoded.push_str("data: ");
78                encoded.push_str(line);
79                encoded.push('\n');
80            }
81        }
82        encoded.push('\n');
83        Bytes::from(encoded)
84    }
85}
86
87/// Builds a streaming EventSource response without proxy or browser caching.
88pub fn sse_response<S>(events: S) -> HttpResponse
89where
90    S: Stream<Item = SseEvent> + 'static,
91{
92    HttpResponse::Ok()
93        .insert_header((CONTENT_TYPE, "text/event-stream"))
94        .insert_header((CACHE_CONTROL, "no-cache, no-transform"))
95        .insert_header(("x-accel-buffering", "no"))
96        .streaming(events.map(|event| Ok::<_, Infallible>(event.encode())))
97}
98
99fn sanitize_single_line(value: String) -> String {
100    value.replace(['\r', '\n'], "")
101}
102
103fn lines(value: &str) -> impl Iterator<Item = &str> {
104    value
105        .split('\n')
106        .map(|line| line.strip_suffix('\r').unwrap_or(line))
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112    use actix_web::{body::to_bytes, http::StatusCode};
113    use futures::stream;
114
115    #[test]
116    fn encodes_named_multiline_events_and_heartbeats() {
117        let event = SseEvent::data("first\nsecond")
118            .with_event("update\nignored")
119            .with_id("42\r\n")
120            .with_retry(Duration::from_millis(1500));
121        assert_eq!(
122            event.encode(),
123            Bytes::from_static(
124                b"event: updateignored\nid: 42\nretry: 1500\ndata: first\ndata: second\n\n"
125            )
126        );
127        assert_eq!(
128            SseEvent::comment("keepalive").encode(),
129            Bytes::from_static(b": keepalive\n\n")
130        );
131    }
132
133    #[actix_rt::test]
134    async fn response_sets_event_stream_headers_and_streams_frames() {
135        let response = sse_response(stream::iter([SseEvent::data("one"), SseEvent::data("two")]));
136
137        assert_eq!(response.status(), StatusCode::OK);
138        assert_eq!(
139            response.headers().get(CONTENT_TYPE).unwrap(),
140            "text/event-stream"
141        );
142        assert_eq!(
143            response.headers().get(CACHE_CONTROL).unwrap(),
144            "no-cache, no-transform"
145        );
146        assert_eq!(response.headers().get("x-accel-buffering").unwrap(), "no");
147        assert_eq!(
148            to_bytes(response.into_body()).await.unwrap(),
149            Bytes::from_static(b"data: one\n\ndata: two\n\n")
150        );
151    }
152}