3 files changed, 410 insertions, 95 deletions

diff --git a/src/Network/BitTorrent/Exchange/Session.hs b/src/Network/BitTorrent/Exchange/Session.hs
index 565ef7ab..0d4f3d02 100644
--- a/src/Network/BitTorrent/Exchange/Session.hs
+++ b/src/Network/BitTorrent/Exchange/Session.hs
@@ -1,7 +1,9 @@
-{-# LANGUAGE TemplateHaskell    #-}
-{-# LANGUAGE DeriveDataTypeable #-}
+{-# LANGUAGE FlexibleInstances    #-}
+{-# LANGUAGE TemplateHaskell      #-}
+{-# LANGUAGE DeriveDataTypeable   #-}
 module Network.BitTorrent.Exchange.Session
        ( Session
+       , LogFun
        , newSession
        , closeSession
 
@@ -12,76 +14,141 @@ import Control.Applicative
 import Control.Concurrent
 import Control.Exception
 import Control.Lens
+import Control.Monad.Logger
 import Control.Monad.Reader
 import Control.Monad.State
+import Data.ByteString as BS
+import Data.ByteString.Lazy as BL
+import Data.Conduit
 import Data.Function
 import Data.IORef
+import Data.List as L
 import Data.Maybe
 import Data.Map as M
+import Data.Monoid
 import Data.Ord
+import Data.Set as S
+import Data.Text as T
 import Data.Typeable
-import Text.PrettyPrint
+import Text.PrettyPrint hiding ((<>))
+import Text.PrettyPrint.Class
+import System.Log.FastLogger (LogStr, ToLogStr (..))
 
 import Data.Torrent (InfoDict (..))
 import Data.Torrent.Bitfield as BF
 import Data.Torrent.InfoHash
+import Data.Torrent.Piece (pieceData, piPieceLength)
+import qualified Data.Torrent.Piece as Torrent (Piece (Piece))
 import Network.BitTorrent.Core
 import Network.BitTorrent.Exchange.Assembler
-import Network.BitTorrent.Exchange.Block
+import Network.BitTorrent.Exchange.Block as Block
 import Network.BitTorrent.Exchange.Message
+import Network.BitTorrent.Exchange.Session.Status as SS
 import Network.BitTorrent.Exchange.Status
 import Network.BitTorrent.Exchange.Wire
 import System.Torrent.Storage
 
+{-----------------------------------------------------------------------
+--  Exceptions
+-----------------------------------------------------------------------}
+
+data ExchangeError
+  = InvalidRequest BlockIx StorageFailure
+  | CorruptedPiece PieceIx
+    deriving (Show, Typeable)
+
+instance Exception ExchangeError
+
+packException :: Exception e => (e -> ExchangeError) -> IO a -> IO a
+packException f m = try m >>= either (throwIO . f) return
+
+{-----------------------------------------------------------------------
+--  Session
+-----------------------------------------------------------------------}
+
+data ConnectionEntry = ConnectionEntry
+  { initiatedBy :: !ChannelSide
+  , connection  :: !(Connection Session)
+  }
 
 data Session = Session
-  { tpeerId     :: PeerId
-  , infohash    :: InfoHash
-  , bitfield    :: Bitfield
-  , assembler   :: Assembler
-  , storage     :: Storage
-  , unchoked    :: [PeerAddr IP]
-  , connections :: MVar (Map (PeerAddr IP) (Connection Session))
-  , broadcast   :: Chan Message
+  { tpeerId      :: PeerId
+  , infohash     :: InfoHash
+  , storage      :: Storage
+  , status       :: MVar SessionStatus
+  , unchoked     :: [PeerAddr IP]
+  , connections  :: MVar (Map (PeerAddr IP) ConnectionEntry)
+  , broadcast    :: Chan Message
+  , logger       :: LogFun
   }
 
-newSession :: PeerAddr (Maybe IP) -- ^ /external/ address of this peer;
+-- | Logger function.
+type LogFun = Loc -> LogSource -> LogLevel -> LogStr -> IO ()
+
+newSession :: LogFun
+           -> PeerAddr (Maybe IP) -- ^ /external/ address of this peer;
            -> FilePath            -- ^ root directory for content files;
            -> InfoDict            -- ^ torrent info dictionary;
            -> IO Session          -- ^
-newSession addr rootPath dict = do
-  connVar <- newMVar M.empty
-  store   <- openInfoDict ReadWriteEx rootPath dict
-  chan    <- newChan
+newSession logFun addr rootPath dict = do
+  connVar     <- newMVar M.empty
+  store       <- openInfoDict ReadWriteEx rootPath dict
+  statusVar   <- newMVar $ sessionStatus (BF.haveNone (totalPieces store))
+                                         (piPieceLength (idPieceInfo dict))
+  chan        <- newChan
   return Session
     { tpeerId     = fromMaybe (error "newSession: impossible") (peerId addr)
     , infohash    = idInfoHash dict
-    , bitfield    = BF.haveNone (totalPieces store)
-    , assembler   = error "newSession"
+    , status      = statusVar
     , storage     = store
     , unchoked    = []
     , connections = connVar
     , broadcast   = chan
+    , logger      = logFun
     }
 
 closeSession :: Session -> IO ()
 closeSession = undefined
 
+instance MonadIO m => MonadLogger (Connected Session m) where
+  monadLoggerLog loc src lvl msg = do
+    conn <- ask
+    ses  <- asks connSession
+    addr <- asks connRemoteAddr
+    let addrSrc = src <> " @ " <> T.pack (render (pretty addr))
+    liftIO $ logger ses loc addrSrc lvl (toLogStr msg)
+
+logMessage :: Message -> Wire Session ()
+logMessage msg = logDebugN $ T.pack (render (pretty msg))
+
+logEvent :: Text -> Wire Session ()
+logEvent = logInfoN
+
+{-----------------------------------------------------------------------
+--  Connections
+-----------------------------------------------------------------------}
+-- TODO unmap storage on zero connections
+
 insert :: PeerAddr IP
        -> {- Maybe Socket
        -> -} Session -> IO ()
 insert addr ses @ Session {..} = do
-  forkIO $ do
-    let caps  = def
-    let ecaps = def
-    let hs = Handshake def caps infohash tpeerId
-    chan <- dupChan broadcast
-    connectWire ses hs addr ecaps chan $ do
-      conn <- getConnection
+    forkIO $ do
+      action `finally` runStatusUpdates status (resetPending addr)
+    return ()
+  where
+    action = do
+      let caps  = def
+      let ecaps = def
+      let hs = Handshake def caps infohash tpeerId
+      chan <- dupChan broadcast
+      connectWire ses hs addr ecaps chan $ do
+        conn <- getConnection
 --      liftIO $ modifyMVar_ connections $ pure . M.insert addr conn
-      exchange
+        resizeBitfield (totalPieces storage)
+        logEvent "Connection established"
+        exchange
 --      liftIO $ modifyMVar_ connections $ pure . M.delete addr
-  return ()
 
 delete :: PeerAddr IP -> Session -> IO ()
 delete = undefined
@@ -90,81 +157,125 @@ deleteAll :: Session -> IO ()
 deleteAll = undefined
 
 {-----------------------------------------------------------------------
---  Query
+--  Helpers
 -----------------------------------------------------------------------}
 
+withStatusUpdates :: StatusUpdates a -> Wire Session a
+withStatusUpdates m = do
+  Session {..} <- getSession
+  liftIO $ runStatusUpdates status m
+
 getThisBitfield :: Wire Session Bitfield
-getThisBitfield = undefined
+getThisBitfield = do
+  ses <- getSession
+  liftIO $ SS.getBitfield (status ses)
+
+readBlock :: BlockIx -> Storage -> IO (Block BL.ByteString)
+readBlock bix @ BlockIx {..} s = do
+  p <- packException (InvalidRequest bix) $ do readPiece ixPiece s
+  let chunk = BL.take (fromIntegral ixLength) $
+              BL.drop (fromIntegral ixOffset) (pieceData p)
+  if BL.length chunk == fromIntegral ixLength
+    then return  $ Block ixPiece ixOffset chunk
+    else throwIO $ InvalidRequest bix (InvalidSize ixLength)
+
+sendBroadcast :: PeerMessage msg => msg -> Wire Session ()
+sendBroadcast msg = do
+  Session {..} <- getSession
+  ecaps <- getExtCaps
+  liftIO $ writeChan broadcast (envelop ecaps msg)
+
+{-----------------------------------------------------------------------
+--  Triggers
+-----------------------------------------------------------------------}
 
-{-
-data PendingSet = PendingSet (Map (PeerAddr IP) [BlockIx])
+fillRequestQueue :: Wire Session ()
+fillRequestQueue = do
+  maxN <- getAdvertisedQueueLength
+  rbf  <- getRemoteBitfield
+  addr <- connRemoteAddr <$> getConnection
+  blks <- withStatusUpdates $ do
+    n <- getRequestQueueLength addr
+    scheduleBlocks addr rbf (maxN - n)
+  mapM_ (sendMessage . Request) blks
 
-empty :: PendingSet
-empty = undefined
+tryFillRequestQueue :: Wire Session ()
+tryFillRequestQueue = do
+  allowed <- canDownload <$> getStatus
+  when allowed $ do
+    fillRequestQueue
 
-member :: PeerAddr IP -> BlockIx -> PendingSet -> Bool
-member addr bix = undefined
+interesting :: Wire Session ()
+interesting = do
+  addr <- connRemoteAddr <$> getConnection
+  logMessage  (Status (Interested True))
+  sendMessage (Interested True)
+  logMessage  (Status (Choking    False))
+  sendMessage (Choking    False)
+  tryFillRequestQueue
 
-insert :: PeerAddr IP -> BlockIx -> PendingSet -> PendingSet
-insert addr bix = undefined
--}
 {-----------------------------------------------------------------------
---  Event loop
+--  Incoming message handling
 -----------------------------------------------------------------------}
-{-
-data ExchangeError
-  = InvalidRequest BlockIx StorageFailure
-  | CorruptedPiece PieceIx
 
-packException :: Exception e => (e -> ExchangeError) -> IO a -> IO a
-packException f m = try >>= either (throwIO . f) m
+handleStatus :: StatusUpdate -> Wire Session ()
+handleStatus s = do
+  updateConnStatus RemotePeer s
+  case s of
+    Interested _     -> return ()
+    Choking    True  -> do
+      addr <- connRemoteAddr <$> getConnection
+      withStatusUpdates (resetPending addr)
+    Choking    False -> tryFillRequestQueue
 
-readBlock :: BlockIx -> Storage -> IO (Block ByteString)
-readBlock bix @ BlockIx {..} s = do
-  p <- packException (InvalidRequest bix) $ do readPiece ixPiece storage
-  let chunk = BS.take ixLength $ BS.drop ixOffset p
-  if BS.length chunk == ixLength
-    then return chunk
-    else throwIO $ InvalidRequest bix (InvalidSize ixLength)
--}
+handleAvailable :: Available -> Wire Session ()
+handleAvailable msg = do
+  updateRemoteBitfield $ case msg of
+    Have     ix -> BF.insert ix
+    Bitfield bf -> const     bf
+
+  thisBf <- getThisBitfield
+  case msg of
+    Have     ix
+      | ix `BF.member`     thisBf -> return ()
+      |     otherwise             -> interesting
+    Bitfield bf
+      | bf `BF.isSubsetOf` thisBf -> return ()
+      |     otherwise             -> interesting
 
 handleTransfer :: Transfer -> Wire Session ()
 handleTransfer (Request bix) = do
---    Session {..} <- getSession
---    addr <- getRemoteAddr
---    when (addr `elem` unchoked && ixPiece bix `BF.member` bitfield) $ do
---      blk <- liftIO $ readBlock bix storage
---      sendMsg (Piece blk)
-    return ()
+  Session {..} <- getSession
+  bitfield <- getThisBitfield
+  upload   <- canUpload <$> getStatus
+  when (upload && ixPiece bix `BF.member` bitfield) $ do
+    blk <- liftIO $ readBlock bix storage
+    sendMessage (Piece blk)
 
 handleTransfer (Piece   blk) = do
-{-
-    Session {..} <- getSession
-    when (blockIx blk `PS.member` pendingSet) $ do
-      insert blk stalledSet
-      sendBroadcast have
-      maybe send not interested
--}
-    return ()
+  Session {..} <- getSession
+  isSuccess <- withStatusUpdates (pushBlock blk storage)
+  case isSuccess of
+    Nothing -> liftIO $ throwIO $ userError "block is not requested"
+    Just isCompleted -> do
+      when isCompleted $ do
+        sendBroadcast (Have (blkPiece blk))
+--        maybe send not interested
+      tryFillRequestQueue
 
 handleTransfer (Cancel  bix) = filterQueue (not . (transferResponse bix))
   where
     transferResponse bix (Transfer (Piece blk)) = blockIx blk == bix
     transferResponse _    _                     = False
 
+{-----------------------------------------------------------------------
+--  Event loop
+-----------------------------------------------------------------------}
+
 handleMessage :: Message -> Wire Session ()
 handleMessage KeepAlive       = return ()
-handleMessage (Status s)      = undefined
-handleMessage (Available msg) = do
-  thisBf <- getThisBitfield
-  case msg of
-    Have     ix
-      | ix `BF.member`     thisBf -> return ()
-      |     otherwise             -> undefined
-    Bitfield bf
-      | bf `BF.isSubsetOf` thisBf -> return ()
-      |     otherwise             -> undefined
-
+handleMessage (Status s)      = handleStatus s
+handleMessage (Available msg) = handleAvailable msg
 handleMessage (Transfer  msg) = handleTransfer msg
 handleMessage (Port      n)   = undefined
 handleMessage (Fast      _)   = undefined
@@ -172,19 +283,17 @@ handleMessage (Extended  _)   = undefined
 
 exchange :: Wire Session ()
 exchange = do
-  e <- recvMessage
-  liftIO $ print e
-
-type Exchange = StateT Session (ReaderT (Connection Session) IO)
-
---runExchange :: Exchange () -> [PeerAddr] -> IO ()
---runExchange exchange peers = do
---  forM_ peers $ \ peer -> do
---    forkIO $ runReaderT (runStateT exchange session )
+  bf <- getThisBitfield
+  sendMessage (Bitfield bf)
+  awaitForever $ \ msg -> do
+    logMessage msg
+    handleMessage msg
 
 data Event = NewMessage (PeerAddr IP) Message
            | Timeout -- for scheduling
 
+type Exchange a = Wire Session a
+
 awaitEvent :: Exchange Event
 awaitEvent = undefined
 
diff --git a/src/Network/BitTorrent/Exchange/Session/Status.hs b/src/Network/BitTorrent/Exchange/Session/Status.hs
new file mode 100644
index 00000000..565c3bf3
--- /dev/null
+++ b/src/Network/BitTorrent/Exchange/Session/Status.hs
@@ -0,0 +1,175 @@
+module Network.BitTorrent.Exchange.Session.Status
+       ( -- * Environment
+         StatusUpdates
+       , runStatusUpdates
+
+         -- * Status
+       , SessionStatus
+       , sessionStatus
+
+         -- * Query
+       , getBitfield
+       , getRequestQueueLength
+
+         -- * Control
+       , scheduleBlocks
+       , resetPending
+       , pushBlock
+       ) where
+
+import Control.Applicative
+import Control.Concurrent
+import Control.Monad.State
+import Data.ByteString.Lazy as BL
+import Data.Default
+import Data.List as L
+import Data.Maybe
+import Data.Map as M
+import Data.Set as S
+import Data.Tuple
+
+import Data.Torrent.Piece
+import Data.Torrent.Bitfield as BF
+import Network.BitTorrent.Core
+import Network.BitTorrent.Exchange.Block as Block
+import System.Torrent.Storage (Storage, writePiece)
+
+
+{-----------------------------------------------------------------------
+--  Piece entry
+-----------------------------------------------------------------------}
+
+data PieceEntry = PieceEntry
+  { pending :: [(PeerAddr IP, BlockIx)]
+  , stalled :: Bucket
+  }
+
+pieceEntry :: PieceSize -> PieceEntry
+pieceEntry s = PieceEntry [] (Block.empty s)
+
+isEmpty :: PieceEntry -> Bool
+isEmpty PieceEntry {..} = L.null pending && Block.null stalled
+
+holes :: PieceIx -> PieceEntry -> [BlockIx]
+holes pix PieceEntry {..} = fmap mkBlockIx (spans defaultTransferSize stalled)
+  where
+    mkBlockIx (off, sz) = BlockIx pix off sz
+
+{-----------------------------------------------------------------------
+--  Session status
+-----------------------------------------------------------------------}
+
+data SessionStatus = SessionStatus
+  { inprogress :: !(Map PieceIx PieceEntry)
+  , bitfield   :: !Bitfield
+  , pieceSize  :: !PieceSize
+  }
+
+sessionStatus :: Bitfield -> PieceSize -> SessionStatus
+sessionStatus bf ps = SessionStatus
+  { inprogress = M.empty
+  , bitfield   = bf
+  , pieceSize  = ps
+  }
+
+type StatusUpdates a = StateT SessionStatus IO a
+
+-- |
+runStatusUpdates :: MVar SessionStatus -> StatusUpdates a -> IO a
+runStatusUpdates var m = modifyMVar var (fmap swap . runStateT m)
+
+getBitfield :: MVar SessionStatus -> IO Bitfield
+getBitfield var = bitfield <$> readMVar var
+
+getRequestQueueLength :: PeerAddr IP -> StatusUpdates Int
+getRequestQueueLength addr = do
+  m <- gets (M.elems . M.map (L.filter ((==) addr . fst) . pending) . inprogress)
+  return $ L.sum $ L.map L.length m
+
+modifyEntry :: PieceIx -> (PieceEntry -> PieceEntry) -> StatusUpdates ()
+modifyEntry pix f = modify $ \ s @ SessionStatus {..} -> s
+    { inprogress = alter (g pieceSize) pix inprogress }
+  where
+    g s = h . f . fromMaybe (pieceEntry s)
+    h e
+      | isEmpty e = Nothing
+      | otherwise = Just e
+
+{-----------------------------------------------------------------------
+--  Piece download
+-----------------------------------------------------------------------}
+
+-- TODO choose block nearest to pending or stalled sets to reduce disk
+-- seeks on remote machines
+chooseBlocks :: [BlockIx] -> Int -> StatusUpdates [BlockIx]
+chooseBlocks xs n = return (L.take n xs)
+
+-- TODO use selection strategies from Exchange.Selector
+choosePiece :: Bitfield -> StatusUpdates (Maybe PieceIx)
+choosePiece bf
+  | BF.null bf = return $ Nothing
+  | otherwise  = return $ Just $ BF.findMin bf
+
+scheduleBlocks :: PeerAddr IP -> Bitfield -> Int -> StatusUpdates [BlockIx]
+scheduleBlocks addr maskBF n = do
+  SessionStatus {..} <- get
+  let wantPieces = maskBF `BF.difference` bitfield
+  let wantBlocks = L.concat $ M.elems $ M.mapWithKey holes $
+                   M.filterWithKey (\ pix _ -> pix `BF.member` wantPieces) inprogress
+
+  bixs <- if L.null wantBlocks
+    then do
+      mpix <- choosePiece wantPieces
+      case mpix of -- TODO return 'n' blocks
+        Nothing  -> return []
+        Just pix -> return [leadingBlock pix defaultTransferSize]
+    else chooseBlocks wantBlocks n
+
+  forM_ bixs $ \ bix -> do
+    modifyEntry (ixPiece bix) $ \ e @ PieceEntry {..} -> e
+      { pending = (addr, bix) : pending }
+
+  return bixs
+
+
+-- | Remove all pending block requests to the remote peer. May be used
+-- when:
+--
+--     * a peer closes connection;
+--
+--     * remote peer choked this peer;
+--
+--     * timeout expired.
+--
+resetPending :: PeerAddr IP -> StatusUpdates ()
+resetPending addr = modify $ \ s -> s { inprogress = reset (inprogress s) }
+  where
+    reset = fmap $ \ e -> e
+            { pending = L.filter (not . (==) addr . fst) (pending e) }
+
+-- | MAY write to storage, if a new piece have been completed.
+pushBlock :: Block BL.ByteString -> Storage -> StatusUpdates (Maybe Bool)
+pushBlock blk @ Block {..} storage = do
+  mpe <- gets (M.lookup blkPiece . inprogress)
+  case mpe of
+    Nothing -> return Nothing
+    Just (pe @ PieceEntry {..})
+      | blockIx blk `L.notElem` fmap snd pending -> return Nothing
+      |             otherwise                    -> do
+       let bkt' = Block.insertLazy blkOffset blkData stalled
+       case toPiece bkt' of
+         Nothing        -> do
+           modifyEntry blkPiece $ \ e @ PieceEntry {..} -> e
+             { pending = L.filter ((==) (blockIx blk) . snd) pending
+             , stalled = bkt'
+             }
+           return (Just False)
+
+         Just pieceData -> do
+           -- TODO verify
+           liftIO $ writePiece (Piece blkPiece pieceData) storage
+           modify $ \ s @ SessionStatus {..} -> s
+             { inprogress =  M.delete blkPiece inprogress
+             , bitfield   = BF.insert blkPiece bitfield
+             }
+           return (Just True)
diff --git a/src/Network/BitTorrent/Exchange/Wire.hs b/src/Network/BitTorrent/Exchange/Wire.hs
index 1fca6a66..8873546d 100644
--- a/src/Network/BitTorrent/Exchange/Wire.hs
+++ b/src/Network/BitTorrent/Exchange/Wire.hs
@@ -15,7 +15,8 @@
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 module Network.BitTorrent.Exchange.Wire
        ( -- * Wire
-         Wire
+         Connected
+       , Wire
 
          -- ** Exceptions
        , ChannelSide   (..)
@@ -38,12 +39,14 @@ module Network.BitTorrent.Exchange.Wire
 
          -- ** Connection
        , Connection
+       , connRemoteAddr
        , connProtocol
        , connCaps
        , connTopic
        , connRemotePeerId
        , connThisPeerId
        , connOptions
+       , connSession
 
          -- ** Setup
        , runWire
@@ -55,6 +58,7 @@ module Network.BitTorrent.Exchange.Wire
        , recvMessage
        , sendMessage
        , filterQueue
+       , getAdvertisedQueueLength
 
          -- ** Query
        , getConnection
@@ -93,6 +97,7 @@ import Network.Socket.ByteString as BS
 import Text.PrettyPrint as PP hiding (($$), (<>))
 import Text.PrettyPrint.Class
 import Text.Show.Functions
+import System.Log.FastLogger (ToLogStr(..))
 import System.Timeout
 
 import Data.BEncode as BE
@@ -190,13 +195,15 @@ errorPenalty (DisallowedMessage  _ _) = 1
 
 -- | Exceptions used to interrupt the current P2P session.
 data WireFailure
+  = ConnectionRefused IOError
+
     -- | Force termination of wire connection.
     --
     --   Normally you should throw only this exception from event loop
     --   using 'disconnectPeer', other exceptions are thrown
     --   automatically by functions from this module.
     --
-  = DisconnectPeer
+  | DisconnectPeer
 
      -- | A peer not responding and did not send a 'KeepAlive' message
      -- for a specified period of time.
@@ -464,10 +471,12 @@ makeLenses ''ConnectionState
 
 -- | Connection keep various info about both peers.
 data Connection s = Connection
-  { -- | /Both/ peers handshaked with this protocol string. The only
+  { connRemoteAddr   :: !(PeerAddr IP)
+
+    -- | /Both/ peers handshaked with this protocol string. The only
     -- value is \"Bittorrent Protocol\" but this can be changed in
     -- future.
-    connProtocol     :: !ProtocolName
+  , connProtocol     :: !ProtocolName
 
     -- | Set of enabled core extensions, i.e. the pre BEP10 extension
     -- mechanism. This value is used to check if a message is allowed
@@ -503,6 +512,17 @@ data Connection s = Connection
 instance Pretty (Connection s) where
   pretty Connection {..} = "Connection"
 
+instance ToLogStr (Connection s) where
+  toLogStr Connection {..} = mconcat
+    [ toLogStr (show connRemoteAddr)
+    , toLogStr (show connProtocol)
+    , toLogStr (show connCaps)
+    , toLogStr (show connTopic)
+    , toLogStr (show connRemotePeerId)
+    , toLogStr (show connThisPeerId)
+    , toLogStr (show connOptions)
+    ]
+
 -- TODO check extended messages too
 isAllowed :: Connection s -> Message -> Bool
 isAllowed Connection {..} msg
@@ -592,6 +612,15 @@ getConnection = lift ask
 getSession :: Wire s s
 getSession = lift (asks connSession)
 
+-- TODO configurable
+defQueueLength :: Int
+defQueueLength = 1
+
+getAdvertisedQueueLength :: Wire s Int
+getAdvertisedQueueLength = do
+  ExtendedHandshake {..} <- getRemoteEhs
+  return $ fromMaybe defQueueLength ehsQueueLength
+
 {-----------------------------------------------------------------------
 --  Wrapper
 -----------------------------------------------------------------------}
@@ -685,8 +714,9 @@ reconnect = undefined
 --
 connectWire :: s -> Handshake -> PeerAddr IP -> ExtendedCaps -> Chan Message
             -> Wire s () -> IO ()
-connectWire session hs addr extCaps chan wire =
-  bracket (peerSocket Stream addr) close $ \ sock -> do
+connectWire session hs addr extCaps chan wire = do
+  let catchRefusal m = try m >>= either (throwIO . ConnectionRefused) return
+  bracket (catchRefusal (peerSocket Stream addr)) close $ \ sock -> do
     hs' <- initiateHandshake sock hs
 
     Prelude.mapM_ (\(t,e) -> unless t $ throwIO $ ProtocolError e) [
@@ -723,7 +753,8 @@ connectWire session hs addr extCaps chan wire =
       (forkIO $ sourceChan kaInterval chan $= conduitPut S.put $$ sinkSocket sock)
       (killThread) $ \ _ ->
       runWire wire' sock chan $ Connection
-        { connProtocol     = hsProtocol hs
+        { connRemoteAddr   = addr
+        , connProtocol     = hsProtocol hs
         , connCaps         = caps
         , connTopic        = hsInfoHash hs
         , connRemotePeerId = hsPeerId   hs'