3 files changed, 132 insertions, 44 deletions

diff --git a/src/Network/BitTorrent/Exchange/Message.hs b/src/Network/BitTorrent/Exchange/Message.hs
index 8fcf582f..8ef9f3da 100644
--- a/src/Network/BitTorrent/Exchange/Message.hs
+++ b/src/Network/BitTorrent/Exchange/Message.hs
@@ -210,9 +210,12 @@ instance Serialize Handshake where
               <*> S.get
               <*> S.get
 
+-- | Show handshake protocol string, caps and fingerprint.
 instance Pretty Handshake where
   pretty Handshake {..}
-    = text (BC.unpack hsProtocol) <+> pretty (fingerprint hsPeerId)
+    = text (BC.unpack hsProtocol) $$
+      pretty hsReserved           $$
+      pretty (fingerprint hsPeerId)
 
 -- | Get handshake message size in bytes from the length of protocol
 -- string.
@@ -227,31 +230,49 @@ handshakeMaxSize = handshakeSize maxBound
 defaultBTProtocol :: BS.ByteString
 defaultBTProtocol = "BitTorrent protocol"
 
--- | Length of info hash and peer id is unchecked, so it /should/ be
--- equal 20.
+-- | Handshake with default protocol string and reserved bitmask.
 defaultHandshake :: InfoHash -> PeerId -> Handshake
 defaultHandshake = Handshake defaultBTProtocol def
 
 {-----------------------------------------------------------------------
-    Regular messages
+--  Regular messages
 -----------------------------------------------------------------------}
 
 class PeerMessage a where
-  envelop :: ExtendedCaps -> a -> Message
+  -- | Construct a message to be /sent/.
+  envelop :: ExtendedCaps -- ^ The /receiver/ extended capabilities;
+          -> a            -- ^ An regular message;
+          -> Message      -- ^ Enveloped message to sent.
 
+{-----------------------------------------------------------------------
+--  Status messages
+-----------------------------------------------------------------------}
+
+-- | Notification that the sender have updated its
+-- 'Network.BitTorrent.Exchange.Status.PeerStatus'.
 data StatusUpdate
-  = Choke
-  | Unchoke
-  | Interested
-  | NotInterested
-    deriving (Show, Eq, Ord, Enum, Bounded)
+    -- | Notification that the sender will not upload data to the
+    -- receiver until unchoking happen.
+  = Choking    !Bool
+
+    -- | Notification that the sender is interested (or not interested)
+    -- in any of the receiver's data pieces.
+  | Interested !Bool
+    deriving (Show, Eq, Ord, Typeable)
 
 instance Pretty StatusUpdate where
-  pretty = text . show
+  pretty (Choking    False) = "not choking"
+  pretty (Choking    True ) = "choking"
+  pretty (Interested False) = "not interested"
+  pretty (Interested True ) = "interested"
 
 instance PeerMessage StatusUpdate where
   envelop _ = Status
 
+{-----------------------------------------------------------------------
+--  Available and transfer messages
+-----------------------------------------------------------------------}
+
 data RegularMessage =
     -- | Zero-based index of a piece that has just been successfully
     -- downloaded and verified via the hash.
@@ -303,6 +324,10 @@ instance PeerMessage BlockIx where
 instance PeerMessage (Block BL.ByteString) where
   envelop c = envelop c . Piece
 
+{-----------------------------------------------------------------------
+--  Fast messages
+-----------------------------------------------------------------------}
+
 -- | BEP6 messages.
 data FastMessage =
     -- | If a peer have all pieces it might send the 'HaveAll' message
@@ -526,12 +551,14 @@ type MessageId = Word8
 --   extension then the client MUST close the connection.
 --
 data Message
-    -- core
+    -- | Peers may close the TCP connection if they have not received
+    -- any messages for a given period of time, generally 2
+    -- minutes. Thus, the "keep-alive" message is sent tot keep the
+    -- connection between two peers alive, if no /other/ message has
+    -- been sentin a given period of time.
   = KeepAlive
   | Status   !StatusUpdate
   | Regular  !RegularMessage
-
-    -- extensions
   | Port     !PortNumber
   | Fast     !FastMessage
   | Extended !ExtendedMessage
