3 files changed, 58 insertions, 60 deletions

diff --git a/src/Network/BitTorrent/Exchange.hs b/src/Network/BitTorrent/Exchange.hs
index 2173cf8b..dda7d304 100644
--- a/src/Network/BitTorrent/Exchange.hs
+++ b/src/Network/BitTorrent/Exchange.hs
@@ -60,10 +60,11 @@ waitMessage se = do
     Nothing -> waitMessage se
     Just msg -> do
       liftIO $ updateIncoming se
+      liftIO $ print msg
       return msg
 
-signalMessage :: Message -> PeerSession -> PeerWire ()
-signalMessage msg se = do
+signalMessage :: PeerSession -> Message -> PeerWire ()
+signalMessage se msg = do
   C.yield msg
   liftIO $ updateOutcoming se
 
@@ -71,58 +72,75 @@ signalMessage msg se = do
 getPieceCount :: PeerSession -> IO PieceCount
 getPieceCount = undefined
 
+canOffer :: PeerSession -> PeerWire Bitfield
+canOffer PeerSession {..} = liftIO $ do
+  pbf <- readIORef $ peerBitfield
+  cbf <- readIORef $ clientBitfield $ swarmSession
+  return $ BF.difference pbf cbf
+
+revise :: PeerSession -> PeerWire ()
+revise se @ PeerSession {..} = do
+  isInteresting <- (not . BF.null) <$> canOffer se
+  SessionStatus {..} <- liftIO $ readIORef peerSessionStatus
+
+  when (isInteresting /= _interested seClientStatus) $
+    signalMessage se $ if isInteresting then Interested else NotInterested
+
+
 nextEvent :: PeerSession -> PeerWire Event
-nextEvent se @ PeerSession {..} = waitMessage se >>= diff
+nextEvent se @ PeerSession {..} = waitMessage se >>= go
   where
-    -- diff finds difference between
---    diff KeepAlive = nextEvent se
-    diff msg = do
-      liftIO $ print (ppMessage msg)
-      nextEvent se
-
-    handleMessage Choke     = do
+    go KeepAlive = nextEvent se
+    go Choke     = do
       SessionStatus {..} <- liftIO $ readIORef peerSessionStatus
-      if psChoking  sePeerStatus
+      if _choking  sePeerStatus
         then nextEvent se
         else undefined
 
-    handleMessage Unchoke   = undefined
+    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
 
-    handleMessage Interested = return undefined
-    handleMessage NotInterested = return undefined
-    handleMessage (Have ix) = do
+    go Interested = return undefined
+    go NotInterested = return undefined
+
+    go (Have ix) = do
       pc <- liftIO $ getPieceCount se
       haveMessage $ have ix (haveNone pc) -- TODO singleton
 
-    handleMessage (Bitfield bf) = undefined
-    handleMessage (Request  bix) = do
+    go (Bitfield bf) = undefined
+    go (Request  bix) = do
       undefined
 
-    handleMessage msg @ (Piece    blk) = undefined
-    handleMessage msg @ (Port     _)
+    go msg @ (Piece    blk) = undefined
+    go msg @ (Port     _)
       = checkExtension msg ExtDHT $ do
       undefined
 
-    handleMessage msg @ HaveAll
+    go msg @ HaveAll
       = checkExtension msg ExtFast $ do
           pc <- liftIO $ getPieceCount se
           haveMessage (haveAll pc)
 
-    handleMessage msg @ HaveNone
+    go msg @ HaveNone
       = checkExtension msg ExtFast $ do
         pc <- liftIO $ getPieceCount se
         haveMessage (haveNone pc)
 
-    handleMessage msg @ (SuggestPiece ix)
+    go msg @ (SuggestPiece ix)
           = checkExtension msg ExtFast $ do
             undefined
 
-    handleMessage msg @ (RejectRequest ix)
+    go msg @ (RejectRequest ix)
           = checkExtension msg ExtFast $ do
             undefined
 
-    handleMessage msg @ (AllowedFast pix)
+    go msg @ (AllowedFast pix)
           = checkExtension msg ExtFast $ do
             undefined
 
