3 files changed, 32 insertions, 29 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
diff --git a/dht/Presence/XMPPServer.hs b/dht/Presence/XMPPServer.hs
index 272f6efe..3bafd33c 100644
--- a/dht/Presence/XMPPServer.hs
+++ b/dht/Presence/XMPPServer.hs
@@ -65,7 +65,7 @@ import Control.Monad.Trans (lift)
 import Control.Monad.IO.Class (MonadIO, liftIO)
 import Control.Monad.Fix (fix)
 import Control.Monad
-import Control.Concurrent.ThreadUtil (forkIO,myThreadId,forkLabeled,labelThread,ThreadId,MVar,putMVar,takeMVar,newMVar)
+import Control.Concurrent.ThreadUtil (forkIO,myThreadId,forkLabeled,labelThread,ThreadId,MVar,putMVar,takeMVar,newMVar,threadDelay)
 import Control.Concurrent.STM
 import Data.List hiding ((\\))
 -- import Control.Concurrent.STM.TChan
@@ -107,6 +107,7 @@ import Stanza.Parse
 import Stanza.Types
 import MUC
 import Chat
+import Network.StreamServer (Local(..), Remote(..))
 
 -- peerport :: PortNumber
 -- peerport = 5269
@@ -117,9 +118,6 @@ my_uuid :: Text
 my_uuid = "154ae29f-98f2-4af4-826d-a40c8a188574"
 
 
-newtype Local a = Local a deriving (Eq,Ord,Show)
-newtype Remote a = Remote a deriving (Eq,Ord,Show)
-
 data XMPPServerParameters =
     XMPPServerParameters
     { -- | Called when a client requests a resource id.  The first Maybe indicates
@@ -1090,13 +1088,15 @@ data PeerState
     | PeerConnected (TChan Stanza)
 -}
 
-peerKey :: SocketLike sock => Maybe SockAddr -> sock -> IO (PeerAddress,ConnectionData)
-peerKey bind_addr sock = do
+peerKey :: SocketLike sock => Maybe SockAddr -> (sock, (Local SockAddr, Remote SockAddr)) -> IO (PeerAddress,ConnectionData)
+peerKey bind_addr (sock,(laddr,Remote raddr)) = do
+    {-
     laddr <- getSocketName sock
     raddr <-
       isValidSocket sock >>= \(sock,c) ->
         if c then getPeerName sock -- addr is normally socketName
              else return laddr     -- Weird hack: addr is would-be peer name
+    -}
     -- Assume remote peers are listening on the same port that we do.
     let peerport = fromIntegral $ fromMaybe 5269 $ do
                         p <- bind_addr >>= sockAddrPort
@@ -1106,15 +1106,15 @@ peerKey bind_addr sock = do
     rname <- atomically $ newTVar Nothing
     -- dput XMan $ "peerKey " ++ show (PeerAddress $ raddr `withPort` peerport,laddr)
     return $ ( PeerAddress $ raddr `withPort` peerport
-             , ConnectionData { cdAddr       = Left (Local laddr)
+             , ConnectionData { cdAddr       = Left laddr
                               , cdType       = XMPP
                               , cdProfile    = "."
                               , cdRemoteName = rname } )
 
-clientKey :: SocketLike sock => sock -> IO (PeerAddress,ConnectionData)
-clientKey sock = do
-    laddr <- getSocketName sock -- [::1]:5222  bind address, same for all clients
-    raddr <- getPeerName sock   -- [::1]:????? unique key
+clientKey :: SocketLike sock => (sock, (Local SockAddr,Remote SockAddr)) -> IO (PeerAddress,ConnectionData)
+clientKey (sock,(laddr,Remote raddr)) = do
+    -- laddr <- getSocketName sock -- [::1]:5222  bind address, same for all clients
+    -- raddr <- getPeerName sock   -- [::1]:????? unique key
     when (Just 0 == sockAddrPort raddr) $ do
         dput XMan $ unwords [ "BUG: XMPP Client"
                             , show (laddr,raddr)
diff --git a/dht/src/Network/StreamServer.hs b/dht/src/Network/StreamServer.hs
index 37f8812a..1da612ce 100644
--- a/dht/src/Network/StreamServer.hs
+++ b/dht/src/Network/StreamServer.hs
@@ -13,6 +13,8 @@ module Network.StreamServer
     , quitListening
   --, dummyServerHandle
     , listenSocket
+    , Local(..)
+    , Remote(..)
     ) where
 
 import Data.Monoid
@@ -89,15 +91,18 @@ bshow e = show e
 warnStderr :: String -> IO ()
 warnStderr str = dput XMisc str >> hFlush stderr
 
+newtype Local a = Local a deriving (Eq,Ord,Show)
+newtype Remote a = Remote a deriving (Eq,Ord,Show)
+
 data ServerConfig = ServerConfig
     { serverWarn    :: String -> IO ()
     -- ^ Action to report warnings and errors.
-    , serverSession :: RestrictedSocket -> Int -> Handle -> IO ()
+    , serverSession :: ( RestrictedSocket, (Local SockAddr, Remote SockAddr)) -> Int -> Handle -> IO ()
     -- ^ Action to handle interaction with a client
     }
 
 -- | Initialize a 'ServerConfig' using the provided session handler.
-withSession :: (RestrictedSocket -> Int -> Handle -> IO ()) -> ServerConfig
+withSession :: ((RestrictedSocket,(Local SockAddr,Remote SockAddr)) -> Int -> Handle -> IO ()) -> ServerConfig
 withSession session = ServerConfig warnStderr session
 
 -- | Launch a thread to listen at the given bind address and dispatch
@@ -138,19 +143,13 @@ streamServer cfg addrs = do
 -- socket must be closed by 'quitListening'.
 acceptLoop :: ServerConfig -> Socket -> Int -> IO ()
 acceptLoop cfg sock n = handle (acceptException cfg n sock) $ do
-    con <- accept sock -- TODO: We need to remember the (snd con) peer address here!!!
-                       -- See also: "-- Weird hack: addr is would-be peer name" in XMPPServer.hs
-                       -- If we remember the peer address here, we won't need that weird
-                       -- hack or to call sIsConnected.
-                       -- Probably we should move the
-                       --   newtype Local a = Local a deriving (Eq,Ord,Show)
-                       --   newtype Remote a = Remote a deriving (Eq,Ord,Show)
-                       -- defines so that they are accessible to this module.
+    (con,raddr) <- accept sock
     let conkey = n + 1
-    h <- socketToHandle (fst con) ReadWriteMode
+    laddr <- Socket.getSocketName con
+    h <- socketToHandle con ReadWriteMode
     forkIO $ do
         myThreadId >>= flip labelThread "StreamServer.session"
-        serverSession cfg (restrictHandleSocket h (fst con)) conkey h
+        serverSession cfg (restrictHandleSocket h con, (Local laddr, Remote raddr)) conkey h
     acceptLoop cfg sock (n + 1)
 
 acceptException :: ServerConfig -> Int -> Socket -> IOError -> IO ()