3 files changed, 257 insertions, 134 deletions

diff --git a/src/Network/BitTorrent/Exchange.hs b/src/Network/BitTorrent/Exchange.hs
index dda7d304..65ec0eb7 100644
--- a/src/Network/BitTorrent/Exchange.hs
+++ b/src/Network/BitTorrent/Exchange.hs
@@ -5,12 +5,19 @@
 --   Stability   :  experimental
 --   Portability :  portable
 --
-{-# LANGUAGE DoAndIfThenElse #-}
+{-# LANGUAGE OverloadedStrings          #-}
+{-# LANGUAGE DoAndIfThenElse            #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
-{-# LANGUAGE MultiParamTypeClasses #-}
-{-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE MultiParamTypeClasses      #-}
+{-# LANGUAGE RecordWildCards            #-}
 module Network.BitTorrent.Exchange
-       ( P2P, withPeer
+       ( -- * Block
+         Block(..), BlockIx(..)
+
+         -- * Event
+       , Event(..)
+
+       , P2P, withPeer
        , awaitEvent, signalEvent
        ) where
 
@@ -18,6 +25,7 @@ import Control.Applicative
 import Control.Concurrent
 import Control.Concurrent.STM
 import Control.Exception
+import Control.Lens
 import Control.Monad.Reader
 import Control.Monad.State
 
@@ -30,6 +38,7 @@ import Data.Conduit as C
 import Data.Conduit.Cereal
 import Data.Conduit.Network
 import Data.Serialize as S
+import Text.PrettyPrint as PP hiding (($$))
 
 import Network
 
@@ -43,8 +52,8 @@ import Data.Torrent
 
 
 data Event = Available Bitfield
-           | Want
-           | Block
+           | Want      BlockIx
+           | Fragment  Block
              deriving Show
 
 {-----------------------------------------------------------------------
@@ -53,88 +62,132 @@ data Event = Available Bitfield
 
 type PeerWire = ConduitM Message Message IO
 
-waitMessage :: PeerSession -> PeerWire Message
-waitMessage se = do
-  mmsg <- await
-  case mmsg of
-    Nothing -> waitMessage se
-    Just msg -> do
-      liftIO $ updateIncoming se
-      liftIO $ print msg
-      return msg
-
-signalMessage :: PeerSession -> Message -> PeerWire ()
-signalMessage se msg = do
+runConduit :: Socket -> PeerWire () -> IO ()
+runConduit sock p2p =
+  sourceSocket sock   $=
+    conduitGet S.get  $=
+      forever p2p     $=
+    conduitPut S.put  $$
+  sinkSocket sock
+
+waitMessage :: P2P Message
+waitMessage = P2P (ReaderT go)
+  where
+    go se = do
+      mmsg <- await
+      case mmsg of
+        Nothing -> go se
+        Just msg -> do
+          liftIO $ updateIncoming se
+          liftIO $ print msg
+          return msg
+
+signalMessage :: Message -> P2P ()
+signalMessage msg = P2P $ ReaderT $ \se -> do
   C.yield msg
   liftIO $ updateOutcoming se
 
 
-getPieceCount :: PeerSession -> IO PieceCount
-getPieceCount = undefined
+peerWant :: P2P Bitfield
+peerWant   = BF.difference <$> getPeerBF    <*> use bitfield
+
+clientWant :: P2P Bitfield
+clientWant = BF.difference <$> use bitfield <*> getPeerBF
 
-canOffer :: PeerSession -> PeerWire Bitfield
-canOffer PeerSession {..} = liftIO $ do
-  pbf <- readIORef $ peerBitfield
-  cbf <- readIORef $ clientBitfield $ swarmSession
-  return $ BF.difference pbf cbf
+peerOffer :: P2P Bitfield
+peerOffer = do
+  sessionStatus <- use status
+  if canDownload sessionStatus then clientWant else emptyBF
 
-revise :: PeerSession -> PeerWire ()
-revise se @ PeerSession {..} = do
-  isInteresting <- (not . BF.null) <$> canOffer se
-  SessionStatus {..} <- liftIO $ readIORef peerSessionStatus
+clientOffer :: P2P Bitfield
+clientOffer = do
+  sessionStatus <- use status
+  if canUpload sessionStatus then peerWant else emptyBF
 
-  when (isInteresting /= _interested seClientStatus) $
-    signalMessage se $ if isInteresting then Interested else NotInterested
+revise :: P2P ()
+revise = do
+  peerInteresting  <- (not . BF.null) <$> clientWant
+  clientInterested <- use (status.clientStatus.interested)
 
+  when (clientInterested /= peerInteresting) $
+    signalMessage $ if peerInteresting then Interested else NotInterested
 
-nextEvent :: PeerSession -> PeerWire Event
-nextEvent se @ PeerSession {..} = waitMessage se >>= go
+requireExtension :: Extension -> P2P ()
+requireExtension required = do
+  enabled <- asks enabledExtensions
+  unless (required `elem` enabled) $
+    sessionError $ ppExtension required <+> "not enabled"
+
+peerHave :: P2P Event
+peerHave = undefined
+
+--    haveMessage bf = do
+--      cbf <- undefined -- liftIO $ readIORef $ clientBitfield swarmSession
+--      if undefined -- ix `member` bf
+--        then nextEvent se
+--        else undefined  -- return $ Available diff
+
+
+-- |
+--   properties:
+--
+--   forall (Fragment block). isPiece block == True
+--
+awaitEvent :: P2P Event
+awaitEvent = waitMessage >>= go
   where
-    go KeepAlive = nextEvent se
+    go KeepAlive = awaitEvent
     go Choke     = do
-      SessionStatus {..} <- liftIO $ readIORef peerSessionStatus
-      if _choking  sePeerStatus
-        then nextEvent se
-        else undefined
+      status.peerStatus.choking .= True
+      awaitEvent
 
     go Unchoke   = do
-      SessionStatus {..} <- liftIO $ readIORef peerSessionStatus
-      if not (_choking sePeerStatus)
-        then nextEvent se
-        else if undefined
-             then undefined
-             else undefined
-    --return $ Available BF.difference
-
-    go Interested = return undefined
-    go NotInterested = return undefined
-
-    go (Have ix) = do
-      pc <- liftIO $ getPieceCount se
-      haveMessage $ have ix (haveNone pc) -- TODO singleton
-
-    go (Bitfield bf) = undefined
+      status.clientStatus.choking .= False
+      awaitEvent
+
+    go Interested    = do
+      status.peerStatus.interested .= True
+      awaitEvent
+
+    go NotInterested = do
+      status.peerStatus.interested .= False
+      awaitEvent
+
+--    go (Have ix)      = peerHave =<< singletonBF ix
+--    go (Bitfield bf)  = peerHave =<< adjustBF bf
     go (Request  bix) = do
-      undefined
+      bf <- use bitfield
+      if ixPiece bix `BF.member` bf
+         then return (Want bix)
+         else do
+           signalMessage (RejectRequest bix)
+           awaitEvent
+
+    go (Piece    blk) = undefined
 
-    go msg @ (Piece    blk) = undefined
+{-
     go msg @ (Port     _)
       = checkExtension msg ExtDHT $ do
       undefined
 
-    go msg @ HaveAll
-      = checkExtension msg ExtFast $ do
-          pc <- liftIO $ getPieceCount se
-          haveMessage (haveAll pc)
-
-    go msg @ HaveNone
-      = checkExtension msg ExtFast $ do
-        pc <- liftIO $ getPieceCount se
-        haveMessage (haveNone pc)
-
-    go msg @ (SuggestPiece ix)
-          = checkExtension msg ExtFast $ do
-            undefined
+    go HaveAll = do
+      requireExtension ExtFast
+      bitfield <~ fullBF
+      revise
+      awaitEvent
+
+    go HaveNone = do
+      requireExtension ExtFast
+      bitfield <~ emptyBF
+      revise
+      awaitEvent
+
+    go (SuggestPiece ix) = do
+      requireExtension ExtFast
+      bf <- use bitfield
+      if ix `BF.notMember` bf
+        then Available <$> singletonBF ix
+        else awaitEvent
 
     go msg @ (RejectRequest ix)
           = checkExtension msg ExtFast $ do
@@ -143,20 +196,16 @@ nextEvent se @ PeerSession {..} = waitMessage se >>= go
     go msg @ (AllowedFast pix)
           = checkExtension msg ExtFast $ do
             undefined
+-}
 
-    haveMessage bf = do
-      cbf <- liftIO $ readIORef $ clientBitfield swarmSession
-      if undefined -- ix `member` bf
-        then nextEvent se
-        else undefined  -- return $ Available diff
+signalEvent  :: Event -> P2P ()
+signalEvent (Available bf) = undefined
+signalEvent _ = undefined
 
-    checkExtension msg requredExtension action
-      | requredExtension `elem` enabledExtensions = action
-      | otherwise = liftIO $ throwIO $ userError errorMsg
-      where
-        errorMsg = show (ppExtension requredExtension)
-                   ++  "not enabled, but peer sent"
-                   ++ show (ppMessage msg)
+--flushBroadcast :: P2P ()
+--flushBroadcast = nextBroadcast >>= maybe (return ()) go
+--  where
+--    go = undefined
 
 {-----------------------------------------------------------------------
     P2P monad
@@ -164,29 +213,9 @@ nextEvent se @ PeerSession {..} = waitMessage se >>= go
 
 newtype P2P a = P2P {
     runP2P :: ReaderT PeerSession PeerWire a
-  } deriving (Monad, MonadReader PeerSession, MonadIO)
-
-instance MonadState SessionStatus P2P where
-  get   = asks peerSessionStatus >>= liftIO . readIORef
-  put x = asks peerSessionStatus >>= liftIO . (`writeIORef` x)
-
-
-runConduit :: Socket -> Conduit Message IO Message -> IO ()
-runConduit sock p2p =
-  sourceSocket sock   $=
-    conduitGet S.get  $=
-      forever p2p     $=
-    conduitPut S.put  $$
-  sinkSocket sock
+  } deriving (Functor, Applicative, Monad, MonadReader PeerSession, MonadIO)
 
 withPeer :: SwarmSession -> PeerAddr -> P2P () -> IO ()
 withPeer se addr p2p =
   withPeerSession se addr $ \(sock, pses) -> do
     runConduit sock (runReaderT (runP2P p2p) pses)
-
-
-awaitEvent :: P2P Event
-awaitEvent = P2P (ReaderT nextEvent)
-
-signalEvent  :: Event -> P2P ()
-signalEvent = undefined
diff --git a/src/Network/BitTorrent/Exchange/Protocol.hs b/src/Network/BitTorrent/Exchange/Protocol.hs
index 46e25fa3..718e339d 100644
--- a/src/Network/BitTorrent/Exchange/Protocol.hs
+++ b/src/Network/BitTorrent/Exchange/Protocol.hs
@@ -26,7 +26,8 @@
 --   <https://wiki.theory.org/BitTorrentSpecification#Peer_wire_protocol_.28TCP.29>
 --
 {-# LANGUAGE OverloadedStrings #-}
-{-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE RecordWildCards   #-}
+{-# LANGUAGE TemplateHaskell   #-}
 module Network.BitTorrent.Exchange.Protocol
        ( -- * Inital handshake
          Handshake(..), ppHandshake
@@ -50,18 +51,22 @@ module Network.BitTorrent.Exchange.Protocol
        , Message(..)
        , ppMessage
 
-         -- * Exchange control
+         -- * control
        , PeerStatus(..)
+       , choking, interested
+
        , SessionStatus(..)
---       , canUpload, canDownload
+       , clientStatus, peerStatus
+       , canUpload, canDownload
 
-       -- ** Defaults
+         -- ** Defaults
        , defaultUnchokeSlots
        ) where
 
 import Control.Applicative
-import Control.Monad
 import Control.Exception
+import Control.Monad
+import Control.Lens
 import           Data.ByteString (ByteString)
 import qualified Data.ByteString as B
 import qualified Data.ByteString.Char8 as BC
@@ -427,29 +432,31 @@ data PeerStatus = PeerStatus {
   , _interested :: Bool
   }
 
+$(makeLenses ''PeerStatus)
+
 instance Default PeerStatus where
   def = PeerStatus True False
 
 -- |
 data SessionStatus = SessionStatus {
-    seClientStatus :: PeerStatus
-  , sePeerStatus   :: PeerStatus
+    _clientStatus :: PeerStatus
+  , _peerStatus   :: PeerStatus
   }
 
+$(makeLenses ''SessionStatus)
+
 instance Default SessionStatus where
   def = SessionStatus def def
 
-{-
--- | Can the /client/ to upload to the /peer/?
+-- | Can the /client/ transfer to the /peer/?
 canUpload :: SessionStatus -> Bool
 canUpload SessionStatus {..}
-  = psInterested sePeerStatus && not (psChoking seClientStatus)
+  = _interested _peerStatus && not (_choking _clientStatus)
 
--- | Can the /client/ download from the /peer/?
+-- | Can the /client/ transfer from the /peer/?
 canDownload :: SessionStatus -> Bool
 canDownload SessionStatus {..}
-  = psInterested seClientStatus && not (psChoking sePeerStatus)
--}
+  = _interested _clientStatus && not (_choking _peerStatus)
 
 -- | Indicates how many peers are allowed to download from the client
 -- by default.
diff --git a/src/Network/BitTorrent/Internal.hs b/src/Network/BitTorrent/Internal.hs
index 38087f0d..2f538652 100644
--- a/src/Network/BitTorrent/Internal.hs
+++ b/src/Network/BitTorrent/Internal.hs
@@ -10,12 +10,33 @@
 --   Network.BitTorrent.Exchange and modules. To hide some internals
 --   of this module we detach it from Exchange.
 --
-{-# LANGUAGE RecordWildCards #-}
+--   Note: expose only static data in data field lists, all dynamic
+--   data should be modified through standalone functions.
+--
+{-# LANGUAGE OverloadedStrings     #-}
+{-# LANGUAGE RecordWildCards       #-}
+{-# LANGUAGE TemplateHaskell       #-}
+{-# LANGUAGE FlexibleInstances     #-}
+{-# LANGUAGE FlexibleContexts      #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE UndecidableInstances  #-}
 module Network.BitTorrent.Internal
        ( Progress(..), startProgress
-       , ClientSession(..), newClient
-       , SwarmSession(..), newLeacher, newSeeder
-       , PeerSession(..), withPeerSession
+
+       , ClientSession (clientPeerID, allowedExtensions)
+       , newClient, getCurrentProgress
+
+       , SwarmSession(SwarmSession, torrentMeta, clientSession)
+       , newLeacher, newSeeder
+
+       , PeerSession(PeerSession, connectedPeerAddr
+                    , swarmSession, enabledExtensions
+                    )
+       , SessionState
+       , bitfield, status
+       , emptyBF, fullBF, singletonBF
+       , getPieceCount, getPeerBF
+       , sessionError, withPeerSession
 
          -- * Timeouts
        , updateIncoming, updateOutcoming
@@ -24,6 +45,9 @@ module Network.BitTorrent.Internal
 import Control.Applicative
 import Control.Concurrent
 import Control.Concurrent.STM
+import Control.Lens
+import Control.Monad.State
+import Control.Monad.Reader
 import Control.Exception
 
 import Data.IORef
@@ -32,10 +56,8 @@ import Data.Function
 import Data.Ord
 import Data.Set as S
 
-import Data.Conduit
-import Data.Conduit.Cereal
-import Data.Conduit.Network
-import Data.Serialize
+import Data.Serialize hiding (get)
+import Text.PrettyPrint
 
 import Network
 import Network.Socket
@@ -70,13 +92,20 @@ startProgress = Progress 0 0
     Client session
 -----------------------------------------------------------------------}
 
--- | In one application you could have many clients.
+-- | In one application we could have many clients with difference
+-- ID's and enabled extensions.
 data ClientSession = ClientSession {
-    clientPeerID      :: PeerID      -- ^
-  , allowedExtensions :: [Extension] -- ^
+    -- | Our peer ID used in handshaked and discovery mechanism.
+    clientPeerID      :: PeerID
+
+    -- | Extensions we should try to use. Hovewer some particular peer
+    -- might not support some extension, so we keep enableExtension in
+    -- 'PeerSession'.
+  , allowedExtensions :: [Extension]
+
   , swarmSessions     :: TVar (Set SwarmSession)
   , eventManager      :: EventManager
-  , currentProgress   :: IORef Progress
+  , currentProgress   :: TVar  Progress
   }
 
 instance Eq ClientSession where
@@ -85,6 +114,9 @@ instance Eq ClientSession where
 instance Ord ClientSession where
   compare = comparing clientPeerID
 
+getCurrentProgress :: MonadIO m => ClientSession -> m Progress
+getCurrentProgress = liftIO . readTVarIO . currentProgress
+
 newClient :: [Extension] -> IO ClientSession
 newClient exts = do
   mgr <- Ev.new
@@ -95,7 +127,7 @@ newClient exts = do
     <*> pure exts
     <*> newTVarIO S.empty
     <*> pure mgr
-    <*> newIORef (startProgress 0)
+    <*> newTVarIO (startProgress 0)
 
 {-----------------------------------------------------------------------
     Swarm session
@@ -106,7 +138,9 @@ newClient exts = do
 data SwarmSession = SwarmSession {
     torrentMeta       :: Torrent
   , clientSession     :: ClientSession
-  , clientBitfield    :: IORef Bitfield
+
+    -- | Modify this carefully updating global progress.
+  , clientBitfield    :: TVar  Bitfield
   , connectedPeers    :: TVar (Set PeerSession)
   }
 
@@ -120,7 +154,7 @@ newSwarmSession :: Bitfield -> ClientSession -> Torrent -> IO SwarmSession
 newSwarmSession bf cs @ ClientSession {..} t @ Torrent {..}
   = SwarmSession <$> pure t
                  <*> pure cs
-                 <*> newIORef bf
+                 <*> newTVarIO bf
                  <*> newTVarIO S.empty
 
 newSeeder :: ClientSession -> Torrent -> IO SwarmSession
@@ -134,13 +168,29 @@ newLeacher cs t @ Torrent {..}
 isLeacher :: SwarmSession -> IO Bool
 isLeacher = undefined
 
+getClientBitfield :: MonadIO m => SwarmSession -> m Bitfield
+getClientBitfield = liftIO . readTVarIO .  clientBitfield
+
+{-
+haveDone :: MonadIO m => PieceIx -> SwarmSession -> m ()
+haveDone ix =
+  liftIO $ atomically $ do
+    bf <- readTVar clientBitfield
+    writeTVar (have ix bf)
+    currentProgress
+-}
 {-----------------------------------------------------------------------
     Peer session
 -----------------------------------------------------------------------}
 
 data PeerSession = PeerSession {
+    -- | Used as unique 'PeerSession' identifier within one
+    -- 'SwarmSession'.
     connectedPeerAddr :: PeerAddr
+
   , swarmSession      :: SwarmSession
+
+    -- | Extensions such that both peer and client support.
   , enabledExtensions :: [Extension]
 
     -- | To dissconnect from died peers appropriately we should check
@@ -163,16 +213,31 @@ data PeerSession = PeerSession {
   , outcomingTimeout   :: TimeoutKey
 
   , broadcastMessages :: Chan   [Message]
-  , peerBitfield      :: IORef  Bitfield
-  , peerSessionStatus :: IORef  SessionStatus
+  , sessionState      :: IORef   SessionState
+  }
+
+data SessionState = SessionState {
+    _bitfield :: Bitfield
+  , _status   :: SessionStatus
   }
 
+$(makeLenses ''SessionState)
+
 instance Eq PeerSession where
   (==) = (==) `on` connectedPeerAddr
 
 instance Ord PeerSession where
   compare = comparing connectedPeerAddr
 
+instance (MonadIO m, MonadReader PeerSession m)
+      => MonadState SessionState m where
+  get   = asks sessionState >>= liftIO . readIORef
+  put s = asks sessionState >>= \ref -> liftIO $ writeIORef ref s
+
+sessionError :: MonadIO m => Doc -> m ()
+sessionError msg
+  = liftIO $ throwIO $ userError $ render $ msg <+> "in session"
+
 -- TODO check if it connected yet peer
 withPeerSession :: SwarmSession -> PeerAddr
                 -> ((Socket, PeerSession) -> IO a)
@@ -196,12 +261,34 @@ withPeerSession ss @ SwarmSession {..} addr
          <*> registerTimeout (eventManager clientSession)
                 maxOutcomingTime (sendKA sock)
          <*> newChan
-         <*> pure clientBitfield
-         <*> newIORef def
+         <*> do {
+           ; tc <- totalCount <$> readTVarIO clientBitfield
+           ; newIORef (SessionState (haveNone tc) def)
+           }
       return (sock, ps)
 
     closeSession = close . fst
 
+getPieceCount :: (MonadReader PeerSession m) => m PieceCount
+getPieceCount = asks (pieceCount . tInfo . torrentMeta . swarmSession)
+
+emptyBF :: (MonadReader PeerSession m) => m Bitfield
+emptyBF = liftM haveNone getPieceCount
+
+fullBF ::  (MonadReader PeerSession m) => m Bitfield
+fullBF = liftM haveAll getPieceCount
+
+singletonBF :: (MonadReader PeerSession m) => PieceIx -> m Bitfield
+singletonBF ix = liftM (BF.singleton ix) getPieceCount
+
+getPeerBF :: (MonadIO m, MonadReader PeerSession m) => m Bitfield
+getPeerBF = asks swarmSession >>= liftIO . readTVarIO . clientBitfield
+
+--data Signal =
+--nextBroadcast :: P2P (Maybe Signal)
+--nextBroadcast =
+
+
 {-----------------------------------------------------------------------
     Timeouts
 -----------------------------------------------------------------------}