From b1ed2c106458573cd3211f7f6d6eb74b5965c393 Mon Sep 17 00:00:00 2001
From: Nikkuss <Scuggo@pm.me>
Date: Fri, 21 Aug 2026 23:25:39 +0400
Subject: [PATCH] fix bug

---
 flake.lock   | 18 +++++++-------
 src/relay.rs | 69 +++++++++++++++++++++++++++++++++++++++++++++++++---
 2 files changed, 74 insertions(+), 13 deletions(-)

diff --git a/flake.lock b/flake.lock
index 4700871..e349791 100644
--- a/flake.lock
+++ b/flake.lock
@@ -2,11 +2,11 @@
   "nodes": {
     "crane": {
       "locked": {
-        "lastModified": 1774313767,
-        "narHash": "sha256-hy0XTQND6avzGEUFrJtYBBpFa/POiiaGBr2vpU6Y9tY=",
+        "lastModified": 1776635034,
+        "narHash": "sha256-OEOJrT3ZfwbChzODfIH4GzlNTtOFuZFWPtW7jIeR8xU=",
         "owner": "ipetkov",
         "repo": "crane",
-        "rev": "3d9df76e29656c679c744968b17fbaf28f0e923d",
+        "rev": "dc7496d8ea6e526b1254b55d09b966e94673750f",
         "type": "github"
       },
       "original": {
@@ -35,11 +35,11 @@
     },
     "nixpkgs": {
       "locked": {
-        "lastModified": 1774106199,
-        "narHash": "sha256-US5Tda2sKmjrg2lNHQL3jRQ6p96cgfWh3J1QBliQ8Ws=",
+        "lastModified": 1776877367,
+        "narHash": "sha256-EHq1/OX139R1RvBzOJ0aMRT3xnWyqtHBRUBuO1gFzjI=",
         "owner": "nixos",
         "repo": "nixpkgs",
-        "rev": "6c9a78c09ff4d6c21d0319114873508a6ec01655",
+        "rev": "0726a0ecb6d4e08f6adced58726b95db924cef57",
         "type": "github"
       },
       "original": {
@@ -64,11 +64,11 @@
         ]
       },
       "locked": {
-        "lastModified": 1774408260,
-        "narHash": "sha256-Jn9d9r85dmf3gTMnSRt6t+DP2nQ5uJns/MMXg2FpzfM=",
+        "lastModified": 1777086717,
+        "narHash": "sha256-vEl3cGHRxEFdVNuP9PbrhAWnmU98aPOLGy9/1JXzSuM=",
         "owner": "oxalica",
         "repo": "rust-overlay",
-        "rev": "d6471ee5a8f470251e6e5b83a20a182eb6c46c9b",
+        "rev": "3be56bd430bfd65d3c468a50626c3a601c7dee03",
         "type": "github"
       },
       "original": {
diff --git a/src/relay.rs b/src/relay.rs
index a18733b..7a90927 100644
--- a/src/relay.rs
+++ b/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,