1 files changed, 124 insertions, 30 deletions

diff --git a/src/Network/BitTorrent/Exchange/Session.hs b/src/Network/BitTorrent/Exchange/Session.hs
index 1e72ba96..f10f601e 100644
--- a/src/Network/BitTorrent/Exchange/Session.hs
+++ b/src/Network/BitTorrent/Exchange/Session.hs
@@ -12,7 +12,7 @@ module Network.BitTorrent.Exchange.Session
 
 import Control.Applicative
 import Control.Concurrent
-import Control.Exception
+import Control.Exception hiding (Handler)
 import Control.Lens
 import Control.Monad.Logger
 import Control.Monad.Reader
@@ -20,6 +20,7 @@ import Control.Monad.State
 import Data.ByteString as BS
 import Data.ByteString.Lazy as BL
 import Data.Conduit
+import Data.Conduit.List as CL (iterM)
 import Data.Function
 import Data.IORef
 import Data.List as L
@@ -34,16 +35,18 @@ import Text.PrettyPrint hiding ((<>))
 import Text.PrettyPrint.Class
 import System.Log.FastLogger (LogStr, ToLogStr (..))
 
+import Data.BEncode as BE
 import Data.Torrent (InfoDict (..))
 import Data.Torrent.Bitfield as BF
 import Data.Torrent.InfoHash
-import Data.Torrent.Piece (pieceData, piPieceLength)
+import Data.Torrent.Piece
 import qualified Data.Torrent.Piece as Torrent (Piece (Piece))
 import Network.BitTorrent.Core
 import Network.BitTorrent.Exchange.Assembler
-import Network.BitTorrent.Exchange.Block as Block
-import Network.BitTorrent.Exchange.Message
-import Network.BitTorrent.Exchange.Session.Status as SS
+import Network.BitTorrent.Exchange.Block   as Block
+import Network.BitTorrent.Exchange.Message as Message
+import Network.BitTorrent.Exchange.Session.Metadata as Metadata
+import Network.BitTorrent.Exchange.Session.Status   as SS
 import Network.BitTorrent.Exchange.Status
 import Network.BitTorrent.Exchange.Wire
 import System.Torrent.Storage
@@ -66,6 +69,14 @@ packException f m = try m >>= either (throwIO . f) return
 --  Session
 -----------------------------------------------------------------------}
 
+data Cached a = Cached
+  { cachedValue :: !a
+  , cachedData  :: BL.ByteString -- keep lazy
+  }
+
+cache :: BEncode a => a -> Cached a
+cache s = Cached s (BE.encode s)
+
 data ConnectionEntry = ConnectionEntry
   { initiatedBy :: !ChannelSide
   , connection  :: !(Connection Session)
@@ -74,12 +85,14 @@ data ConnectionEntry = ConnectionEntry
 data Session = Session
   { tpeerId      :: PeerId
   , infohash     :: InfoHash
+  , metadata     :: MVar Metadata.Status
   , storage      :: Storage
   , status       :: MVar SessionStatus
   , unchoked     :: [PeerAddr IP]
   , connections  :: MVar (Map (PeerAddr IP) ConnectionEntry)
   , broadcast    :: Chan Message
   , logger       :: LogFun
+  , infodict     :: MVar (Cached InfoDict)
   }
 
 -- | Logger function.
@@ -108,7 +121,13 @@ newSession logFun addr rootPath dict = do
     }
 
 closeSession :: Session -> IO ()
-closeSession = undefined
+closeSession ses = do
+  deleteAll ses
+  undefined
+
+{-----------------------------------------------------------------------
+--  Logging
+-----------------------------------------------------------------------}
 
 instance MonadLogger (Connected Session) where
   monadLoggerLog loc src lvl msg = do
@@ -118,11 +137,11 @@ instance MonadLogger (Connected Session) where
     let addrSrc = src <> " @ " <> T.pack (render (pretty addr))
     liftIO $ logger ses loc addrSrc lvl (toLogStr msg)
 
-logMessage :: Message -> Wire Session ()
+logMessage :: MonadLogger m => Message -> m ()
 logMessage msg = logDebugN $ T.pack (render (pretty msg))
 
