use super::body::{HyperResponseBody, StreamBody};
use super::handle::{ConnCtx, run_head_gate, to_hyper_full};
use super::record::record_request;
use super::response::HeaderPair;
use super::server_lifecycle::ConnectionLifecycle;
use super::sse::SseWriter;
use super::{Request, Response};
pub(super) async fn handle_sse(
handler: super::router::SseHandler,
req: Request,
buffer_size: usize,
lifecycle: &ConnectionLifecycle,
) -> Result<hyper::Response<HyperResponseBody>, std::convert::Infallible> {
if let Some(script) = lifecycle.script() {
script
.pause(super::mock::LifecycleCheckpoint::SseBufferConfigured(
buffer_size,
))
.await;
}
let (tx, rx) = tokio::sync::mpsc::channel::<bytes::Bytes>(buffer_size);
let _task = crate::task::spawn(move || {
let mut writer = SseWriter::new(tx);
if let Err(e) = handler(&req, &mut writer) {
tracing::warn!(error = %e, "SSE handler returned error");
}
});
let body = StreamBody { rx };
let builder = hyper::Response::builder()
.status(200)
.header("Content-Type", "text/event-stream")
.header("Cache-Control", "no-cache");
Ok(streaming_response_or_empty(builder, body))
}
fn empty_stream_body() -> StreamBody {
StreamBody {
rx: tokio::sync::mpsc::channel(1).1,
}
}
fn streaming_response_or_empty(
builder: hyper::http::response::Builder,
body: StreamBody,
) -> hyper::Response<HyperResponseBody> {
builder
.body(HyperResponseBody::Streaming(body))
.unwrap_or_else(|err| {
tracing::error!("failed to build streaming response: {err}");
hyper::Response::new(HyperResponseBody::Streaming(empty_stream_body()))
})
}
fn build_streaming_response(
status: u16,
headers: &[HeaderPair],
rx: tokio::sync::mpsc::Receiver<bytes::Bytes>,
is_head: bool,
) -> hyper::Response<HyperResponseBody> {
let mut builder = hyper::Response::builder().status(status);
for (name, value) in headers {
builder = builder.header(name.as_ref(), value.as_ref());
}
let body = match is_head {
true => empty_stream_body(),
false => StreamBody { rx },
};
streaming_response_or_empty(builder, body)
}
pub(super) fn handle_stream_response(
stream_resp: super::stream::StreamResponse,
req: Request,
ctx: &ConnCtx,
start: std::time::Instant,
) -> Result<hyper::Response<HyperResponseBody>, std::convert::Infallible> {
let is_head = req.is_head();
let parts = stream_resp.into_parts();
record_request(ctx, req.method(), req.path(), parts.status, start);
Ok(build_streaming_response(
parts.status,
&parts.headers,
parts.rx,
is_head,
))
}
pub(super) async fn handle_proxy_stream_response(
req: Request,
backend: &str,
prefix: &str,
ctx: &ConnCtx,
start: std::time::Instant,
) -> Result<hyper::Response<HyperResponseBody>, std::convert::Infallible> {
let proxy_req = super::async_proxy::ProxyRequest::from_request(&req);
let is_head = req.is_head();
let upstream =
match super::async_proxy::forward_request_streaming(proxy_req, backend, prefix).await {
Ok(u) => u,
Err(e) => {
tracing::warn!(error = %e, "streaming proxy upstream failed");
record_request(ctx, req.method(), req.path(), 502, start);
return Ok(to_hyper_full(Response::text_raw(
502,
"proxy upstream failed",
)));
}
};
record_request(ctx, req.method(), req.path(), upstream.status, start);
Ok(build_streaming_response(
upstream.status,
&upstream.headers,
upstream.rx,
is_head,
))
}
pub(super) async fn dispatch_streaming_proxy(
hyper_req: hyper::Request<hyper::body::Incoming>,
dispatch: &super::router::ServerDispatch,
ctx: &ConnCtx,
remote_addr: Option<std::net::IpAddr>,
backend: &str,
prefix: &str,
params: super::request::Params,
) -> Result<hyper::Response<HyperResponseBody>, std::convert::Infallible> {
let start = std::time::Instant::now();
let method =
super::method::Method::from_hyper(hyper_req.method()).unwrap_or(super::method::Method::Get);
let method_str = method.as_str();
let is_head = matches!(method, super::method::Method::Head);
let gate_blocked =
run_head_gate(&hyper_req, dispatch, remote_addr, ctx.is_tls, Some(params)).await;
if let Some(blocked) = gate_blocked {
record_request(
ctx,
method_str,
hyper_req.uri().path(),
blocked.status(),
start,
);
return Ok(to_hyper_full(blocked));
}
let scheme = match ctx.is_tls {
true => "https",
false => "http",
};
let (hyper_parts, body) = hyper_req.into_parts();
let path: Box<str> = hyper_parts.uri.path().into();
let proxy_parts = super::async_proxy::IncomingProxyParts {
method,
path_and_query: hyper_parts
.uri
.path_and_query()
.map_or("/", |pq| pq.as_str())
.into(),
headers: hyper_parts.headers,
remote_addr,
scheme,
};
let upstream =
match super::async_proxy::forward_incoming_streaming(proxy_parts, body, backend, prefix)
.await
{
Ok(u) => u,
Err(e) => {
tracing::warn!(error = %e, "streaming proxy upstream failed");
record_request(ctx, method_str, &path, 502, start);
return Ok(to_hyper_full(Response::text_raw(
502,
"proxy upstream failed",
)));
}
};
record_request(ctx, method_str, &path, upstream.status, start);
Ok(build_streaming_response(
upstream.status,
&upstream.headers,
upstream.rx,
is_head,
))
}