SSE 与重定向

    fn into_response(self) -> Response {
        (
            [
                (http::header::CONTENT_TYPE, mime::TEXT_EVENT_STREAM.as_ref()),
                (http::header::CACHE_CONTROL, "no-cache"),
            ],
            Body::new(SseBody {
                event_stream: SyncWrapper::new(self.stream),
            }),
        )
            .into_response()
    }

    fn field(&mut self, name: &str, value: impl AsRef<[u8]>) {
        let value = value.as_ref();
        assert_eq!(
            memchr::memchr2(b'\r', b'\n', value),
            None,
            "SSE field value cannot contain newlines or carriage returns",
        );

        let buffer = self.buffer.as_mut();
        buffer.extend_from_slice(name.as_bytes());
        buffer.put_u8(b':');
        buffer.put_u8(b' ');
        buffer.extend_from_slice(value);
        buffer.put_u8(b'\n');
    }

        let mut last_split = 0;
        for delimiter in memchr::memchr2_iter(b'\n', b'\r', buf) {
            let _ = writer.write_all(&buf[last_split..=delimiter]);
            let _ = writer.write_all(b"data: ");
            last_split = delimiter + 1;
        }
        let _ = writer.write_all(&buf[last_split..]);

use axum::response::{IntoResponse, sse::{Event, KeepAlive, Sse}};
use std::panic::{UnwindSafe, catch_unwind};

fn try_build(name: &str, f: impl FnOnce() -> Event + UnwindSafe) {
  match catch_unwind(f) {
    Ok(ev) => {
      // finalize happens when the body is polled; build a one-event SSE body to read the bytes
      let res = Sse::new(futures::stream::iter([Ok::<_, std::convert::Infallible>(ev)])).into_response();
      let bytes = futures::executor::block_on(http_body_util::BodyExt::collect(res.into_body())).unwrap().to_bytes();
      println!("{name:<18} ok    {:?}", String::from_utf8_lossy(&bytes));
    }
    Err(p) => {
      let msg = p.downcast_ref::<String>().cloned().or_else(|| p.downcast_ref::<&str>().map(|s| s.to_string())).unwrap();
      println!("{name:<18} panic {msg:?}");
    }
  }
}

fn main() {
  std::panic::set_hook(Box::new(|_| {}));
  try_build("data with \\r\\n", || Event::default().data("a\r\nb"));
  try_build("event with \\n", || Event::default().event("a\nb"));
  try_build("comment with \\r", || Event::default().comment("a\rb"));
  try_build("id with \\0", || Event::default().id("a\0b"));
  try_build("event twice", || Event::default().event("a").event("b"));
  try_build("data twice", || Event::default().data("a").data("b"));
  let r = catch_unwind(|| KeepAlive::new().text("x\ny"));
  println!("{:<18} {}", "KeepAlive::text", if r.is_err() { "panic" } else { "ok" });
}

data with \r\n     ok    "data: a\rdata: \ndata: b\n\n"
event with \n      panic "assertion `left == right` failed: SSE field value cannot contain newlines or carriage returns\n  left: Some(1)\n right: None"
comment with \r    panic "assertion `left == right` failed: SSE field value cannot contain newlines or carriage returns\n  left: Some(1)\n right: None"
id with \0         panic "assertion `left == right` failed: Event ID cannot contain null characters\n  left: Some(1)\n right: None"
event twice        panic "Called `Event::event` multiple times"
data twice         panic "Called `Event::data*` multiple times"
KeepAlive::text    panic

use axum::{
  Router,
  response::sse::{Event, KeepAlive, Sse},
  routing::get,
};
use futures::stream::{self, Stream};
use std::{convert::Infallible, time::Duration};
use tokio::{
  io::{AsyncReadExt, AsyncWriteExt},
  net::{TcpListener, TcpStream},
  time::{Instant, sleep},
};

fn events() -> impl Stream<Item = Result<Event, Infallible>> {
  stream::unfold(0u32, |step| async move {
    let ev = match step {
      0 => Event::default().event("greet").id("1").data("line1\nline2"),
      1 => Event::default().json_data(serde_json::json!({"n": 2})).unwrap(),
      _ => {
        sleep(Duration::from_secs(3600)).await;
        return None;
      }
    };
    Some((Ok(ev), step + 1))
  })
}

