From e352c25b32a9f5fffdb48b5734b02211c46a02cd Mon Sep 17 00:00:00 2001
From: Simon Wicky <simon@nymtech.net>
Date: Mon, 25 Mar 2024 11:31:52 +0100
Subject: [PATCH] mark packet_router's shutdown as success before drop (#4491)

---
 common/client-libs/gateway-client/src/packet_router.rs | 4 ++++
 common/client-libs/gateway-client/src/socket_state.rs  | 3 ++-
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/common/client-libs/gateway-client/src/packet_router.rs b/common/client-libs/gateway-client/src/packet_router.rs
index bf73ecb8ca..cd3cc23d51 100644
--- a/common/client-libs/gateway-client/src/packet_router.rs
+++ b/common/client-libs/gateway-client/src/packet_router.rs
@@ -69,6 +69,10 @@ impl PacketRouter {
         }
         Ok(())
     }
+
+    pub fn mark_as_success(&mut self) {
+        self.shutdown.mark_as_success();
+    }
 }
 
 impl GatewayPacketRouter for PacketRouter {
diff --git a/common/client-libs/gateway-client/src/socket_state.rs b/common/client-libs/gateway-client/src/socket_state.rs
index e834fe3eee..e58ebbcabd 100644
--- a/common/client-libs/gateway-client/src/socket_state.rs
+++ b/common/client-libs/gateway-client/src/socket_state.rs
@@ -97,7 +97,7 @@ impl PartiallyDelegated {
 
     pub(crate) fn split_and_listen_for_mixnet_messages(
         conn: WsConn,
-        packet_router: PacketRouter,
+        mut packet_router: PacketRouter,
         shared_key: Arc<SharedKeys>,
         mut shutdown: TaskClient,
     ) -> Self {
@@ -140,6 +140,7 @@ impl PartiallyDelegated {
             if match ret_err {
                 Err(err) => stream_sender.send(Err(err)),
                 Ok(_) => {
+                    packet_router.mark_as_success();
                     shutdown.mark_as_success();
                     stream_sender.send(Ok(stream))
                 }