From 92c01bfc99f5781ff37fbb5a98a73af376b9af44 Mon Sep 17 00:00:00 2001
From: Joe Crayne <joe@jerkface.net>
Date: Tue, 3 Dec 2019 23:45:23 -0500
Subject: Better TCP cleanup.

---
 dht/src/Network/QueryResponse/TCP.hs | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

(limited to 'dht/src/Network/QueryResponse/TCP.hs')

diff --git a/dht/src/Network/QueryResponse/TCP.hs b/dht/src/Network/QueryResponse/TCP.hs
index defdf6a2..67c19512 100644
--- a/dht/src/Network/QueryResponse/TCP.hs
+++ b/dht/src/Network/QueryResponse/TCP.hs
@@ -21,7 +21,7 @@ import Data.Time.Clock.POSIX
 import Data.Word
 import Data.String (IsString(..))
 import Network.BSD
-import Network.Socket
+import Network.Socket as Socket
 import System.Timeout
 import System.IO
 import System.IO.Error
@@ -94,23 +94,25 @@ acquireConnection mvar tcpcache stream addr bDoCon = do
             Nothing | bDoCon -> writeTVar (lru tcpcache)
                                 $ MM.insert' (TCPAddress $ streamAddr stream addr) PendingTCPSession (Down now) c
                     | otherwise -> return ()
-            Just (tm, v) -> modifyTVar' (lru tcpcache) $ MM.insert' (TCPAddress $ streamAddr stream addr) v (Down now)
+            Just (tm, v) -> writeTVar (lru tcpcache)
+                                $ MM.insert' (TCPAddress $ streamAddr stream addr) v (Down now) c
         return v
     -- dput XTCP $ "acquireConnection 1 " ++ show (streamAddr stream addr, fmap (second showStat) entry)
     case entry of
         Nothing -> fmap join $ forM (guard bDoCon) $ \() -> do
             proto <- getProtocolNumber "tcp"
+            sock <- socket (socketFamily $ streamAddr stream addr) Stream proto
             mh <- catchIOError (do h <- timeout tcp_timeout $ do
-                                       sock <- socket (socketFamily $ streamAddr stream addr) Stream proto
                                        connect sock (streamAddr stream addr) `catchIOError` (\e -> close sock)
                                        h <- socketToHandle sock ReadWriteMode
                                        hSetBuffering h NoBuffering
                                        return h
                                    return h)
                                $ \e -> return Nothing
-            when (isNothing mh)
-              $ atomically $ modifyTVar' (lru tcpcache)
+            when (isNothing mh) $ do
+                atomically $ modifyTVar' (lru tcpcache)
                            $ MM.delete (TCPAddress $ streamAddr stream addr)
+                Socket.close sock
             ret <- fmap join $ forM mh $ \h -> do
              mst <- catchIOError (Just <$> streamHello stream addr h)
                                  (\e -> return Nothing)
@@ -147,7 +149,7 @@ acquireConnection mvar tcpcache stream addr bDoCon = do
                             case mreport of
                                 Just treport -> dput XTCP treport
                                 Nothing      -> dput XTCP "TCP ERROR: threadReport timed out."
-                            hClose h
+                            hClose h `catchIOError` \e -> return ()
                 let v = TCPSession
                         { tcpHandle = h
                         , tcpState  = st
@@ -167,6 +169,7 @@ acquireConnection mvar tcpcache stream addr bDoCon = do
                     case r of TCPSession {tcpState=st,tcpHandle=h} -> do
                                     streamGoodbye st
                                     hClose h
+                                `catchIOError` \e -> return ()
                               _ -> return ()
 
                 return $ Just $ streamEncode st
@@ -190,7 +193,8 @@ closeAll tcpcache stream = do
     cache <- atomically $ swapTVar (lru tcpcache) MM.empty
     forM_ (MM.toList cache) $ \(MM.Binding (TCPAddress addr) r tm) -> do
         killSession r
-        case r of TCPSession{tcpState=st,tcpHandle=h} -> streamGoodbye st >> hClose h
+        case r of TCPSession{tcpState=st,tcpHandle=h} -> catchIOError (streamGoodbye st >> hClose h)
+                                                                      (\e -> return ())
                   _                                   -> return ()
 
 -- Use a cache of TCP client connections for sending (and receiving) packets.
-- 
cgit v1.2.3