@@ -148,7 +166,10 @@ newtype P2P a = P2P {
     runP2P :: ReaderT PeerSession PeerWire a
   } deriving (Monad, MonadReader PeerSession, MonadIO)
 
-instance MonadState Bitfield P2P where
+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 =
diff --git a/src/Network/BitTorrent/Exchange/Protocol.hs b/src/Network/BitTorrent/Exchange/Protocol.hs
index dc25a9c9..46e25fa3 100644
--- a/src/Network/BitTorrent/Exchange/Protocol.hs
+++ b/src/Network/BitTorrent/Exchange/Protocol.hs
@@ -51,16 +51,9 @@ module Network.BitTorrent.Exchange.Protocol
        , ppMessage
 
          -- * Exchange control
-         -- ** Peer status
        , PeerStatus(..)
-       , setChoking, setInterested
-       , initPeerStatus
-
-         -- ** Session status
        , SessionStatus(..)
-       , initSessionStatus
-       , setClientStatus, setPeerStatus
-       , canUpload, canDownload
+--       , canUpload, canDownload
 
        -- ** Defaults
        , defaultUnchokeSlots
@@ -73,6 +66,7 @@ import           Data.ByteString (ByteString)
 import qualified Data.ByteString as B
 import qualified Data.ByteString.Char8 as BC
 import qualified Data.ByteString.Lazy as Lazy
+import Data.Default
 import Data.Serialize as S
 import Data.Int
 import Data.Word
@@ -429,21 +423,12 @@ ppMessage msg = text (show msg)
 
 -- |
 data PeerStatus = PeerStatus {
-    psChoking    :: Bool
-  , psInterested :: Bool
+    _choking    :: Bool
+  , _interested :: Bool
   }
 
--- | Any session between peers starts as choking and not interested.
-initPeerStatus :: PeerStatus
-initPeerStatus = PeerStatus True False
-
--- | Update choking field.
-setChoking :: Bool -> PeerStatus -> PeerStatus
-setChoking b ps = ps { psChoking = b }
-
--- | Update interested field.
-setInterested :: Bool -> PeerStatus -> PeerStatus
-setInterested b ps = ps { psInterested = b }
+instance Default PeerStatus where
+  def = PeerStatus True False
 
 -- |
 data SessionStatus = SessionStatus {
@@ -451,20 +436,10 @@ data SessionStatus = SessionStatus {
   , sePeerStatus   :: PeerStatus
   }
 
--- | Initial session status after two peers handshaked.
-initSessionStatus :: SessionStatus
-initSessionStatus = SessionStatus initPeerStatus initPeerStatus
-
--- | Update client status.
-setClientStatus :: (PeerStatus -> PeerStatus)
-                -> SessionStatus -> SessionStatus
-setClientStatus f ss = ss { seClientStatus = f (seClientStatus ss) }
-
--- | Update peer status.
-setPeerStatus :: (PeerStatus -> PeerStatus)
-              -> SessionStatus -> SessionStatus
-setPeerStatus f ss = ss { sePeerStatus = f (sePeerStatus ss) }
+instance Default SessionStatus where
+  def = SessionStatus def def
 
+{-
 -- | Can the /client/ to upload to the /peer/?
 canUpload :: SessionStatus -> Bool
 canUpload SessionStatus {..}
@@ -474,6 +449,7 @@ canUpload SessionStatus {..}
 canDownload :: SessionStatus -> Bool
 canDownload SessionStatus {..}
   = psInterested seClientStatus && not (psChoking sePeerStatus)
+-}
 
 -- | 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 91dc35d5..38087f0d 100644
--- a/src/Network/BitTorrent/Internal.hs
+++ b/src/Network/BitTorrent/Internal.hs
@@ -27,6 +27,7 @@ import Control.Concurrent.STM
 import Control.Exception
 
 import Data.IORef
+import Data.Default
 import Data.Function
 import Data.Ord
 import Data.Set as S
@@ -196,7 +197,7 @@ withPeerSession ss @ SwarmSession {..} addr
                 maxOutcomingTime (sendKA sock)
          <*> newChan
          <*> pure clientBitfield
-         <*> newIORef initSessionStatus
+         <*> newIORef def
       return (sock, ps)
 
     closeSession = close . fst