2 files changed, 296 insertions, 78 deletions

diff --git a/src/Network/BitTorrent/Exchange/Session.hs b/src/Network/BitTorrent/Exchange/Session.hs
index 88554807..1537efe1 100644
--- a/src/Network/BitTorrent/Exchange/Session.hs
+++ b/src/Network/BitTorrent/Exchange/Session.hs
@@ -7,11 +7,16 @@ module Network.BitTorrent.Exchange.Session
        , newSession
        , closeSession
 
+         -- * Connections
        , Network.BitTorrent.Exchange.Session.insert
+       , Network.BitTorrent.Exchange.Session.attach
+       , Network.BitTorrent.Exchange.Session.delete
+       , Network.BitTorrent.Exchange.Session.deleteAll
        ) where
 
 import Control.Applicative
 import Control.Concurrent
+import Control.Concurrent.STM
 import Control.Exception hiding (Handler)
 import Control.Lens
 import Control.Monad.Logger
@@ -23,6 +28,7 @@ import Data.Conduit.List as CL (iterM)
 import Data.Maybe
 import Data.Map as M
 import Data.Monoid
+import Data.Set  as S
 import Data.Text as T
 import Data.Typeable
 import Text.PrettyPrint hiding ((<>))
@@ -70,11 +76,6 @@ data Cached a = Cached
 cache :: BEncode a => a -> Cached a
 cache s = Cached s (BE.encode s)
 
-data ConnectionEntry = ConnectionEntry
-  { initiatedBy :: !ChannelSide
-  , connection  :: !(Connection Session)
-  }
-
 data Session = Session
   { tpeerId      :: PeerId
   , infohash     :: InfoHash
@@ -82,12 +83,15 @@ data Session = Session
   , storage      :: Storage
   , status       :: MVar SessionStatus
   , unchoked     :: [PeerAddr IP]
-  , connections  :: MVar (Map (PeerAddr IP) ConnectionEntry)
+  , pendingConnections     :: TVar (Set (PeerAddr IP))
+  , establishedConnections :: TVar (Map (PeerAddr IP) (Connection Session))
   , broadcast    :: Chan Message
   , logger       :: LogFun
   , infodict     :: MVar (Cached InfoDict)
   }
 
+instance Ord IP
+
 -- | Logger function.
 type LogFun = Loc -> LogSource -> LogLevel -> LogStr -> IO ()
 
@@ -97,18 +101,21 @@ newSession :: LogFun
            -> InfoDict            -- ^ torrent info dictionary;
            -> IO Session          -- ^
 newSession logFun addr rootPath dict = do
-  connVar     <- newMVar M.empty
+  pconnVar    <- newTVarIO S.empty
+  econnVar    <- newTVarIO 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
-    , status      = statusVar
-    , storage     = store
-    , unchoked    = []
-    , connections = connVar
+    { tpeerId                = fromMaybe (error "newSession: impossible")
+                                         (peerId addr)
+    , infohash               = idInfoHash dict
+    , status                 = statusVar
+    , storage                = store
+    , unchoked               = []
+    , pendingConnections     = pconnVar
+    , establishedConnections = econnVar
     , broadcast   = chan
     , logger      = logFun
     }
@@ -137,31 +144,79 @@ logEvent   :: MonadLogger m => Text    -> m ()
 logEvent       = logInfoN
 
 {-----------------------------------------------------------------------
+--  Connection slots
+-----------------------------------------------------------------------}
+--- pending -> established -> closed
+---    |                        /|\
+---    \-------------------------|
+
+pendingConnection :: PeerAddr IP -> Session -> IO Bool
+pendingConnection addr Session {..} = atomically $ do
+  pSet <- readTVar pendingConnections
+  eSet <- readTVar establishedConnections
+  if (addr `S.member` pSet) || (addr `M.member` eSet)
+    then return False
+    else do
+      modifyTVar' pendingConnections (S.insert addr)
+      return True
+
+establishedConnection :: Connected Session ()
+establishedConnection = undefined --atomically $ do
+--  pSet <- readTVar pendingConnections
+--  eSet <- readTVar
+  undefined
+
+finishedConnection :: Connected Session ()
+finishedConnection = return ()
+
+-- | There are no state for this connection, remove it.
+closedConnection :: PeerAddr IP -> Session -> IO ()
+closedConnection addr Session {..} = atomically $ do
+  modifyTVar pendingConnections     $ S.delete addr
+  modifyTVar establishedConnections $ M.delete addr
+
+{-----------------------------------------------------------------------
 --  Connections
 -----------------------------------------------------------------------}
 -- TODO unmap storage on zero connections
 
