1 files changed, 65 insertions, 1 deletions

diff --git a/src/Network/BitTorrent/Exchange/Protocol.hs b/src/Network/BitTorrent/Exchange/Protocol.hs
index 8cfcc79d..505f6ac6 100644
--- a/src/Network/BitTorrent/Exchange/Protocol.hs
+++ b/src/Network/BitTorrent/Exchange/Protocol.hs
@@ -34,6 +34,21 @@ module Network.BitTorrent.Exchange.Protocol
          -- * Regular messages
        , Message(..)
        , ppMessage
+
+         -- * Exchange control
+         -- ** Peer status
+       , PeerStatus(..)
+       , setChoking, setInterested
+       , initPeerStatus
+
+         -- ** Session status
+       , SessionStatus(..)
+       , initSessionStatus
+       , setClientStatus, setPeerStatus
+       , canUpload, canDownload
+
+       -- ** Defaults
+       , defaultUnchokeSlots
        ) where
 
 import Control.Applicative
@@ -384,4 +399,53 @@ ppMessage (Piece blk)        = "Piece"    <+> ppBlock blk
 ppMessage (Cancel ix)        = "Cancel"   <+> ppBlockIx ix
 ppMessage (SuggestPiece pix) = "Suggest"  <+> int pix
 ppMessage (RejectRequest ix) = "Reject"   <+> ppBlockIx ix
-ppMessage msg = text (show msg)
\ No newline at end of file
+ppMessage msg = text (show msg)
+
+{-----------------------------------------------------------------------
+    Peer Status
+-----------------------------------------------------------------------}
+
+data PeerStatus = PeerStatus {
+    psChoking    :: Bool
+  , psInterested :: Bool
+  }
+
+-- | Any session between peers starts as choking and not interested.
+initPeerStatus :: PeerStatus
+initPeerStatus = PeerStatus True False
+
+setChoking :: Bool -> PeerStatus -> PeerStatus
+setChoking b ps = ps { psChoking = b }
+
+setInterested :: Bool -> PeerStatus -> PeerStatus
+setInterested b ps = ps { psInterested = b }
+
+
+
+data SessionStatus = SessionStatus {
+    seClientStatus :: PeerStatus
+  , sePeerStatus   :: PeerStatus
+  }
+
+initSessionStatus :: SessionStatus
+initSessionStatus = SessionStatus initPeerStatus initPeerStatus
+
+setClientStatus :: (PeerStatus -> PeerStatus) -> SessionStatus -> SessionStatus
+setClientStatus f ss = ss { seClientStatus = f (seClientStatus ss) }
+
+setPeerStatus :: (PeerStatus -> PeerStatus) -> SessionStatus -> SessionStatus
+setPeerStatus f ss = ss { sePeerStatus = f (sePeerStatus ss) }
+
+-- | Can the /client/ to upload to the /peer/?
+canUpload :: SessionStatus -> Bool
+canUpload SessionStatus { seClientStatus = client, sePeerStatus = peer}  =
+  psInterested peer && not (psChoking client)
+
+-- | Can the /client/ download from the /peer/?
+canDownload :: SessionStatus -> Bool
+canDownload SessionStatus { seClientStatus = client, sePeerStatus = peer }  =
+  psInterested client && not (psChoking peer)
+
+-- | Indicates have many peers are allowed to download from the client.
+defaultUnchokeSlots :: Int
+defaultUnchokeSlots = 4
\ No newline at end of file