Skip to content

Commit

Permalink
fix(sse): skip sse incompatible chars of serde_json::RawValue
Browse files Browse the repository at this point in the history
  • Loading branch information
Threated committed Oct 16, 2024
1 parent 0ddc63f commit 1510729
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 2 deletions.
2 changes: 1 addition & 1 deletion axum/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -117,7 +117,7 @@ quickcheck = "1.0"
quickcheck_macros = "1.0"
reqwest = { version = "0.12", default-features = false, features = ["json", "stream", "multipart"] }
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
serde_json = { version = "1.0", features = ["raw_value"]}
time = { version = "0.3", features = ["serde-human-readable"] }
tokio = { package = "tokio", version = "1.25.0", features = ["macros", "rt", "rt-multi-thread", "net", "test-util"] }
tokio-stream = "0.1"
Expand Down
32 changes: 31 additions & 1 deletion axum/src/response/sse.rs
Original file line number Diff line number Diff line change
Expand Up @@ -213,7 +213,24 @@ impl Event {
}

self.buffer.extend_from_slice(b"data: ");
serde_json::to_writer((&mut self.buffer).writer(), &data).map_err(axum_core::Error::new)?;
struct IgnoreNewLines<'a>(bytes::buf::Writer<&'a mut BytesMut>);
impl std::io::Write for IgnoreNewLines<'_> {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
let mut last_split = 0;
for delimiter in memchr::memchr2_iter(b'\n', b'\r', buf) {
self.0.write(&buf[last_split..delimiter])?;
last_split = delimiter + 1;
}
self.0.write(&buf[last_split..])?;
Ok(buf.len())
}

fn flush(&mut self) -> std::io::Result<()> {
self.0.flush()
}
}
serde_json::to_writer(IgnoreNewLines((&mut self.buffer).writer()), &data)
.map_err(axum_core::Error::new)?;
self.buffer.put_u8(b'\n');

self.flags.insert(EventFlags::HAS_DATA);
Expand Down Expand Up @@ -515,6 +532,7 @@ mod tests {
use super::*;
use crate::{routing::get, test_helpers::*, Router};
use futures_util::stream;
use serde_json::value::RawValue;
use std::{collections::HashMap, convert::Infallible};
use tokio_stream::StreamExt as _;

Expand All @@ -527,6 +545,18 @@ mod tests {
assert_eq!(&*leading_space.finalize(), b"data: foobar\n\n");
}

#[test]
fn valid_json_raw_value_chars_stripped() {
let json_string = "{\r\"foo\": \n\r\r \"bar\\n\"\n}";
let json_raw_value_event = Event::default()
.json_data(serde_json::from_str::<&RawValue>(&json_string).unwrap())
.unwrap();
assert_eq!(
&*json_raw_value_event.finalize(),
format!("data: {}\n\n", json_string.replace(['\n', '\r'], "")).as_bytes()
);
}

#[crate::test]
async fn basic() {
let app = Router::new().route(
Expand Down

0 comments on commit 1510729

Please sign in to comment.