From f87fdcc65008fb40e8f0f37f33f6bbae1afa162a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Roman=20Smr=C5=BE?= <roman.smrz@seznam.cz>
Date: Wed, 16 Jul 2025 21:46:56 +0200
Subject: Close streams used for tunnel on connection drop

---
 src/Erebos/Discovery.hs | 3 +++
 src/Erebos/Network.hs   | 6 ++++++
 2 files changed, 9 insertions(+)

(limited to 'src')

diff --git a/src/Erebos/Discovery.hs b/src/Erebos/Discovery.hs
index b27c617..5788ab8 100644
--- a/src/Erebos/Discovery.hs
+++ b/src/Erebos/Discovery.hs
@@ -569,6 +569,9 @@ instance PeerAddressType TunnelAddress where
     sendBytesToAddress TunnelAddress {..} bytes = do
         writeStream tunnelWriter bytes
 
+    connectionToAddressClosed TunnelAddress {..} = do
+        closeStream tunnelWriter
+
 relayStream :: StreamReader -> StreamWriter -> IO ()
 relayStream r w = do
     p <- readStreamPacket r
diff --git a/src/Erebos/Network.hs b/src/Erebos/Network.hs
index a90f5d8..6126e60 100644
--- a/src/Erebos/Network.hs
+++ b/src/Erebos/Network.hs
@@ -163,6 +163,8 @@ instance Eq Peer where
 
 class (Eq addr, Ord addr, Show addr, Typeable addr) => PeerAddressType addr where
     sendBytesToAddress :: addr -> ByteString -> IO ()
+    connectionToAddressClosed :: addr -> IO ()
+    connectionToAddressClosed _ = return ()
 
 data PeerAddress
     = forall addr. PeerAddressType addr => CustomPeerAddress addr
@@ -384,6 +386,10 @@ startServer serverOptions serverOrigHead logd' serverServices = do
                                         handlePacket identity secure peer chanSvc svcs header prefs
                                         peerLoop
                                     Nothing -> do
+                                        case paddr of
+                                            DatagramAddress _ -> return ()
+                                            CustomPeerAddress caddr -> connectionToAddressClosed caddr
+
                                         dropPeer peer
                                         atomically $ writeTChan serverChanPeer peer
                             peerLoop
-- 
cgit v1.2.3