-insert :: PeerAddr IP
-       -> {- Maybe Socket
-       -> -} Session -> IO ()
+mainWire :: Wire Session ()
+mainWire = do
+  lift establishedConnection
+  Session {..} <- asks connSession
+  lift $ resizeBitfield (totalPieces storage)
+  logEvent "Connection established"
+  iterM logMessage =$= exchange =$= iterM logMessage
+  lift finishedConnection
+
+getConnectionConfig :: Session -> IO (ConnectionConfig Session)
+getConnectionConfig s @ Session {..} = undefined --ConnectionConfig
+--      let caps  = def
+--      let ecaps = def
+--      let hs = Handshake def caps infohash tpeerId
+--      chan <- dupChan broadcast
+
+--  { cfgPrefs   = undefined
+--  , cfgSession = ConnectionSession undefined undefined s
+--  , cfgWire    = mainWire
+--  }
+
+insert :: PeerAddr IP -> Session -> IO ()
 insert addr ses @ Session {..} = void $ forkIO (action `finally` cleanup)
   where
+    action = do
+      pendingConnection addr ses
+      cfg <- getConnectionConfig ses
+      connectWire addr cfg
+
     cleanup = do
       runStatusUpdates status (SS.resetPending addr)
       -- TODO Metata.resetPending addr
+      closedConnection addr ses
 
-    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 <- ask
---      liftIO $ modifyMVar_ connections $ pure . M.insert addr conn
-        lift $ resizeBitfield (totalPieces storage)
-        logEvent "Connection established"
-        iterM logMessage =$= exchange =$= iterM logMessage
---      liftIO $ modifyMVar_ connections $ pure . M.delete addr
+-- TODO closePending on error
+attach :: PendingConnection -> Session -> IO ()
+attach = undefined
 
 delete :: PeerAddr IP -> Session -> IO ()
 delete = undefined
diff --git a/src/Network/BitTorrent/Exchange/Wire.hs b/src/Network/BitTorrent/Exchange/Wire.hs
index 27915724..d1bed146 100644
--- a/src/Network/BitTorrent/Exchange/Wire.hs
+++ b/src/Network/BitTorrent/Exchange/Wire.hs
@@ -7,6 +7,11 @@
 --   Stability   :  experimental
 --   Portability :  portable
 --
