From 4993e60385147901833e5bb97ef0851fc2b14d00 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Philipp=20Kr=C3=BCger?= <philipp.krueger1@gmail.com>
Date: Fri, 6 Dec 2024 18:12:30 +0100
Subject: [PATCH] fix: Actually close the connections properly

---
 src/doctor.rs | 44 +++++++++++++++++++++++++++++++++-----------
 1 file changed, 33 insertions(+), 11 deletions(-)

diff --git a/src/doctor.rs b/src/doctor.rs
index e932e64..43ee212 100644
--- a/src/doctor.rs
+++ b/src/doctor.rs
@@ -500,6 +500,13 @@ async fn active_side(
             let _d = echo_test(&connection, config, pb).await?;
         }
     }
+
+    // Close the connection gracefully.
+    // We're always the ones last receiving data, because
+    // `echo_test` waits for data on the connection as the last thing.
+    connection.close(0u32.into(), b"done");
+    connection.closed().await;
+
     Ok(())
 }
 
@@ -591,19 +598,34 @@ async fn recv_test(
 
 /// Accepts connections and answers requests (echo, drain or send) as passive side.
 async fn passive_side(gui: Gui, connection: Connection) -> anyhow::Result<()> {
-    loop {
-        match connection.accept_bi().await {
-            Ok((send, recv)) => {
-                if let Err(cause) = handle_test_request(send, recv, &gui).await {
-                    eprintln!("Error handling test request {cause}");
+    let conn = connection.clone();
+    let accept_loop = async move {
+        let result = loop {
+            match conn.accept_bi().await {
+                Ok((send, recv)) => {
+                    if let Err(cause) = handle_test_request(send, recv, &gui).await {
+                        eprintln!("Error handling test request {cause}");
+                    }
                 }
-            }
-            Err(cause) => {
-                eprintln!("error accepting bidi stream {cause}");
-                break Err(cause.into());
-            }
+                Err(cause) => {
+                    eprintln!("error accepting bidi stream {cause}");
+                    break Err(cause.into());
+                }
+            };
         };
-    }
+
+        conn.close(0u32.into(), b"internal err");
+        conn.closed().await;
+        eprintln!("Connection closed.");
+
+        result
+    };
+    let conn_closed = async move {
+        connection.closed().await;
+        eprintln!("Connection closed.");
+        anyhow::Ok(())
+    };
+    futures_lite::future::race(conn_closed, accept_loop).await
 }
 
 /// Configures a relay map with some default values.