6 files changed, 81 insertions, 11 deletions

diff --git a/bittorrent.cabal b/bittorrent.cabal
index a95d6ac2..28188ce3 100644
--- a/bittorrent.cabal
+++ b/bittorrent.cabal
@@ -86,6 +86,7 @@ library
                      , conduit              == 1.*
                      , network-conduit      == 1.*
                      , cereal-conduit       >= 0.5
+                     , resourcet
 
                        -- Misc
                      , lens
diff --git a/exsamples/Main.hs b/exsamples/Main.hs
index a5b865b0..88796642 100644
--- a/exsamples/Main.hs
+++ b/exsamples/Main.hs
@@ -6,7 +6,7 @@ import Data.Bitfield
 import Network.BitTorrent
 import System.Environment
 import Control.Monad.Reader
-
+import Data.IORef
 
 main :: IO ()
 main = do
@@ -16,6 +16,8 @@ main = do
   client  <- newClient []
   swarm   <- newLeacher  client torrent
 
+  ref <- newIORef 0
+
   discover swarm $ do
     addr <- asks connectedPeerAddr
     liftIO $ print $ "connected to" ++ show addr
@@ -27,7 +29,14 @@ main = do
           | Just m <- findMin bf -> yieldEvent (Want (BlockIx m 0 10))
           |     otherwise        -> return ()
         Want     bix -> liftIO $ print bix
-        Fragment blk -> liftIO $ print (ppBlock blk)
+        Fragment blk -> do
+
+          liftIO $ do
+            readIORef ref >>= print
+            modifyIORef ref succ
+            print (ppBlock blk)
+
+          yieldEvent (Want (BlockIx 0 0 (16 * 1024)))
 
 
   print "Bye-bye! =_="
\ No newline at end of file
diff --git a/src/Network/BitTorrent.hs b/src/Network/BitTorrent.hs
index f7c4c004..7ec0a067 100644
--- a/src/Network/BitTorrent.hs
+++ b/src/Network/BitTorrent.hs
@@ -29,9 +29,14 @@ module Network.BitTorrent
        , awaitEvent, yieldEvent
        ) where
 
+import Control.Concurrent
+import Control.Exception
 import Control.Monad
+
 import Data.IORef
 
+import Network
+
 import Data.Torrent
 import Network.BitTorrent.Internal
 import Network.BitTorrent.Exchange
@@ -45,6 +50,8 @@ import Network.BitTorrent.Tracker
 
 discover :: SwarmSession -> P2P () -> IO ()
 discover swarm action = do
+  port <- listener swarm action
+
   let conn = TConnection (tAnnounce (torrentMeta swarm))
                          (tInfoHash (torrentMeta swarm))
                          (clientPeerID (clientSession swarm))
@@ -57,8 +64,14 @@ discover swarm action = do
     forever $ do
       addr <- getPeerAddr tses
       putStrLn "connecting to peer"
-      withPeer swarm addr action
-
+      handle handler (withPeer swarm addr action)
 
+  where
+    handler :: IOException -> IO ()
+    handler _ = return ()
 
-port = 10000
+listener :: SwarmSession -> P2P () -> IO PortNumber
+listener _ _ = do
+  -- TODO:
+--  forkIO loop
+  return 10000
diff --git a/src/Network/BitTorrent/Exchange.hs b/src/Network/BitTorrent/Exchange.hs
index de13d4ce..0e2799b2 100644
--- a/src/Network/BitTorrent/Exchange.hs
+++ b/src/Network/BitTorrent/Exchange.hs
@@ -28,6 +28,7 @@ import Control.Exception
 import Control.Lens
 import Control.Monad.Reader
 import Control.Monad.State
+import Control.Monad.Trans.Resource
 
 import Data.IORef
 import Data.Function
@@ -77,7 +78,7 @@ awaitMessage = P2P (ReaderT go)
       liftIO $ putStrLn "trying recv:"
       mmsg <- await
       case mmsg of
-        Nothing -> go se
+        Nothing  -> monadThrow SessionException
         Just msg -> do
 --          liftIO $ updateIncoming se
           liftIO $ print ("recv:" <+> ppMessage msg)
@@ -250,8 +251,18 @@ awaitEvent = awaitMessage >>= go
 --   @
 --
 yieldEvent  :: Event -> P2P ()
-yieldEvent (Available bf) = undefined
-yieldEvent _ = undefined
+yieldEvent (Available bf)  = undefined
+yieldEvent (Want      bix) = do
+  offer <- peerOffer
+  if ixPiece bix `BF.member` offer
+     then yieldMessage (Request bix)
+     else return ()
+
+yieldEvent (Fragment  blk) = do
+  offer <- clientOffer
+  if blkPiece blk `BF.member` offer
+    then yieldMessage (Piece blk)
+    else return ()
 
 --flushBroadcast :: P2P ()
 --flushBroadcast = nextBroadcast >>= maybe (return ()) go