-logEvent :: Text -> Wire Session ()
-logEvent = logInfoN
+logEvent   :: MonadLogger m => Text    -> m ()
+logEvent       = logInfoN
 
 {-----------------------------------------------------------------------
 --  Connections
@@ -132,11 +151,12 @@ logEvent = logInfoN
 insert :: PeerAddr IP
        -> {- Maybe Socket
        -> -} Session -> IO ()
-insert addr ses @ Session {..} = do
-    forkIO $ do
-      action `finally` runStatusUpdates status (resetPending addr)
-    return ()
+insert addr ses @ Session {..} = void $ forkIO (action `finally` cleanup)
   where
+    cleanup = do
+      runStatusUpdates status (SS.resetPending addr)
+      -- TODO Metata.resetPending addr
+
     action = do
       let caps  = def
       let ecaps = def
@@ -147,7 +167,7 @@ insert addr ses @ Session {..} = do
 --      liftIO $ modifyMVar_ connections $ pure . M.insert addr conn
         lift $ resizeBitfield (totalPieces storage)
         logEvent "Connection established"
-        exchange
+        iterM logMessage =$= exchange =$= iterM logMessage
 --      liftIO $ modifyMVar_ connections $ pure . M.delete addr
 
 delete :: PeerAddr IP -> Session -> IO ()
@@ -160,11 +180,26 @@ deleteAll = undefined
 --  Helpers
 -----------------------------------------------------------------------}
 
+waitMVar :: MVar a -> IO ()
+waitMVar m = withMVar m (const (return ()))
+
+-- This function appear in new GHC "out of box". (moreover it is atomic)
+tryReadMVar :: MVar a -> IO (Maybe a)
+tryReadMVar m = do
+  ma <- tryTakeMVar m
+  maybe (return ()) (putMVar m) ma
+  return ma
+
 withStatusUpdates :: StatusUpdates a -> Wire Session a
 withStatusUpdates m = do
   Session {..} <- asks connSession
   liftIO $ runStatusUpdates status m
 
+withMetadataUpdates :: Updates a -> Connected Session a
+withMetadataUpdates m = do
+  Session {..} <- asks connSession
+  liftIO $ runUpdates metadata m
+
 getThisBitfield :: Wire Session Bitfield
 getThisBitfield = do
   ses <- asks connSession
@@ -179,6 +214,16 @@ readBlock bix @ BlockIx {..} s = do
     then return  $ Block ixPiece ixOffset chunk
     else throwIO $ InvalidRequest bix (InvalidSize ixLength)
 
+-- |
+tryReadMetadataBlock :: PieceIx
+   -> Connected Session (Maybe (Torrent.Piece BS.ByteString, Int))
+tryReadMetadataBlock pix = do
+  Session {..} <- asks connSession
+  mcached      <- liftIO (tryReadMVar infodict)
+  case mcached of
+    Nothing            -> undefined
+    Just (Cached {..}) -> undefined
+
 sendBroadcast :: PeerMessage msg => msg -> Wire Session ()
 sendBroadcast msg = do
   Session {..} <- asks connSession
@@ -208,9 +253,7 @@ tryFillRequestQueue = do
 interesting :: Wire Session ()
 interesting = do
   addr <- asks connRemoteAddr
-  logMessage  (Status (Interested True))
   sendMessage (Interested True)
-  logMessage  (Status (Choking    False))
   sendMessage (Choking    False)
   tryFillRequestQueue
 
@@ -218,17 +261,19 @@ interesting = do
 --  Incoming message handling
 -----------------------------------------------------------------------}
 
-handleStatus :: StatusUpdate -> Wire Session ()
+type Handler msg = msg -> Wire Session ()
+
+handleStatus :: Handler StatusUpdate
 handleStatus s = do
   connStatus %= over remoteStatus (updateStatus s)
   case s of
     Interested _     -> return ()
     Choking    True  -> do
       addr <- asks connRemoteAddr
-      withStatusUpdates (resetPending addr)
+      withStatusUpdates (SS.resetPending addr)
     Choking    False -> tryFillRequestQueue
 
-handleAvailable :: Available -> Wire Session ()
+handleAvailable :: Handler Available
 handleAvailable msg = do
   connBitfield %= case msg of
     Have     ix -> BF.insert ix
@@ -243,18 +288,18 @@ handleAvailable msg = do
       | bf `BF.isSubsetOf` thisBf -> return ()
       |     otherwise             -> interesting
 
