1 files changed, 16 insertions, 5 deletions

diff --git a/dht/src/Network/QueryResponse.hs b/dht/src/Network/QueryResponse.hs
index c4ff50e3..f72763de 100644
--- a/dht/src/Network/QueryResponse.hs
+++ b/dht/src/Network/QueryResponse.hs
@@ -553,9 +553,14 @@ sockAddrFamily _                       = AF_CAN -- SockAddrCan constructor depre
 -- | Packets with an empty payload may trigger EOF exception.
 -- 'udpTransport' uses this function to avoid throwing in that
 -- case.
-ignoreEOF :: a -> IOError -> IO a
-ignoreEOF def e | isEOFError e = pure def
-                | otherwise    = throwIO e
+ignoreEOF :: Socket -> MVar () -> a -> IOError -> IO (Maybe a)
+ignoreEOF sock isClosed def e = do
+    done <- tryReadMVar isClosed
+    case done of
+        Just () -> do close sock
+                      pure Nothing
+        _ -> if isEOFError e then pure $ Just def
+                             else throwIO e
 
 -- | Hard-coded maximum packet size for incoming UDP Packets received via
 -- 'udpTransport'.
@@ -589,9 +594,10 @@ udpTransport' bind_address = do
     setSocketOption sock IPv6Only 0
   setSocketOption sock Broadcast 1
   bind sock bind_address
+  isClosed <- newEmptyMVar
   let tr = Transport {
       awaitMessage = \kont -> do
-        r <- handle (ignoreEOF $ Just $ Right (B.empty, SockAddrInet 0 0)) $ do
+        r <- handle (ignoreEOF sock isClosed $ Right (B.empty, SockAddrInet 0 0)) $ do
                 Just . Right <$!> B.recvFrom sock udpBufferSize
         kont $! r
     , sendMessage = case family of
@@ -609,7 +615,12 @@ udpTransport' bind_address = do
             addr@(SockAddrInet6 {}) -> \bs -> dput XMisc ("Discarding packet to "++show addr)
             addr4 -> \bs -> saferSendTo sock bs addr4
         _ -> \addr bs -> saferSendTo sock bs addr
-    , closeTransport = close sock
+    , closeTransport = do
+        -- close sock
+        tryTakeMVar isClosed >> putMVar isClosed ()
+        -- set O_NONBLOCK using fcntl
+        setNonBlockIfNeeded (fdSocket sock) -- setSocketOption sock ??? 0
+
     }
   return (tr, sock)