117 lines
3.7 KiB
Rust
117 lines
3.7 KiB
Rust
use color_eyre::Result;
|
|
use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
|
|
use tracing::debug;
|
|
|
|
/// Bidirectional relay between two async streams.
|
|
/// Returns once *both* directions have finished.
|
|
///
|
|
/// The two halves are deliberately independent: `copy_bidirectional` aborts the
|
|
/// whole relay as soon as one direction errors, which loses an HTTP response
|
|
/// that is already in flight whenever the backend answers early and closes
|
|
/// without draining the request body (e.g. a 401 on a large upload). Each
|
|
/// direction here runs to completion and shuts its own writer down, so the
|
|
/// response half still drains — and the QUIC send stream is finished cleanly
|
|
/// instead of being reset by drop.
|
|
pub async fn relay<A, B>(a: A, b: B) -> Result<()>
|
|
where
|
|
A: AsyncRead + AsyncWrite + Unpin,
|
|
B: AsyncRead + AsyncWrite + Unpin,
|
|
{
|
|
let (mut ar, mut aw) = tokio::io::split(a);
|
|
let (mut br, mut bw) = tokio::io::split(b);
|
|
|
|
let a_to_b = async {
|
|
let r = tokio::io::copy(&mut ar, &mut bw).await;
|
|
let _ = bw.shutdown().await;
|
|
r
|
|
};
|
|
let b_to_a = async {
|
|
let r = tokio::io::copy(&mut br, &mut aw).await;
|
|
let _ = aw.shutdown().await;
|
|
r
|
|
};
|
|
|
|
let (ab, ba) = tokio::join!(a_to_b, b_to_a);
|
|
if let Err(e) = ab {
|
|
debug!("relay a->b ended: {e}");
|
|
}
|
|
if let Err(e) = ba {
|
|
debug!("relay b->a ended: {e}");
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
|
|
|
/// The upload-502 regression: the backend answers and hangs up without
|
|
/// draining the request body, so the relay's request half fails on its
|
|
/// first write. The response half must still deliver.
|
|
#[tokio::test]
|
|
async fn response_survives_backend_early_close() {
|
|
const RESPONSE: &[u8] = b"HTTP/1.1 401 Unauthorized\r\n\r\n";
|
|
|
|
let (a, a_peer) = tokio::io::duplex(1024);
|
|
let (b, mut b_peer) = tokio::io::duplex(1024);
|
|
|
|
// Backend: reply, then hang up with the request body still unread.
|
|
b_peer.write_all(RESPONSE).await.unwrap();
|
|
drop(b_peer);
|
|
|
|
// Downstream: request bytes already buffered, so the request half has
|
|
// something to copy on its very first poll.
|
|
let (mut a_rd, mut a_wr) = tokio::io::split(a_peer);
|
|
a_wr.write_all(&[0u8; 512]).await.unwrap();
|
|
|
|
relay(a, b).await.unwrap();
|
|
|
|
let mut got = Vec::new();
|
|
a_rd.read_to_end(&mut got).await.unwrap();
|
|
assert_eq!(got, RESPONSE);
|
|
}
|
|
}
|
|
|
|
/// Wrapper to combine a QUIC send+recv into a single AsyncRead+AsyncWrite.
|
|
pub struct QuicBiStream {
|
|
pub send: quinn::SendStream,
|
|
pub recv: quinn::RecvStream,
|
|
}
|
|
|
|
impl tokio::io::AsyncRead for QuicBiStream {
|
|
fn poll_read(
|
|
mut self: std::pin::Pin<&mut Self>,
|
|
cx: &mut std::task::Context<'_>,
|
|
buf: &mut tokio::io::ReadBuf<'_>,
|
|
) -> std::task::Poll<std::io::Result<()>> {
|
|
std::pin::Pin::new(&mut self.recv).poll_read(cx, buf)
|
|
}
|
|
}
|
|
|
|
impl tokio::io::AsyncWrite for QuicBiStream {
|
|
fn poll_write(
|
|
mut self: std::pin::Pin<&mut Self>,
|
|
cx: &mut std::task::Context<'_>,
|
|
buf: &[u8],
|
|
) -> std::task::Poll<std::io::Result<usize>> {
|
|
std::pin::Pin::new(&mut self.send)
|
|
.poll_write(cx, buf)
|
|
.map(|r| r.map_err(std::io::Error::other))
|
|
}
|
|
|
|
fn poll_flush(
|
|
mut self: std::pin::Pin<&mut Self>,
|
|
cx: &mut std::task::Context<'_>,
|
|
) -> std::task::Poll<std::io::Result<()>> {
|
|
std::pin::Pin::new(&mut self.send).poll_flush(cx)
|
|
}
|
|
|
|
fn poll_shutdown(
|
|
mut self: std::pin::Pin<&mut Self>,
|
|
cx: &mut std::task::Context<'_>,
|
|
) -> std::task::Poll<std::io::Result<()>> {
|
|
std::pin::Pin::new(&mut self.send).poll_shutdown(cx)
|
|
}
|
|
}
|