#[tokio::main]
async fn main() {
  let app = Router::new().route(
    "/sse",
    get(|| async { Sse::new(events()).keep_alive(KeepAlive::new().interval(Duration::from_millis(400))) }),
  );
  let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
  let addr = listener.local_addr().unwrap();
  tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });

  let mut conn = TcpStream::connect(addr).await.unwrap();
  conn.write_all(b"GET /sse HTTP/1.1\r\nhost: x\r\n\r\n").await.unwrap();
  let start = Instant::now();
  let mut buf = vec![0u8; 4096];
  while start.elapsed() < Duration::from_millis(900) {
    let left = Duration::from_millis(900) - start.elapsed();
    match tokio::time::timeout(left, conn.read(&mut buf)).await {
      Ok(Ok(n)) if n > 0 => {
        let text = String::from_utf8_lossy(&buf[..n]);
        let ms = start.elapsed().as_millis() / 100 * 100;
        if text.starts_with("HTTP/1.1") {
          let head: Vec<&str> = text.split("\r\n").filter(|l| !l.starts_with("date:")).collect();
          println!("~{ms:>4} ms head+: {:?}", head.join("\r\n"));
        } else {
          println!("~{ms:>4} ms {text:?}");
        }
      }
      _ => break,
    }
  }
}

~   0 ms head+: "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncache-control: no-cache\r\ntransfer-encoding: chunked\r\n\r\n2C\r\nevent: greet\nid: 1\ndata: line1\ndata: line2\n\n\r\nF\r\ndata: {\"n\":2}\n\n\r\n"
~ 400 ms "3\r\n:\n\n\r\n"
~ 800 ms "3\r\n:\n\n\r\n"

        match this.inner.as_mut().poll_next(cx) {
            Poll::Ready(Some(Ok(event))) => {
                self.reset();

                Poll::Ready(Some(Ok(event)))
            }
            Poll::Ready(Some(Err(error))) => Poll::Ready(Some(Err(error))),
            Poll::Ready(None) => Poll::Ready(None),
            Poll::Pending => {
                ready!(this.alive_timer.poll(cx));

                let event = this.keep_alive.event.clone();

                self.reset();

                Poll::Ready(Some(Ok(event)))
            }
        }

    // This is intentionally not public since other kinds of redirects might not
    // use the `Location` header, namely `304 Not Modified`.
    //
    // We're open to adding more constructors upon request, if they make sense :)
    fn with_status_code(status_code: StatusCode, uri: &str) -> Self {
        assert!(
            status_code.is_redirection(),
            "not a redirection status code"
        );

        Self {
            status_code,
            location: uri.to_owned(),
        }
    }

    fn into_response(self) -> Response {
        match HeaderValue::try_from(self.location) {
            Ok(location) => (self.status_code, [(LOCATION, location)]).into_response(),
            Err(error) => (StatusCode::INTERNAL_SERVER_ERROR, error.to_string()).into_response(),
        }
    }

use axum::response::{IntoResponse, Redirect};

fn main() {
  for (name, r) in [
    ("to", Redirect::to("/done")),
    ("temporary", Redirect::temporary("/v2/upload")),
    ("permanent", Redirect::permanent("https://example.com/new")),
    ("bad location", Redirect::to("/a\nb")),
  ] {
    let res = r.into_response();
    println!("{name:<13} {} location={:?}", res.status(), res.headers().get("location"));
  }
}

to            303 See Other location=Some("/done")
temporary     307 Temporary Redirect location=Some("/v2/upload")
permanent     308 Permanent Redirect location=Some("https://example.com/new")
bad location  500 Internal Server Error location=None

~   0 ms head+: "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncache-control: no-cache\r\ntransfer-encoding: chunked\r\n\r\n2C\r\nevent: greet\nid: 1\ndata: line1\ndata: line2\n\n\r\nF\r\ndata: {\"n\":2}\n\n\r\n"
~ 400 ms "3\r\n:\n\n\r\n"
~ 800 ms "3\r\n:\n\n\r\n"
~1000 ms "1E\r\n: hi\nretry: 2500\ndata: third\n\n\r\n"
~1400 ms "3\r\n:\n\n\r\n"