@@ -265,11 +276,21 @@ checkPiece = undefined
     P2P monad
 -----------------------------------------------------------------------}
 
+-- |
+--   Exceptions:
+--
+--     * SessionException: is visible with one peer session. Use this
+--     exception to terminate P2P session, but not the swarm session.
+--
 newtype P2P a = P2P {
     runP2P :: ReaderT PeerSession PeerWire a
-  } deriving (Functor, Applicative, Monad, MonadReader PeerSession, MonadIO)
+  } deriving ( Functor, Applicative, Monad
+             , MonadIO, MonadThrow, MonadActive
+             , MonadReader PeerSession
+             )
 
 withPeer :: SwarmSession -> PeerAddr -> P2P () -> IO ()
 withPeer se addr p2p =
   withPeerSession se addr $ \(sock, pses) -> do
-    runConduit sock (runReaderT (runP2P p2p) pses)
+    handle putSessionException $
+      runConduit sock (runReaderT (runP2P p2p) pses)
diff --git a/src/Network/BitTorrent/Internal.hs b/src/Network/BitTorrent/Internal.hs
index da104240..a3e1a5cd 100644
--- a/src/Network/BitTorrent/Internal.hs
+++ b/src/Network/BitTorrent/Internal.hs
@@ -16,6 +16,7 @@
 {-# LANGUAGE OverloadedStrings     #-}
 {-# LANGUAGE RecordWildCards       #-}
 {-# LANGUAGE TemplateHaskell       #-}
+{-# LANGUAGE DeriveDataTypeable    #-}
 {-# LANGUAGE FlexibleInstances     #-}
 {-# LANGUAGE FlexibleContexts      #-}
 {-# LANGUAGE MultiParamTypeClasses #-}
@@ -23,20 +24,31 @@
 module Network.BitTorrent.Internal
        ( Progress(..), startProgress
 
+         -- * Client
        , ClientSession (clientPeerID, allowedExtensions)
        , newClient, getCurrentProgress
 
+         -- * Swarm
        , SwarmSession(SwarmSession, torrentMeta, clientSession)
        , newLeacher, newSeeder
 
+         -- * Peer
        , PeerSession(PeerSession, connectedPeerAddr
                     , swarmSession, enabledExtensions
                     )
        , SessionState
+       , withPeerSession
+
+         -- ** Exceptions
+       , SessionException(..)
+       , isSessionException
+       , putSessionException
+       , sessionError
+
+         -- ** Properties
        , bitfield, status
        , emptyBF, fullBF, singletonBF, adjustBF
        , getPieceCount, getClientBF
-       , sessionError, withPeerSession
 
          -- * Timeouts
        , updateIncoming, updateOutcoming
@@ -55,6 +67,7 @@ import Data.Default
 import Data.Function
 import Data.Ord
 import Data.Set as S
+import Data.Typeable
 
 import Data.Serialize hiding (get)
 import Text.PrettyPrint
@@ -212,6 +225,7 @@ data PeerSession = PeerSession {
     -- to avoid reduntant KA messages.
   , outcomingTimeout   :: TimeoutKey
 
+    -- TODO use dupChan for broadcasting
   , broadcastMessages :: Chan   [Message]
   , sessionState      :: IORef   SessionState
   }
@@ -234,6 +248,17 @@ instance (MonadIO m, MonadReader PeerSession m)
   get   = asks sessionState >>= liftIO . readIORef
   put s = asks sessionState >>= \ref -> liftIO $ writeIORef ref s
 
+data SessionException = SessionException
+                        deriving (Show, Typeable)
+
+instance Exception SessionException
+
+isSessionException :: Monad m => SessionException -> m ()
+isSessionException _ = return ()
+
+putSessionException :: SessionException -> IO ()
+putSessionException = print
+
 sessionError :: MonadIO m => Doc -> m ()
 sessionError msg
   = liftIO $ throwIO $ userError $ render $ msg <+> "in session"
diff --git a/src/Network/BitTorrent/Tracker.hs b/src/Network/BitTorrent/Tracker.hs
index 179dae1d..11bc52de 100644
--- a/src/Network/BitTorrent/Tracker.hs
+++ b/src/Network/BitTorrent/Tracker.hs
@@ -145,6 +145,7 @@ data TSession = TSession {
   , sePeers      :: Chan PeerAddr
     -- TODO use something like 'TVar (Set PeerAddr)'
     -- otherwise we might get space leak
+    -- TODO or maybe BoundedChan?
   }
 
 newSession :: Progress -> Int -> [PeerAddr] -> IO TSession