fix bug

src/relay.rs

@@ -1,17 +1,78 @@
 use color_eyre::Result;
-use tokio::io::{AsyncRead, AsyncWrite, copy_bidirectional};
+use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
+use tracing::debug;
 
 /// Bidirectional relay between two async streams.
-/// Returns when either direction hits EOF or an error.
-pub async fn relay<A, B>(mut a: A, mut b: B) -> Result<()>
+/// 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,
 {
-    copy_bidirectional(&mut a, &mut b).await?;
+    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,