@@ -581,10 +608,10 @@ instance Serialize Message where
       else do
         mid <- S.getWord8
         case mid of
-          0x00 -> return $ Status Choke
-          0x01 -> return $ Status Unchoke
-          0x02 -> return $ Status Interested
-          0x03 -> return $ Status NotInterested
+          0x00 -> return $ Status (Choking True)
+          0x01 -> return $ Status (Choking False)
+          0x02 -> return $ Status (Interested True)
+          0x03 -> return $ Status (Interested False)
           0x04 -> (Regular . Have)    <$> getInt
           0x05 -> (Regular . Bitfield . fromBitmap)
                                       <$> S.getByteString (pred len)
@@ -616,8 +643,12 @@ instance Serialize Message where
   put (Fast    msg) = putFast    msg
   put (Extended m ) = putExtendedMessage m
 
+statusUpdateId :: StatusUpdate -> MessageId
+statusUpdateId (Choking    choking) = fromIntegral (0 + fromEnum choking)
+statusUpdateId (Interested choking) = fromIntegral (2 + fromEnum choking)
+
 putStatus :: Putter StatusUpdate
-putStatus su = putInt 1  >> S.putWord8 (fromIntegral (fromEnum su))
+putStatus su = putInt 1  >> S.putWord8 (statusUpdateId su)
 
 putRegular :: Putter RegularMessage
 putRegular (Have i)      = putInt 5  >> S.putWord8 0x04 >> putInt i
