Add stats method to PeerMessage class

6 files changed, 210 insertions, 54 deletions

diff --git a/bittorrent.cabal b/bittorrent.cabal
index 3ad1d7e4..dbb3c3be 100644
--- a/bittorrent.cabal
+++ b/bittorrent.cabal
@@ -152,7 +152,8 @@ test-suite spec
   type:                exitcode-stdio-1.0
   hs-source-dirs:      tests
   main-is:             Spec.hs
-  other-modules:       Data.Torrent.InfoHashSpec
+  other-modules:       Data.Torrent.BitfieldSpec
+                       Data.Torrent.InfoHashSpec
                        Data.Torrent.LayoutSpec
                        Data.Torrent.MagnetSpec
                        Data.Torrent.MetainfoSpec
diff --git a/src/Network/BitTorrent/Exchange/Message.hs b/src/Network/BitTorrent/Exchange/Message.hs
index 38a8ac33..b3100269 100644
--- a/src/Network/BitTorrent/Exchange/Message.hs
+++ b/src/Network/BitTorrent/Exchange/Message.hs
@@ -46,10 +46,15 @@ module Network.BitTorrent.Exchange.Message
        , handshakeSize
        , handshakeMaxSize
 
+         -- * Stats
+       , ByteCount
+       , ByteStats   (..)
+       , byteLength
+
          -- * Messages
        , Message        (..)
-       , PeerMessage    (..)
        , defaultKeepAliveInterval
+       , PeerMessage    (..)
 
          -- ** Core messages
        , StatusUpdate   (..)