+--   Each peer wire connection is identified by triple @(topic,
+--   remote_addr, this_addr)@. This means that connections are the
+--   same if and only if their 'ConnectionId' are the same. Of course,
+--   you /must/ avoid duplicated connections.
+--
 --   This module control /integrity/ of data send and received.
 --
 {-# LANGUAGE DeriveDataTypeable #-}
@@ -21,6 +26,7 @@ module Network.BitTorrent.Exchange.Wire
 
          -- * Connection
        , Connection
+       , connInitiatedBy
 
          -- ** Identity
        , connRemoteAddr
@@ -44,9 +50,23 @@ module Network.BitTorrent.Exchange.Wire
        , connStats
 
          -- * Setup
-       , runWire
+       , ConnectionPrefs   (..)
+       , ConnectionSession (..)
+       , ConnectionConfig  (..)
+
+         -- ** Initiate
        , connectWire
+
+         -- ** Accept
+       , PendingConnection
+       , newPendingConnection
+       , pendingPeer
+       , pendingCaps
+       , pendingTopic
+       , closePending
        , acceptWire
+
+         -- ** Post setup actions
        , resizeBitfield
 
          -- * Messaging
@@ -143,7 +163,7 @@ data ProtocolError
 
     -- | /Remote/ peer replied with invalid 'hsInfoHash' which do not
     -- match with 'hsInfoHash' /this/ peer have sent. Can occur in
-    -- 'connectWire' only.
+    -- 'connectWire' or 'acceptWire' only.
   | UnexpectedTopic   InfoHash
 
     -- | Some trackers or DHT can return 'PeerId' of a peer. If a
@@ -470,9 +490,20 @@ data ConnectionState = ConnectionState {
 
 makeLenses ''ConnectionState
 
+instance Default ConnectionState where
+   def = ConnectionState
+     { _connExtCaps      = def
+     , _connRemoteEhs    = def
+     , _connStats        = def
+     , _connStatus       = def
+     , _connBitfield     = BF.haveNone 0
+     }
+
 -- | Connection keep various info about both peers.
 data Connection s = Connection
-  { connRemoteAddr   :: !(PeerAddr IP)
+  { connInitiatedBy  :: !ChannelSide
+
+  , connRemoteAddr   :: !(PeerAddr IP)
 
     -- | /Both/ peers handshaked with this protocol string. The only
     -- value is \"Bittorrent Protocol\" but this can be changed in
@@ -561,6 +592,29 @@ initiateHandshake sock hs = do
   sendHandshake sock hs
   recvHandshake sock
 
+data HandshakePair = HandshakePair
+  { handshakeSent :: !Handshake
+  , handshakeRecv :: !Handshake
+  } deriving (Show, Eq)
+
+validatePair :: HandshakePair -> PeerAddr IP -> IO ()
+validatePair (HandshakePair hs hs') addr = Prelude.mapM_ checkProp
+  [ (def            == hsProtocol hs', InvalidProtocol    $ hsProtocol hs')
+  , (hsProtocol hs  == hsProtocol hs', UnexpectedProtocol $ hsProtocol hs')
+  , (hsInfoHash hs  == hsInfoHash hs', UnexpectedTopic    $ hsInfoHash hs')
+  , (hsPeerId   hs' == fromMaybe (hsPeerId hs') (peerId addr)
+    , UnexpectedPeerId $ hsPeerId hs')
+  ]
+  where
+    checkProp (t, e) = unless t $ throwIO $ ProtocolError e
+
+-- | Connection state /right/ after handshaking.
+establishedStats :: HandshakePair -> ConnectionStats
+establishedStats HandshakePair {..} = ConnectionStats
+  { outcomingFlow = FlowStats 1 $ handshakeStats handshakeSent
+  , incomingFlow  = FlowStats 1 $ handshakeStats handshakeRecv
+  }
+
 {-----------------------------------------------------------------------
 --  Wire
 -----------------------------------------------------------------------}
@@ -670,53 +724,129 @@ rehandshake caps = undefined
 reconnect :: Wire s ()
 reconnect = undefined
 
--- | Initiate 'Wire' connection and handshake with a peer. This function will
--- also do the BEP10 extension protocol handshake if 'ExtExtended' is enabled on
--- both sides.
+data ConnectionId    = ConnectionId
+  { topic      :: !InfoHash
+  , remoteAddr :: !(PeerAddr IP)
+  , thisAddr   :: !(PeerAddr (Maybe IP)) -- ^ foreign address of this node.
+  }
+
+-- | /Preffered/ settings of wire. To get the real use 'ask'.
+data ConnectionPrefs = ConnectionPrefs
+  { prefOptions  :: !Options
+  , prefProtocol :: !ProtocolName
+  , prefCaps     :: !Caps
+  , prefExtCaps  :: !ExtendedCaps
+  }
+
+instance Default ConnectionPrefs where
+  def = ConnectionPrefs
+    { prefOptions  = def
+    , prefProtocol = def
+    , prefCaps     = def
+    , prefExtCaps  = def
+    }
+
+normalize :: ConnectionPrefs -> ConnectionPrefs
+normalize = undefined
+
+data ConnectionSession s = ConnectionSession
+  { sessionTopic      :: !(InfoHash)
+  , sessionPeerId     :: !(PeerId)
+  , metadataSize      :: !(Maybe Int)
+  , outputChan        :: !(Maybe (Chan Message))
+  , connectionSession :: !(s)
+  }
+
+data ConnectionConfig s = ConnectionConfig
+  { cfgPrefs   :: !(ConnectionPrefs)
+  , cfgSession :: !(ConnectionSession s)
+  , cfgWire    :: !(Wire s ())
+  }
+
+configHandshake :: ConnectionConfig s -> Handshake
+configHandshake ConnectionConfig {..} = Handshake
+  { hsProtocol = prefProtocol  cfgPrefs
+  , hsReserved = prefCaps      cfgPrefs
+  , hsInfoHash = sessionTopic  cfgSession
+  , hsPeerId   = sessionPeerId cfgSession
+  }
+
+{-----------------------------------------------------------------------
+--  Pending connections
+-----------------------------------------------------------------------}
+
+-- | Connection in half opened state. A normal usage scenario:
 --
--- This function can throw 'WireFailure' exception.
+--    * Opened using 'newPendingConnection', usually in the listener
+--    loop;
 --
-connectWire :: s -> Handshake -> PeerAddr IP -> ExtendedCaps -> Chan Message
-            -> Wire s () -> IO ()
-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
+--    * Closed using 'closePending' if 'pendingPeer' is banned,
+--    'pendingCaps' is prohibited or pendingTopic is unknown;
+--
+--    * Accepted using 'acceptWire' otherwise.
+--
+data PendingConnection = PendingConnection
+  { pendingSock  :: Socket
+  , pendingPeer  :: PeerAddr IP -- ^ 'peerId' is always non empty;
+  , pendingCaps  :: Caps        -- ^ advertised by the peer;
+  , pendingTopic :: InfoHash    -- ^ possible non-existent topic.
+  }
 
-    Prelude.mapM_ (\(t,e) -> unless t $ throwIO $ ProtocolError e) [
-      (def == hsProtocol hs'
-      , InvalidProtocol $ hsProtocol hs'),
-      (hsProtocol hs == hsProtocol hs'
-      , UnexpectedProtocol $ hsProtocol hs'),
-      (hsInfoHash hs == hsInfoHash hs'
-      , UnexpectedTopic $ hsInfoHash hs'),
-      (hsPeerId hs' == fromMaybe (hsPeerId hs') (peerId addr)
-      , UnexpectedPeerId $ hsPeerId hs')
-      ]
-
-    let caps = hsReserved hs <> hsReserved hs'
-        wire' = if ExtExtended `allowed` caps
-                then extendedHandshake extCaps >> wire
-                else wire
+-- | Reconstruct handshake sent by the remote peer.
+pendingHandshake :: PendingConnection -> Handshake
+pendingHandshake PendingConnection {..} = Handshake
+  { hsProtocol = def
+  , hsReserved = pendingCaps
+  , hsInfoHash = pendingTopic
+  , hsPeerId   = fromMaybe (error "pendingHandshake: impossible")
+                           (peerId pendingPeer)
+  }
 
-    cstate <- newIORef $ ConnectionState {
-        _connExtCaps      = def
-      , _connRemoteEhs    = def
-      , _connStats        = ConnectionStats {
-                             outcomingFlow = FlowStats 1 $ handshakeStats hs
-                           , incomingFlow  = FlowStats 1 $ handshakeStats hs'
-                           }
-      , _connStatus       = def
-      , _connBitfield     = BF.haveNone 0
-      }
+-- |
+--
+--   This function can throw 'WireFailure' exception.
+--
+newPendingConnection :: Socket -> PeerAddr IP -> IO PendingConnection
+newPendingConnection sock addr = do
+  Handshake {..} <- recvHandshake sock
+  unless (hsProtocol == def) $ do
+    throwIO $ ProtocolError $ InvalidProtocol hsProtocol
+  return PendingConnection
+    { pendingSock  = sock
+    , pendingPeer  = addr { peerId = Just hsPeerId }
+    , pendingCaps  = hsReserved
+    , pendingTopic = hsInfoHash
+    }
 
-    -- TODO make KA interval configurable
-    let kaInterval = defaultKeepAliveInterval
-    bracket
-      (forkIO $ sourceChan kaInterval chan $= conduitPut S.put $$ sinkSocket sock)
-      (killThread) $ \ _ ->
-      runWire wire' sock chan $ Connection
-        { connRemoteAddr   = addr
+-- | Release all resources associated with the given connection. Note
+-- that you /must not/ 'closePending' if you 'acceptWire'.
+closePending :: PendingConnection -> IO ()
+closePending PendingConnection {..} = do
+  close pendingSock
+
+{-----------------------------------------------------------------------
+--  Connection setup
+-----------------------------------------------------------------------}
+
+chanToSock :: Int -> Chan Message -> Socket -> IO ()
+chanToSock ka chan sock =
+  sourceChan ka chan $= conduitPut S.put $$ sinkSocket sock
+
+afterHandshaking :: ChannelSide -> PeerAddr IP -> Socket -> HandshakePair
+                 -> ConnectionConfig s -> IO ()
+afterHandshaking initiator addr sock
+  hpair @ (HandshakePair hs hs')
+          (ConnectionConfig
+           { cfgPrefs   = ConnectionPrefs   {..}
+           , cfgSession = ConnectionSession {..}
+           , cfgWire    = wire
+           }) = do
+    let caps  = hsReserved hs <> hsReserved hs'
+    cstate <- newIORef def { _connStats = establishedStats hpair }
+    chan   <- maybe newChan return outputChan
+    let conn  = Connection {
+          connInitiatedBy  = initiator
+        , connRemoteAddr   = addr
         , connProtocol     = hsProtocol hs
         , connCaps         = caps
         , connTopic        = hsInfoHash hs
@@ -724,20 +854,53 @@ connectWire session hs addr extCaps chan wire = do
         , connThisPeerId   = hsPeerId   hs
         , connOptions      = def
         , connState        = cstate
-        , connSession      = session
+        , connSession      = connectionSession
         , connChan         = chan
         }
 
+    -- TODO make KA interval configurable
+    let kaInterval = defaultKeepAliveInterval
+        wire' = if ExtExtended `allowed` caps
+                then extendedHandshake prefExtCaps >> wire
+                else wire
+
+    bracket (forkIO (chanToSock kaInterval chan sock))
+            (killThread)
+            (\ _ -> runWire wire' sock chan conn)
+
+-- | Initiate 'Wire' connection and handshake with a peer. This function will
+-- also do the BEP10 extension protocol handshake if 'ExtExtended' is enabled on
+-- both sides.
+--
+-- This function can throw 'WireFailure' exception.
+--
+connectWire :: PeerAddr IP -> ConnectionConfig s -> IO ()
+connectWire addr cfg = do
+  let catchRefusal m = try m >>= either (throwIO . ConnectionRefused) return
+  bracket (catchRefusal (peerSocket Stream addr)) close $ \ sock -> do
+    let hs = configHandshake cfg
+    hs' <- initiateHandshake sock hs
+    let hpair = HandshakePair hs hs'
+    validatePair hpair addr
+    afterHandshaking ThisPeer addr sock hpair cfg
+
 -- | Accept 'Wire' connection using already 'Network.Socket.accept'ed
 --   socket. For peer listener loop the 'acceptSafe' should be
 --   prefered against 'accept'. The socket will be closed at exit.
 --
 --   This function can throw 'WireFailure' exception.
 --
-acceptWire :: Socket -> PeerAddr IP -> Wire s () -> IO ()
-acceptWire sock peerAddr wire = do
-  bracket (return sock) close $ \ _ -> do
-    error "acceptWire: not implemented"
+acceptWire :: PendingConnection -> ConnectionConfig s -> IO ()
+acceptWire pc @ PendingConnection {..} cfg = do
+  bracket (return pendingSock) close $ \ _ -> do
+    unless (sessionTopic (cfgSession cfg) == pendingTopic) $ do
+      throwIO (ProtocolError (UnexpectedTopic pendingTopic))
+
+    let hs = configHandshake cfg
+    sendHandshake pendingSock hs
+    let hpair = HandshakePair hs (pendingHandshake pc)
+
+    afterHandshaking RemotePeer pendingPeer pendingSock hpair cfg
 
 -- | Used when size of bitfield becomes known.
 resizeBitfield :: Int -> Connected s ()