1 files changed, 10 insertions, 6 deletions

diff --git a/dht/Connection/Tcp.hs b/dht/Connection/Tcp.hs
index 2572eba6..06290a1c 100644
--- a/dht/Connection/Tcp.hs
+++ b/dht/Connection/Tcp.hs
@@ -37,6 +37,7 @@ import           Data.Map (Map)
 import Data.Monoid ( (<>) )
 import Control.Concurrent.ThreadUtil
 
+import Control.Arrow
 import Control.Concurrent.STM
 -- import Control.Concurrent.STM.TMVar
 -- import Control.Concurrent.STM.TChan
@@ -76,6 +77,7 @@ import Network.StreamServer
 import Network.SocketLike hiding (sClose)
 import qualified Connection as G
          ;import Connection (Manager (..), PeerAddress (..), Policy (..))
+import Network.Address (localhost4)
 import DPut
 import DebugTag
 
@@ -96,7 +98,7 @@ data ConnectionParameters conkey u =
     -- ^ The miliseconds of idle after 'RequiresPing' is signaled
     -- that are necessary for the connection to be considered
     -- lost and signalling 'EOF'.
-    , makeConnKey :: RestrictedSocket -> IO (conkey,u)
+    , makeConnKey :: (RestrictedSocket,(Local SockAddr, Remote SockAddr)) -> IO (conkey,u)
     -- ^ This action creates a lookup key for a new connection.  If 'duplex'
     -- is 'True' and the result is already assocatied with an established
     -- connection, then an 'EOF' will be forced before the the new
@@ -122,7 +124,7 @@ data ConnectionParameters conkey u =
 --  * 'duplex' = True
 --
 connectionDefaults
-  :: (RestrictedSocket -> IO (conkey,u)) -> ConnectionParameters conkey u
+  :: ((RestrictedSocket,(Local SockAddr,Remote SockAddr)) -> IO (conkey,u)) -> ConnectionParameters conkey u
 connectionDefaults f = ConnectionParameters
     { pingInterval = 28000
     , timeout = 2000
@@ -337,14 +339,15 @@ server allocate sessionConduits = do
         handle (\e -> do -- let t = ioeGetErrorType e
                          when (isDoesNotExistError e) $ return () -- warn "GOTCHA"
                          -- warn $ "connect-error: " <> bshow e
-                         (conkey,u) <- makeConnKey params (restrictSocket sock) -- XXX: ?
+                         (conkey,u) <- makeConnKey params (restrictSocket sock,(Local localhost4, Remote addr)) -- XXX: ?
                          Socket.close sock
                          atomically
                             $ writeTChan (serverEvent server)
                                 $ ((conkey,u),ConnectFailure addr))
                $ do
         connect sock addr
-        (conkey,u) <- makeConnKey params (restrictSocket sock)
+        laddr <- Socket.getSocketName sock
+        (conkey,u) <- makeConnKey params (restrictSocket sock, (Local laddr, Remote addr))
         h <- socketToHandle sock ReadWriteMode
         newConnection server sessionConduits params conkey u h Out
         return ()
@@ -375,13 +378,14 @@ server allocate sessionConduits = do
             sock <- socket (socketFamily addr) Stream proto
             handle (\(SomeException e) -> do
                 -- Any thing else goes wrong and we broadcast ConnectFailure.
-                do (conkey,u) <- makeConnKey params (restrictSocket sock)
+                do (conkey,u) <- makeConnKey params (restrictSocket sock,(Local localhost4, Remote addr))
                    Socket.close sock
                    atomically $ writeTChan (serverEvent server) ((conkey,u),ConnectFailure addr)
                  `onException` return ()
                 atomically $ readTVar retryVar) $ do
             connect sock addr
-            (conkey,u) <- makeConnKey params (restrictSocket sock)
+            laddr <- Socket.getSocketName sock
+            (conkey,u) <- makeConnKey params (restrictSocket sock, (Local laddr, Remote addr))
             h <- socketToHandle sock ReadWriteMode
             threads <- newConnection server sessionConduits params conkey u h Out
             atomically $ do threadsWait threads