@@ -292,6 +297,47 @@ defaultHandshake :: InfoHash -> PeerId -> Handshake
 defaultHandshake = Handshake def def
 
 {-----------------------------------------------------------------------
+--  Stats
+-----------------------------------------------------------------------}
+
+-- | Number of bytes.
+type ByteCount = Int
+
+-- | Summary of encoded message byte layout can be used to collect
+-- stats about message flow in both directions. This data can be
+-- retrieved using 'stats' function.
+data ByteStats = ByteStats
+  { -- | Number of bytes used to help encode 'control' and 'payload'
+    -- bytes: message size, message ID's, etc
+    overhead :: {-# UNPACK #-} !ByteCount
+
+    -- | Number of bytes used to exchange peers state\/options: piece
+    -- and block indexes, infohash, port numbers, peer ID\/IP, etc.
+  , control  :: {-# UNPACK #-} !ByteCount
+
+    -- | Number of payload bytes: torrent data blocks and infodict
+    -- metadata.
+  , payload  :: {-# UNPACK #-} !ByteCount
+  } deriving Show
+
+-- | Empty byte sequences.
+instance Default ByteStats where
+  def = ByteStats 0 0 0
+
+-- | Monoid under addition.
+instance Monoid ByteStats where
+  mempty      = def
+  mappend a b = ByteStats
+    { overhead = overhead a + overhead b
+    , control  = control  a + control  b
+    , payload  = payload  a + payload  b
+    }
+
+-- | Sum of the all byte sequences.
+byteLength :: ByteStats -> Int
+byteLength ByteStats {..} = overhead + control + payload
+
+{-----------------------------------------------------------------------
 --  Regular messages
 -----------------------------------------------------------------------}
 
@@ -311,6 +357,17 @@ class PeerMessage a where
   requires :: a -> Maybe Extension
   requires _ = Nothing
 
+  -- | Get sizes of overhead\/control\/payload byte sequences of
+  -- binary message representation without encoding message to binary
+  -- bytestring.
+  --
+  --   This function should obey one law:
+  --
+  --     * 'byteLength' ('stats' msg) == 'BL.length' ('encode' msg)
+  --
+  stats :: a -> ByteStats
+  stats _ = ByteStats 4 0 0
+
 {-----------------------------------------------------------------------
 --  Status messages
 -----------------------------------------------------------------------}
@@ -337,6 +394,9 @@ instance PeerMessage StatusUpdate where
   envelop _ = Status
   {-# INLINE envelop #-}
 
+  stats _ = ByteStats 4 1 0
+  {-# INLINE stats #-}
+
 {-----------------------------------------------------------------------
 --  Available messages
 -----------------------------------------------------------------------}
@@ -361,12 +421,14 @@ instance Pretty Available where
 
 instance PeerMessage Available where
   envelop _ = Available
-
--- | BITFIELD message.
-instance PeerMessage Bitfield where
-  envelop c = envelop c . Bitfield
   {-# INLINE envelop #-}
 
+  stats (Have      _) = ByteStats (4 + 1) 4 0
+  stats (Bitfield bf) = ByteStats (4 + 1) (q + trailing)  0
+    where
+      trailing = if r == 0 then 0 else 1
+      (q, r) = quotRem (totalCount bf) 8
+
 {-----------------------------------------------------------------------
 --  Transfer messages
 -----------------------------------------------------------------------}
@@ -395,15 +457,9 @@ instance PeerMessage Transfer where
   envelop _ = Transfer
   {-# INLINE envelop #-}
 
--- | REQUEST message.
-instance PeerMessage BlockIx where
-  envelop c = envelop c . Request
-  {-# INLINE envelop #-}
-
--- | PIECE message.
-instance PeerMessage (Block BL.ByteString) where
-  envelop c = envelop c . Piece
-  {-# INLINE envelop #-}
+  stats (Request _  ) = ByteStats (4 + 1) (3 * 4) 0
+  stats (Piece   pi ) = ByteStats (4 + 1) (4 + 4 + blockSize pi) 0
+  stats (Cancel  _  ) = ByteStats (4 + 1) (3 * 4) 0
 
 {-----------------------------------------------------------------------
 --  Fast messages
@@ -424,11 +480,12 @@ data FastMessage =
     -- amount of IO.
   | SuggestPiece  !PieceIx
 
-    -- | Notifies a requesting peer that its request will not be satisfied.
+    -- | Notifies a requesting peer that its request will not be
+    -- satisfied.
   | RejectRequest !BlockIx
 
-    -- | This is an advisory messsage meaning "if you ask for this
-    -- piece, I'll give it to you even if you're choked." Used to
+    -- | This is an advisory messsage meaning \"if you ask for this
+    -- piece, I'll give it to you even if you're choked.\" Used to
     -- shorten starting phase.
   | AllowedFast   !PieceIx
     deriving (Show, Eq)
@@ -447,6 +504,12 @@ instance PeerMessage FastMessage where
   requires _ = Just ExtFast
   {-# INLINE requires #-}
 
+  stats  HaveAll          = ByteStats 4 1  0
+  stats  HaveNone         = ByteStats 4 1  0
+  stats (SuggestPiece  _) = ByteStats 5 4  0
+  stats (RejectRequest _) = ByteStats 5 12 0
+  stats (AllowedFast   _) = ByteStats 5 4  0
+
 {-----------------------------------------------------------------------
 --  Extension protocol
 -----------------------------------------------------------------------}
@@ -588,7 +651,7 @@ extHandshakeId = 0
 
 -- | Default 'Request' queue size.
 defaultQueueLength :: Int
-defaultQueueLength = 0
+defaultQueueLength = 1
 
 -- | All fields are empty.
 instance Default ExtendedHandshake where
@@ -619,6 +682,7 @@ instance BEncode ExtendedHandshake where
 instance Pretty ExtendedHandshake where
   pretty = PP.text . show
 
+-- | NOTE: Approximated 'stats'.
 instance PeerMessage ExtendedHandshake where
   envelop  c = envelop c . EHandshake
   {-# INLINE envelop #-}
@@ -626,6 +690,9 @@ instance PeerMessage ExtendedHandshake where
   requires _ = Just ExtExtended
   {-# INLINE requires #-}
 
+  stats _ = ByteStats (4 + 1 + 1) 100 {- is it ok? -} 0 -- FIXME
+  {-# INLINE stats #-}
+
 -- | Set default values and the specified 'ExtendedCaps'.
 nullExtendedHandshake :: ExtendedCaps -> ExtendedHandshake
 nullExtendedHandshake caps = ExtendedHandshake
@@ -721,6 +788,7 @@ instance Pretty ExtendedMetadata where
   pretty (MetadataReject  pix  ) = "Reject"  <+> PP.int    pix
   pretty (MetadataUnknown bval ) = "Unknown" <+> ppBEncode bval
 
+-- | NOTE: Approximated 'stats'.
 instance PeerMessage ExtendedMetadata where
   envelop  c = envelop c . EMetadata (remoteMessageId ExtMetadata c)
   {-# INLINE envelop #-}
@@ -728,6 +796,14 @@ instance PeerMessage ExtendedMetadata where
   requires _ = Just ExtExtended
   {-# INLINE requires #-}
 
+  stats (MetadataRequest _) = ByteStats (4 + 1 + 1) {- ~ -} 25 0
+  stats (MetadataData pi t) = ByteStats (4 + 1 + 1) {- ~ -} 41 $
+                              BS.length (Data.pieceData pi)
+  stats (MetadataReject  _) = ByteStats (4 + 1 + 1) {- ~ -} 25 0
+  stats (MetadataUnknown _) = ByteStats (4 + 1 + 1) {- ? -} 0  0
+
+-- | All 'Piece's in 'MetadataData' messages MUST have size equal to
+-- this value. The last trailing piece can be shorter.
 metadataPieceSize :: Int
 metadataPieceSize = 16 * 1024
 
@@ -791,6 +867,10 @@ instance PeerMessage ExtendedMessage where
   requires _ = Just ExtExtended
   {-# INLINE requires #-}
 
+  stats (EHandshake  hs)  = stats hs
+  stats (EMetadata _ msg) = stats msg
+  stats (EUnknown  _ msg) = ByteStats (4 + 1 + 1) (BS.length msg) 0
+
 {-----------------------------------------------------------------------
 -- The message datatype
 -----------------------------------------------------------------------}
@@ -849,6 +929,14 @@ instance PeerMessage Message where
   requires (Fast      _) = Just ExtFast
   requires (Extended  _) = Just ExtExtended
 
+  stats  KeepAlive    = ByteStats 4 0 0
+  stats (Status    m) = stats m
+  stats (Available m) = stats m
+  stats (Transfer  m) = stats m
+  stats (Port      _) = ByteStats 5 2 0
+  stats (Fast      m) = stats m
+  stats (Extended  m) = stats m
+
 -- | PORT message.
 instance PeerMessage PortNumber where
   envelop  _ = Port
diff --git a/src/Network/BitTorrent/Exchange/Wire.hs b/src/Network/BitTorrent/Exchange/Wire.hs
index a6ee35d8..fe4086bc 100644
--- a/src/Network/BitTorrent/Exchange/Wire.hs
+++ b/src/Network/BitTorrent/Exchange/Wire.hs
@@ -33,9 +33,12 @@ module Network.BitTorrent.Exchange.Wire
        , getExtCaps
 
          -- ** Messaging
+       , recvMessage
+       , sendMessage
+
        , validate
        , validateBoth
-       , keepStats
+       , trackStats
 
          -- ** Stats
        , ConnectionStats (..)
@@ -70,12 +73,16 @@ import Text.PrettyPrint.Class
 import Data.Torrent.InfoHash
 import Network.BitTorrent.Core
 import Network.BitTorrent.Exchange.Message
+import Data.Torrent
+import Data.Torrent.Piece
+import Data.BEncode as BE
 
 -- TODO handle port message?
 -- TODO handle limits?
 -- TODO filter not requested PIECE messages
 -- TODO metadata piece request flood protection
 -- TODO piece request flood protection
+-- TODO protect against flood attacks
 {-----------------------------------------------------------------------
 --  Exceptions
 -----------------------------------------------------------------------}
@@ -130,50 +137,34 @@ isWireFailure _ = return ()
 --  Stats
 -----------------------------------------------------------------------}
 
-type ByteCount = Int
-
-data MessageStats = MessageStats
-  { overhead :: {-# UNPACK #-} !ByteCount
-  , payload  :: {-# UNPACK #-} !ByteCount
+data FlowStats = FlowStats
+  { messageBytes :: {-# UNPACK #-} !ByteStats
+  , messageCount :: {-# UNPACK #-} !Int
+    -- msgTypes :: Map MessageType Int
   } deriving Show
 
-instance Default MessageStats where
-  def = MessageStats 0 0
-
-instance Monoid MessageStats where
-  mempty      = mempty
-  mappend a b = MessageStats
-    { overhead = overhead a + overhead b
-    , payload  = payload  a + payload  b
-    }
-
-
-messageSize :: MessageStats -> Int
-messageSize MessageStats {..} = overhead + payload
-
-messageStats :: Message -> MessageStats
-messageStats = undefined
-
+-- | Note that this is stats is completely different from Progress:
+--   TODO explain why.
 data ConnectionStats = ConnectionStats
-  { incomingFlow  :: !MessageStats
-  , outcomingFlow :: !MessageStats
+  { incomingFlow  :: !ByteStats
+  , outcomingFlow :: !ByteStats
   } deriving Show
 
 instance Default ConnectionStats where
   def = ConnectionStats def def
 
-addStats :: ChannelSide -> MessageStats -> ConnectionStats -> ConnectionStats
+addStats :: ChannelSide -> ByteStats -> ConnectionStats -> ConnectionStats
 addStats ThisPeer   x s = s { outcomingFlow = outcomingFlow s <> x }
 addStats RemotePeer x s = s { incomingFlow  = incomingFlow  s <> x }
 
 recvBytes :: ConnectionStats -> Int
-recvBytes = messageSize . incomingFlow
+recvBytes = byteLength . incomingFlow
 
 sentBytes :: ConnectionStats -> Int
-sentBytes = messageSize . outcomingFlow
+sentBytes = byteLength . outcomingFlow
 
 wastedBytes :: ConnectionStats -> Int
-wastedBytes (ConnectionStats _in out) = overhead _in + overhead out
+wastedBytes (ConnectionStats _in out)  = overhead _in + overhead out
 
 payloadBytes :: ConnectionStats -> Int
 payloadBytes (ConnectionStats _in out) = payload _in + payload out
@@ -184,11 +175,11 @@ payloadBytes (ConnectionStats _in out) = payload _in + payload out
 
 data Connection = Connection
   { connCaps         :: !Caps
-  , connExtCaps      :: !(IORef ExtendedCaps)
   , connTopic        :: !InfoHash
   , connRemotePeerId :: !PeerId
   , connThisPeerId   :: !PeerId
   , connStats        :: !(IORef ConnectionStats)
+  , connExtCaps      :: !(IORef ExtendedCaps)
   }
 
 instance Pretty Connection where
@@ -278,7 +269,7 @@ askStats :: (ConnectionStats -> a) -> Wire a
 askStats f = f <$> getStats
 
 putStats :: ChannelSide -> Message -> Wire ()
-putStats side msg = modifyRef connStats (addStats side (messageStats msg))
+putStats side msg = modifyRef connStats (addStats side (stats msg))
 
 
 getConnection :: Wire Connection
@@ -301,8 +292,8 @@ validateBoth action = do
   action
   validate ThisPeer
 
-keepStats :: Wire ()
-keepStats = do
+trackStats :: Wire ()
+trackStats = do
   mmsg <- await
   case mmsg of
     Nothing  -> return ()
@@ -329,7 +320,7 @@ extendedHandshake caps = do
   sendMessage $ nullExtendedHandshake caps
   msg <- recvMessage
   case msg of
-    Extended (EHandshake ExtendedHandshake {..}) ->
+    Extended (EHandshake ExtendedHandshake {..}) -> do
       setExtCaps $ ehsCaps <> caps
     _ -> protocolError HandshakeRefused
 
@@ -356,10 +347,10 @@ connectWire hs addr extCaps wire =
     statsRef   <- newIORef def
     runWire wire' sock $ Connection
       { connCaps         = caps
-      , connExtCaps      = extCapsRef
       , connTopic        = hsInfoHash hs
       , connRemotePeerId = hsPeerId   hs'
       , connThisPeerId   = hsPeerId   hs
+      , connExtCaps      = extCapsRef
       , connStats        = statsRef
       }
 
diff --git a/tests/Data/Torrent/BitfieldSpec.hs b/tests/Data/Torrent/BitfieldSpec.hs
new file mode 100644
index 00000000..093f6f19
--- /dev/null
+++ b/tests/Data/Torrent/BitfieldSpec.hs
@@ -0,0 +1,13 @@
+{-# OPTIONS -fno-warn-orphans #-}
+module Data.Torrent.BitfieldSpec (spec) where
+import Control.Applicative
+import Test.Hspec
+import Test.QuickCheck
+
+import Data.Torrent.Bitfield
+
+instance Arbitrary Bitfield where
+  arbitrary = fromBitmap <$> arbitrary
+
+spec :: Spec
+spec = return ()
\ No newline at end of file
diff --git a/tests/Network/BitTorrent/Exchange/BlockSpec.hs b/tests/Network/BitTorrent/Exchange/BlockSpec.hs
new file mode 100644
index 00000000..0712a21d
--- /dev/null
+++ b/tests/Network/BitTorrent/Exchange/BlockSpec.hs
@@ -0,0 +1,16 @@
+module Network.BitTorrent.Exchange.BlockSpec (spec) where
+import Control.Applicative
+import Test.Hspec
+import Test.QuickCheck
+
+import Network.BitTorrent.Exchange.Block
+
+
+instance Arbitrary a => Arbitrary (Block a) where
+  arbitrary = Block <$> arbitrary <*> arbitrary <*> arbitrary
+
+instance Arbitrary BlockIx where
+  arbitrary = BlockIx <$> arbitrary <*> arbitrary <*> arbitrary
+
+spec :: Spec
+spec = return ()
\ No newline at end of file
diff --git a/tests/Network/BitTorrent/Exchange/MessageSpec.hs b/tests/Network/BitTorrent/Exchange/MessageSpec.hs
index 38a20112..5d332eaa 100644
--- a/tests/Network/BitTorrent/Exchange/MessageSpec.hs
+++ b/tests/Network/BitTorrent/Exchange/MessageSpec.hs
@@ -2,6 +2,7 @@ module Network.BitTorrent.Exchange.MessageSpec (spec) where
 import Control.Applicative
 import Control.Exception
 import Data.ByteString as BS
+import Data.ByteString.Lazy as BL
 import Data.Default
 import Data.List as L
 import Data.Set as S
@@ -10,9 +11,11 @@ import Data.String
 import Test.Hspec
 import Test.QuickCheck
 
+import Data.Torrent.BitfieldSpec ()
 import Data.Torrent.InfoHashSpec ()
 import Network.BitTorrent.CoreSpec ()
 import Network.BitTorrent.Core
+import Network.BitTorrent.Exchange.BlockSpec ()
 import Network.BitTorrent.Exchange.Message
 
 instance Arbitrary Extension where
@@ -28,6 +31,45 @@ instance Arbitrary Handshake where
   arbitrary = Handshake <$> arbitrary <*> arbitrary
                         <*> arbitrary <*> arbitrary
 
+instance Arbitrary StatusUpdate where
+  arbitrary = frequency
+    [ (1, Choking    <$> arbitrary)
+    , (1, Interested <$> arbitrary)
+    ]
+
+instance Arbitrary Available where
+  arbitrary = frequency
+    [ (1, Have     <$> arbitrary)
+    , (1, Bitfield <$> arbitrary)
+    ]
+
+instance Arbitrary Transfer where
+  arbitrary = frequency
+    [ (1, Request <$> arbitrary)
+    , (1, Piece   <$> arbitrary)
+    , (1, Cancel  <$> arbitrary)
+    ]
+
+instance Arbitrary FastMessage where
+  arbitrary = frequency
+    [ (1, pure HaveAll)
+    , (1, pure HaveNone)
+    , (1, SuggestPiece  <$> arbitrary)
+    , (1, RejectRequest <$> arbitrary)
+    , (1, AllowedFast   <$> arbitrary)
+    ]
+
+instance Arbitrary Message where
+  arbitrary = frequency
+    [ (1, pure KeepAlive)
+    , (1, Status    <$> arbitrary)
+    , (1, Available <$> arbitrary)
+    , (1, Transfer  <$> arbitrary)
+    , (1, Fast      <$> arbitrary)
+    ]
+
+-- TODO test extension protocol
+
 spec :: Spec
 spec = do
   describe "Caps" $ do
@@ -38,6 +80,11 @@ spec = do
       S.fromList (fromCaps (toCaps (S.toList extSet) :: Caps))
         `shouldBe` extSet
 
+  describe "ByteStats" $ do
+    it "preserve size" $ property $ \ msg ->
+      byteLength (stats msg) `shouldBe`
+        fromIntegral (BS.length (S.encode (msg :: Message)))
+
   describe "ProtocolString" $ do
     it "fail to construct invalid string" $ do
       let str = L.replicate 500 'x'