use crate::error::{Error, Result};
use crate::response::Response;
use crate::stream::{BodySender, Disconnected};
use serde::Serialize;
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct Event {
event: Option<String>,
id: Option<String>,
retry_ms: Option<u64>,
data: Option<String>,
comment: Option<String>,
}
impl Event {
pub fn data(data: impl Into<String>) -> Self {
Self {
data: Some(data.into()),
..Self::default()
}
}
pub fn json<T: Serialize>(value: &T) -> Result<Self> {
let data = serde_json::to_string(value)
.map_err(|e| Error::internal(format!("SSE JSON: {}", e)))?;
Ok(Self::data(data))
}
pub fn comment(text: impl Into<String>) -> Self {
Self {
comment: Some(text.into()),
..Self::default()
}
}
pub fn event(mut self, name: impl Into<String>) -> Self {
self.event = Some(name.into());
self
}
pub fn id(mut self, id: impl Into<String>) -> Self {
self.id = Some(id.into());
self
}
pub fn retry(mut self, millis: u64) -> Self {
self.retry_ms = Some(millis);
self
}
pub fn encode(&self) -> String {
fn single_line(s: &str) -> String {
s.chars().filter(|c| *c != '\n' && *c != '\r').collect()
}
let mut out = String::new();
if let Some(c) = &self.comment {
for line in c.lines() {
out.push_str(": ");
out.push_str(line);
out.push('\n');
}
}
if let Some(e) = &self.event {
out.push_str(&format!("event: {}\n", single_line(e)));
}
if let Some(id) = &self.id {
out.push_str(&format!("id: {}\n", single_line(id)));
}
if let Some(ms) = self.retry_ms {
out.push_str(&format!("retry: {}\n", ms));
}
if let Some(d) = &self.data {
let mut any = false;
for line in d.split('\n') {
any = true;
out.push_str("data: ");
out.push_str(line.trim_end_matches('\r'));
out.push('\n');
}
debug_assert!(any);
}
out.push('\n');
out
}
}
#[derive(Clone)]
pub struct SseSender {
inner: BodySender,
}
impl SseSender {
pub async fn send(&self, event: Event) -> std::result::Result<(), Disconnected> {
self.inner.send(event.encode()).await
}
pub fn is_closed(&self) -> bool {
self.inner.is_closed()
}
pub async fn closed(&self) {
self.inner.closed().await
}
}
pub fn sse() -> (SseSender, Response) {
let (tx, res) = Response::stream();
let res = res
.with_header("content-type", "text/event-stream; charset=utf-8")
.with_header("cache-control", "no-cache")
.with_header("x-accel-buffering", "no");
(SseSender { inner: tx }, res)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn encodes_fields_in_order_and_ends_with_a_blank_line() {
let e = Event::data("hello").event("greeting").id("7").retry(1500);
assert_eq!(
e.encode(),
"event: greeting\nid: 7\nretry: 1500\ndata: hello\n\n"
);
}
#[test]
fn multiline_data_becomes_multiple_data_lines() {
assert_eq!(
Event::data("a\nb\r\nc").encode(),
"data: a\ndata: b\ndata: c\n\n"
);
assert_eq!(Event::data("").encode(), "data: \n\n");
assert_eq!(
Event::data("trailing\n").encode(),
"data: trailing\ndata: \n\n"
);
}
#[test]
fn newlines_cannot_inject_fields() {
let e = Event::data("x").event("a\nretry: 1").id("1\r\ndata: evil");
let wire = e.encode();
assert_eq!(wire, "event: aretry: 1\nid: 1data: evil\ndata: x\n\n");
assert_eq!(wire.matches("\n\n").count(), 1);
}
#[test]
fn comments_and_json() {
assert_eq!(Event::comment("keep-alive").encode(), ": keep-alive\n\n");
assert_eq!(Event::comment("a\nb").encode(), ": a\n: b\n\n");
let e = Event::json(&serde_json::json!({ "up": true })).unwrap();
assert_eq!(e.encode(), "data: {\"up\":true}\n\n");
}
}