-handleTransfer :: Transfer -> Wire Session ()
+handleTransfer :: Handler Transfer
 handleTransfer (Request bix) = do
   Session {..} <- asks connSession
   bitfield <- getThisBitfield
   upload   <- canUpload <$> use connStatus
   when (upload && ixPiece bix `BF.member` bitfield) $ do
     blk <- liftIO $ readBlock bix storage
-    sendMessage (Piece blk)
+    sendMessage (Message.Piece blk)
 
-handleTransfer (Piece   blk) = do
+handleTransfer (Message.Piece   blk) = do
   Session {..} <- asks connSession
-  isSuccess <- withStatusUpdates (pushBlock blk storage)
+  isSuccess <- withStatusUpdates (SS.pushBlock blk storage)
   case isSuccess of
     Nothing -> liftIO $ throwIO $ userError "block is not requested"
     Just isCompleted -> do
@@ -265,29 +310,78 @@ handleTransfer (Piece   blk) = do
 
 handleTransfer (Cancel  bix) = filterQueue (not . (transferResponse bix))
   where
-    transferResponse bix (Transfer (Piece blk)) = blockIx blk == bix
-    transferResponse _    _                     = False
+    transferResponse bix (Transfer (Message.Piece blk)) = blockIx blk == bix
+    transferResponse _    _                             = False
+
+{-----------------------------------------------------------------------
+--  Metadata exchange
+-----------------------------------------------------------------------}
+-- TODO introduce new metadata exchange specific exceptions
+
+tryRequestMetadataBlock :: Wire Session ()
+tryRequestMetadataBlock = do
+  addr <- asks   connRemoteAddr
+  mpix <- lift $ withMetadataUpdates (Metadata.scheduleBlock addr)
+  case mpix of
+    Nothing  -> undefined
+    Just pix -> sendMessage (MetadataRequest pix)
+
+handleMetadata :: Handler ExtendedMetadata
+handleMetadata (MetadataRequest pix) =
+    lift (tryReadMetadataBlock pix) >>= sendMessage . mkResponse
+  where
+    mkResponse  Nothing              = MetadataReject pix
+    mkResponse (Just (piece, total)) = MetadataData   piece total
+
+handleMetadata (MetadataData   {..}) = do
+  addr <- asks connRemoteAddr
+  ih   <- asks connTopic
+  lift $ withMetadataUpdates (Metadata.pushBlock addr piece ih)
+  tryRequestMetadataBlock
+
+handleMetadata (MetadataReject  pix) = do
+  lift $ withMetadataUpdates (Metadata.cancelPending pix)
+
+handleMetadata (MetadataUnknown _  ) = do
+  logInfoN "Unknown metadata message"
+
+waitForMetadata :: Wire Session ()
+waitForMetadata = do
+  Session {..} <- asks connSession
+  needFetch    <- liftIO (isEmptyMVar infodict)
+  when needFetch $ do
+    canFetch   <- allowed ExtMetadata <$> use connExtCaps
+    if canFetch
+      then tryRequestMetadataBlock
+      else liftIO (waitMVar infodict)
 
 {-----------------------------------------------------------------------
 --  Event loop
 -----------------------------------------------------------------------}
 
-handleMessage :: Message -> Wire Session ()
+acceptRehandshake :: ExtendedHandshake -> Wire s ()
+acceptRehandshake ehs = undefined
+
+handleExtended :: Handler ExtendedMessage
+handleExtended (EHandshake     ehs) = acceptRehandshake ehs
+handleExtended (EMetadata  _   msg) = handleMetadata msg
+handleExtended (EUnknown   _   _  ) = logWarnN "Unknown extension message"
+
+handleMessage :: Handler Message
 handleMessage KeepAlive       = return ()
 handleMessage (Status s)      = handleStatus s
 handleMessage (Available msg) = handleAvailable msg
 handleMessage (Transfer  msg) = handleTransfer msg
 handleMessage (Port      n)   = undefined
 handleMessage (Fast      _)   = undefined
-handleMessage (Extended  _)   = undefined
+handleMessage (Extended  msg) = handleExtended msg
 
 exchange :: Wire Session ()
 exchange = do
+  waitForMetadata
   bf <- getThisBitfield
   sendMessage (Bitfield bf)
-  awaitForever $ \ msg -> do
-    logMessage msg
-    handleMessage msg
+  awaitForever handleMessage
 
 data Event = NewMessage (PeerAddr IP) Message
            | Timeout -- for scheduling