1 files changed, 11 insertions, 12 deletions

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 ()