343 lines
14 KiB
Rust
343 lines
14 KiB
Rust
use crate::{
|
|
HttpRequest, HttpResponse, Router, WebMTPConnection, WebServerError, WebServerMetrics,
|
|
transport::accept_web_connection,
|
|
};
|
|
use bytes::{Buf, Bytes};
|
|
use http::{Request, Response, StatusCode};
|
|
use mtp_host::HostConfig;
|
|
use std::{net::SocketAddr, sync::Arc, time::Duration};
|
|
use tokio::sync::{Semaphore, watch};
|
|
|
|
pub(crate) struct DriverConfig {
|
|
pub(crate) router: Router,
|
|
pub(crate) mtp_path: String,
|
|
pub(crate) max_request_body: usize,
|
|
pub(crate) request_timeout: Duration,
|
|
pub(crate) drain_timeout: Duration,
|
|
pub(crate) send_pongs: bool,
|
|
pub(crate) policy: mtp_transport::Policy,
|
|
pub(crate) host_config: Arc<HostConfig>,
|
|
pub(crate) metrics: Option<Arc<dyn WebServerMetrics>>,
|
|
pub(crate) auth_semaphore: Arc<Semaphore>,
|
|
}
|
|
|
|
pub(crate) async fn run_driver(
|
|
endpoint: quinn::Endpoint,
|
|
config: DriverConfig,
|
|
mtp_tx: tokio::sync::mpsc::Sender<Result<WebMTPConnection, mtp_host::AcceptError>>,
|
|
connection_semaphore: Arc<Semaphore>,
|
|
mut shutdown_rx: watch::Receiver<()>,
|
|
) {
|
|
let DriverConfig {
|
|
router,
|
|
mtp_path,
|
|
max_request_body,
|
|
request_timeout,
|
|
drain_timeout,
|
|
send_pongs,
|
|
policy,
|
|
host_config,
|
|
metrics,
|
|
auth_semaphore,
|
|
} = config;
|
|
let mut connection_tasks = tokio::task::JoinSet::new();
|
|
loop {
|
|
tokio::select! {
|
|
biased;
|
|
_ = shutdown_rx.changed() => {
|
|
break;
|
|
}
|
|
incoming = endpoint.accept() => {
|
|
let Some(incoming) = incoming else {
|
|
break;
|
|
};
|
|
// Do not await capacity here: doing so would prevent this loop
|
|
// from observing shutdown while all connection slots are in use.
|
|
let permit = match connection_semaphore.clone().try_acquire_owned() {
|
|
Ok(permit) => permit,
|
|
Err(_) => {
|
|
tracing::debug!("rejecting QUIC connection at configured connection limit");
|
|
continue;
|
|
}
|
|
};
|
|
let router = router.clone();
|
|
let mtp_path = mtp_path.clone();
|
|
let mtp_tx = mtp_tx.clone();
|
|
let metrics = metrics.clone();
|
|
let host_config = host_config.clone();
|
|
let auth_semaphore = auth_semaphore.clone();
|
|
connection_tasks.spawn(async move {
|
|
let _permit = permit;
|
|
let connect_start = std::time::Instant::now();
|
|
let connection = match incoming.await {
|
|
Ok(connection) => connection,
|
|
Err(error) => {
|
|
tracing::debug!(%error, "QUIC connection failed during handshake");
|
|
return;
|
|
}
|
|
};
|
|
let mut builder = h3::server::builder();
|
|
builder.enable_extended_connect(true);
|
|
builder.enable_webtransport(true);
|
|
builder.enable_datagram(true);
|
|
builder.max_webtransport_sessions(16);
|
|
let mut h3 = match builder
|
|
.build(h3_quinn::Connection::new(connection.clone()))
|
|
.await
|
|
{
|
|
Ok(connection) => connection,
|
|
Err(error) => {
|
|
tracing::debug!(%error, "HTTP/3 connection setup failed");
|
|
return;
|
|
}
|
|
};
|
|
|
|
if let Some(ref m) = metrics {
|
|
m.connection_accepted();
|
|
}
|
|
let remote_addr = connection.remote_address();
|
|
|
|
let mut tasks = tokio::task::JoinSet::new();
|
|
loop {
|
|
let resolver = match h3.accept().await {
|
|
Ok(Some(resolver)) => resolver,
|
|
Ok(None) => break,
|
|
Err(error) => {
|
|
tracing::debug!(%error, "HTTP/3 request accept failed");
|
|
break;
|
|
}
|
|
};
|
|
let (request, mut stream) = match resolver.resolve_request().await {
|
|
Ok(request) => request,
|
|
Err(error) => {
|
|
tracing::debug!(%error, "HTTP/3 request parse failed");
|
|
continue;
|
|
}
|
|
};
|
|
if request.method() == http::Method::CONNECT && request.uri().path() == mtp_path {
|
|
if request.extensions().get::<h3::ext::Protocol>()
|
|
!= Some(&h3::ext::Protocol::WEB_TRANSPORT)
|
|
{
|
|
let _ = stream
|
|
.send_response(
|
|
Response::builder()
|
|
.status(StatusCode::METHOD_NOT_ALLOWED)
|
|
.body(())
|
|
.unwrap(),
|
|
)
|
|
.await;
|
|
let _ = stream.finish().await;
|
|
continue;
|
|
}
|
|
let session = match h3_webtransport::server::WebTransportSession::accept(
|
|
request, stream, h3,
|
|
)
|
|
.await
|
|
{
|
|
Ok(session) => Arc::new(session),
|
|
Err(error) => {
|
|
tracing::debug!(%error, "WebTransport session accept failed");
|
|
return;
|
|
}
|
|
};
|
|
tokio::spawn(run_session_requests(
|
|
session.clone(),
|
|
router.clone(),
|
|
max_request_body,
|
|
request_timeout,
|
|
metrics.clone(),
|
|
remote_addr,
|
|
));
|
|
let mtp_tx = mtp_tx.clone();
|
|
let auth_semaphore = auth_semaphore.clone();
|
|
let host_config = host_config.clone();
|
|
let connection = connection.clone();
|
|
tokio::spawn(async move {
|
|
let result =
|
|
accept_web_connection(session, mtp_path, connection, send_pongs, policy, host_config, auth_semaphore)
|
|
.await;
|
|
match mtp_tx.try_send(result) {
|
|
Ok(()) => {}
|
|
Err(tokio::sync::mpsc::error::TrySendError::Full(result)) => {
|
|
tracing::warn!("MTP connection backlog is full; dropping connection");
|
|
if let Ok(connection) = result {
|
|
connection.sender.close();
|
|
}
|
|
}
|
|
Err(tokio::sync::mpsc::error::TrySendError::Closed(result)) => {
|
|
if let Ok(connection) = result {
|
|
connection.sender.close();
|
|
}
|
|
}
|
|
}
|
|
});
|
|
return;
|
|
}
|
|
let router = router.clone();
|
|
let metrics = metrics.clone();
|
|
tasks.spawn(async move {
|
|
let path = request.uri().path().to_string();
|
|
let response = crate::http::run_request(
|
|
&path,
|
|
request_timeout,
|
|
metrics.as_ref(),
|
|
handle_http_request(request, &mut stream, &router, max_request_body, remote_addr),
|
|
).await;
|
|
if let Err(error) = write_response(&mut stream, response).await
|
|
&& let Some(metrics) = &metrics
|
|
{
|
|
metrics.error_occurred(&WebServerError::Http(format!("HTTP/3 response write failed: {error}")));
|
|
}
|
|
});
|
|
}
|
|
tasks.join_all().await;
|
|
if let Some(ref m) = metrics {
|
|
m.connection_closed(connect_start.elapsed(), "normal");
|
|
}
|
|
});
|
|
}
|
|
}
|
|
}
|
|
|
|
// --- Drain phase: wait for in-flight connections ---
|
|
|
|
endpoint.close(quinn::VarInt::from_u32(0), b"mtp-webserver shutdown");
|
|
|
|
let drain_start = std::time::Instant::now();
|
|
while !connection_tasks.is_empty() {
|
|
tokio::select! {
|
|
Some(result) = connection_tasks.join_next() => {
|
|
if let Err(e) = result {
|
|
tracing::warn!("Connection task panicked: {}", e);
|
|
}
|
|
}
|
|
_ = tokio::time::sleep(drain_timeout.saturating_sub(drain_start.elapsed())) => {
|
|
tracing::warn!(
|
|
"Drain timeout expired with {} connections still in flight",
|
|
connection_tasks.len()
|
|
);
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
connection_tasks.shutdown().await;
|
|
}
|
|
|
|
async fn handle_http_request<S>(
|
|
request: Request<()>,
|
|
stream: &mut h3::server::RequestStream<S, Bytes>,
|
|
router: &Router,
|
|
max_request_body: usize,
|
|
remote_addr: SocketAddr,
|
|
) -> Result<HttpResponse, WebServerError>
|
|
where
|
|
S: h3::quic::BidiStream<Bytes>,
|
|
{
|
|
let (request, too_large) = read_request(request, stream, max_request_body, remote_addr)
|
|
.await
|
|
.map_err(|e| WebServerError::Http(format!("request body read failed: {e}")))?;
|
|
if too_large {
|
|
return Err(WebServerError::PayloadTooLarge);
|
|
}
|
|
Ok(crate::http::dispatch_request(request, router).await)
|
|
}
|
|
|
|
async fn run_session_requests(
|
|
session: Arc<h3_webtransport::server::WebTransportSession<h3_quinn::Connection, Bytes>>,
|
|
router: Router,
|
|
max_request_body: usize,
|
|
request_timeout: Duration,
|
|
metrics: Option<Arc<dyn WebServerMetrics>>,
|
|
remote_addr: SocketAddr,
|
|
) {
|
|
loop {
|
|
match session.accept_bi().await {
|
|
Ok(Some(h3_webtransport::server::AcceptedBi::Request(request, mut stream))) => {
|
|
let router = router.clone();
|
|
let metrics = metrics.clone();
|
|
tokio::spawn(async move {
|
|
let path = request.uri().path().to_string();
|
|
let response = crate::http::run_request(
|
|
&path,
|
|
request_timeout,
|
|
metrics.as_ref(),
|
|
handle_http_request(
|
|
request,
|
|
&mut stream,
|
|
&router,
|
|
max_request_body,
|
|
remote_addr,
|
|
),
|
|
)
|
|
.await;
|
|
if let Err(error) = write_response(&mut stream, response).await
|
|
&& let Some(metrics) = &metrics
|
|
{
|
|
metrics.error_occurred(&WebServerError::Http(format!(
|
|
"HTTP/3 response write failed: {error}"
|
|
)));
|
|
}
|
|
});
|
|
}
|
|
Ok(Some(h3_webtransport::server::AcceptedBi::BidiStream(_, _))) => {}
|
|
Ok(None) | Err(_) => break,
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn read_request<S>(
|
|
request: Request<()>,
|
|
stream: &mut h3::server::RequestStream<S, Bytes>,
|
|
max_body: usize,
|
|
remote_addr: SocketAddr,
|
|
) -> Result<(HttpRequest, bool), h3::error::StreamError>
|
|
where
|
|
S: h3::quic::BidiStream<Bytes>,
|
|
{
|
|
let (parts, _) = request.into_parts();
|
|
let mut body = Vec::new();
|
|
let mut too_large = false;
|
|
while let Some(chunk) = stream.recv_data().await? {
|
|
if body.len().saturating_add(chunk.remaining()) > max_body {
|
|
too_large = true;
|
|
break;
|
|
}
|
|
body.extend_from_slice(chunk.chunk());
|
|
}
|
|
Ok((
|
|
HttpRequest {
|
|
method: parts.method,
|
|
uri: parts.uri,
|
|
headers: parts.headers,
|
|
body: (!body.is_empty()).then(|| Bytes::from(body)),
|
|
remote_addr,
|
|
},
|
|
too_large,
|
|
))
|
|
}
|
|
|
|
async fn write_response<S>(
|
|
stream: &mut h3::server::RequestStream<S, Bytes>,
|
|
response: HttpResponse,
|
|
) -> Result<(), h3::error::StreamError>
|
|
where
|
|
S: h3::quic::BidiStream<Bytes>,
|
|
{
|
|
let mut builder = Response::builder().status(response.status);
|
|
for (name, value) in &response.headers {
|
|
builder = builder.header(name, value);
|
|
}
|
|
stream
|
|
.send_response(builder.body(()).expect("valid HTTP response"))
|
|
.await?;
|
|
for chunk in response.body {
|
|
stream.send_data(chunk).await?;
|
|
}
|
|
if let Some(mut chunks) = response.stream {
|
|
while let Some(chunk) = chunks.recv().await {
|
|
stream.send_data(chunk).await?;
|
|
}
|
|
}
|
|
stream.finish().await
|
|
}
|