diff --git a/src/Network/BitTorrent/Exchange/Status.hs b/src/Network/BitTorrent/Exchange/Status.hs
index ae323e09..42766428 100644
--- a/src/Network/BitTorrent/Exchange/Status.hs
+++ b/src/Network/BitTorrent/Exchange/Status.hs
@@ -1,80 +1,134 @@
+-- |
+--   Copyright   :  (c) Sam Truzjan 2013
+--   License     :  BSD3
+--   Maintainer  :  pxqr.sta@gmail.com
+--   Stability   :  experimental
+--   Portability :  portable
+--
+--   Each P2P connection endpoint should keep track status of both
+--   sides.
+--
 {-# LANGUAGE TemplateHaskell   #-}
 module Network.BitTorrent.Exchange.Status
        ( -- * Peer status
          PeerStatus(..)
        , choking
        , interested
+
+         -- ** Query
        , updateStatus
+       , statusUpdates
 
          -- * Session status
        , SessionStatus(..)
        , clientStatus
-       , peerStatus
+       , remoteStatus
 
          -- ** Query
        , canUpload
        , canDownload
 
          -- * Extra
-       , inverseStatus
        , defaultUnchokeSlots
+       , defaultRechokeInterval
        ) where
 
 import Control.Lens
 import Data.Aeson.TH
-import Data.List as L
 import Data.Default
+import Data.List as L
+import Data.Maybe
+import Data.Monoid
+import Text.PrettyPrint as PP hiding ((<>))
+import Text.PrettyPrint.Class
 
 import Network.BitTorrent.Exchange.Message
 
+{-----------------------------------------------------------------------
+--  Peer status
+-----------------------------------------------------------------------}
 
--- |
-data PeerStatus = PeerStatus {
+-- | Connections contain two bits of state on either end: choked or
+-- not, and interested or not.
+data PeerStatus = PeerStatus
+  { -- | Choking is a notification that no data will be sent until
+    -- unchoking happens.
     _choking    :: !Bool
+
+    -- |
   , _interested :: !Bool
-  } deriving (Show, Eq)
+  } deriving (Show, Eq, Ord)
 
 $(makeLenses ''PeerStatus)
 $(deriveJSON defaultOptions { fieldLabelModifier = L.tail } ''PeerStatus)
 
+instance Pretty PeerStatus where
+  pretty PeerStatus {..} =
+    pretty (Choking _choking) <+> "and" <+> pretty (Interested _interested)
+
+-- | Connections start out choked and not interested.
 instance Default PeerStatus where
   def = PeerStatus True False
 
+instance Monoid PeerStatus where
+  mempty      = def
+  mappend a b = PeerStatus
+    { _choking    = _choking    a && _choking    b
+    , _interested = _interested a || _interested b
+    }
+
+-- | Can be used to update remote peer status using incoming 'Status'
+-- message.
 updateStatus :: StatusUpdate -> PeerStatus -> PeerStatus
-updateStatus Choke         = choking    .~ True
-updateStatus Unchoke       = choking    .~ False
-updateStatus Interested    = interested .~ True
-updateStatus NotInterested = interested .~ False
+updateStatus (Choking    b) = choking    .~ b
+updateStatus (Interested b) = interested .~ b
 
+-- | Can be used to generate outcoming messages.
 statusUpdates :: PeerStatus -> PeerStatus -> [StatusUpdate]
-statusUpdates a b = undefined
-
--- |
-data SessionStatus = SessionStatus {
-    _clientStatus :: !PeerStatus
-  , _peerStatus   :: !PeerStatus
+statusUpdates a b = catMaybes $
+  [ if _choking    a == _choking    b then Nothing
+    else Just $ Choking    $ _choking    b
+  , if _interested a == _interested b then Nothing
+    else Just $ Interested $ _interested b
+  ]
+
+{-----------------------------------------------------------------------
+--  Session status
+-----------------------------------------------------------------------}
+
+-- | Status of the both endpoints.
+data SessionStatus = SessionStatus
+  { _clientStatus :: !PeerStatus
+  , _remoteStatus :: !PeerStatus
   } deriving (Show, Eq)
 
 $(makeLenses ''SessionStatus)
 $(deriveJSON defaultOptions { fieldLabelModifier = L.tail } ''SessionStatus)
 
+instance Pretty SessionStatus where
+  pretty SessionStatus {..} =
+    "this  " <+> pretty _clientStatus $$
+    "remote" <+> pretty _remoteStatus
+
+-- | Connections start out choked and not interested.
 instance Default SessionStatus where
   def = SessionStatus def def
 
--- | Can the /client/ transfer to the /peer/?
+-- | Can the client transfer to the remote peer?
 canUpload :: SessionStatus -> Bool
 canUpload SessionStatus {..}
-  = _interested _peerStatus && not (_choking _clientStatus)
+  = _interested _remoteStatus && not (_choking _clientStatus)
 
--- | Can the /client/ transfer from the /peer/?
+-- | Can the client transfer from the remote peer?
 canDownload :: SessionStatus -> Bool
 canDownload SessionStatus {..}
-  = _interested _clientStatus && not (_choking _peerStatus)
-
-inverseStatus :: SessionStatus -> SessionStatus
-inverseStatus SessionStatus {..} = SessionStatus _peerStatus _clientStatus
+  = _interested _clientStatus && not (_choking _remoteStatus)
 
 -- | Indicates how many peers are allowed to download from the client
 -- by default.
 defaultUnchokeSlots :: Int
-defaultUnchokeSlots = 4
\ No newline at end of file
+defaultUnchokeSlots = 4
+
+-- |
+defaultRechokeInterval :: Int
+defaultRechokeInterval = 10 * 1000 * 1000
\ No newline at end of file
diff --git a/src/Network/BitTorrent/Exchange/Wire.hs b/src/Network/BitTorrent/Exchange/Wire.hs
index 6a161762..680da059 100644
--- a/src/Network/BitTorrent/Exchange/Wire.hs
+++ b/src/Network/BitTorrent/Exchange/Wire.hs
@@ -61,13 +61,16 @@ data ChannelSide
   | RemotePeer
     deriving (Show, Eq, Enum)
 
+-- TODO pretty instance
+
 -- | Errors occur when a remote peer violates protocol specification.
 data ProtocolError
   = UnexpectedTopic   InfoHash -- ^ peer replied with unexpected infohash.
   | UnexpectedPeerId  PeerId   -- ^ peer replied with unexpected peer id.
   | UnknownTopic      InfoHash -- ^ peer requested unknown torrent.
   | HandshakeRefused           -- ^ peer do not send an extended handshake back.
-  | InvalidMessage
+  | BitfieldAlreadSend ChannelSide
+  | InvalidMessage -- TODO caps violation
     { violentSender     :: ChannelSide -- ^ endpoint sent invalid message
     , extensionRequired :: Extension   -- ^
     }