6 files changed, 126 insertions, 102 deletions

diff --git a/bittorrent.cabal b/bittorrent.cabal
index 51da7d4f..acb28fa7 100644
--- a/bittorrent.cabal
+++ b/bittorrent.cabal
@@ -59,9 +59,10 @@ library
 --                   , Network.BitTorrent.DHT.Protocol
 --                   , Network.BitTorrent.DHT.Session
 --                   , Network.BitTorrent.Exchange
---                   , Network.BitTorrent.Exchange.Message
+                     , Network.BitTorrent.Exchange.Assembler
+                     , Network.BitTorrent.Exchange.Message
 --                   , Network.BitTorrent.Exchange.Session
---                   , Network.BitTorrent.Exchange.Status
+--                     , Network.BitTorrent.Exchange.Status
 --                   , Network.BitTorrent.Exchange.Wire
 --                   , Network.BitTorrent.Extension
 --                   , Network.BitTorrent.Tracker
diff --git a/src/Data/Torrent/Bitfield.hs b/src/Data/Torrent/Bitfield.hs
index a87a7ccf..c7ef0998 100644
--- a/src/Data/Torrent/Bitfield.hs
+++ b/src/Data/Torrent/Bitfield.hs
@@ -190,6 +190,7 @@ isSubsetOf a b = bfSet a `S.isSubsetOf` bfSet b
 -- download.
 type Frequency = Int
 
+-- TODO rename to availability
 -- | How many times each piece index occur in the given bitfield set.
 frequencies :: [Bitfield] -> Vector Frequency
 frequencies [] = V.fromList []
diff --git a/src/Network/BitTorrent/Exchange/Bus.hs b/src/Network/BitTorrent/Exchange/Bus.hs
index 4800c4a0..7de91180 100644
--- a/src/Network/BitTorrent/Exchange/Bus.hs
+++ b/src/Network/BitTorrent/Exchange/Bus.hs
@@ -34,6 +34,9 @@ flushPending = {-# SCC flushPending #-} do
     P2P monad
 -----------------------------------------------------------------------}
 
+filterMeaninless :: P2P Message Message
+filterMeaninless = undefined
+
 -- |
 --   Exceptions:
 --
diff --git a/src/Network/BitTorrent/Exchange/Extension.hs b/src/Network/BitTorrent/Exchange/Extension.hs
deleted file mode 100644
index e81cdb87..00000000
--- a/src/Network/BitTorrent/Exchange/Extension.hs
+++ /dev/null
@@ -1,70 +0,0 @@
--- |
---   Copyright   :  (c) Sam Truzjan 2013
---   License     :  BSD3
---   Maintainer  :  pxqr.sta@gmail.com
---   Stability   :  experimental
---   Portability :  portable
---
---   This module provides peer capabilities detection.
---
---   See <http://www.bittorrent.org/beps/bep_0004.html> for more
---   information.
---
-module Network.BitTorrent.Exchange.Extension
-       ( -- * Capabilities
-         Caps
-
-         -- * Extensions
-       , Extension(..)
-       ) where
-
-import Data.Bits
-import Data.Default
-import Data.Monoid
-import Data.Word
-import Text.PrettyPrint
-import Text.PrettyPrint.Class
-
-class (Enum a, Bounded a) => Capability a where
-  capMask :: a -> Word64
-  capRequires :: a -> Word64
-
-newtype Caps a = Caps Word64
-
-instance (Pretty a, Capability a) => Pretty (Caps a) where
-  pretty = hcat . punctuate ", " . map pretty . toList
-
-instance Default (Caps a) where
-  def = Caps 0
-  {-# INLINE def #-}
-
-instance Monoid (Caps a) where
-  mempty  = Caps (-1)
-  {-# INLINE mempty #-}
-
-  mappend (Caps a) (Caps b) = Caps (a .&. b)
-  {-# INLINE mappend #-}
-
-allowed :: Capability a => a -> Caps a -> Bool
-allowed = member
-fromList :: Capability a => [a] -> Caps a
-fromList = Caps . foldr (.&.) 0 . map capMask
-
-toList :: Capability a => Caps a -> [a]
-toList (Caps rb) = filter (testMask rb . capMask) [minBound..maxBound]
-  where
-    testMask bits x = bits .&. x > 0
-
-
-data Extension
-  = ExtDHT  -- ^ BEP 5
-  | ExtFast -- ^ BEP 6
-    deriving (Show, Eq, Ord, Enum, Bounded)
-
-instance Pretty Extension where
-  pretty ExtDHT  = "DHT"
-  pretty ExtFast = "Fast Extension"
-
-instance Capability Extension where
-  capMask ExtDHT  = 0x01
-  capMask ExtFast = 0x04
diff --git a/src/Network/BitTorrent/Exchange/Message.hs b/src/Network/BitTorrent/Exchange/Message.hs
index 546288b2..6f649030 100644
--- a/src/Network/BitTorrent/Exchange/Message.hs
+++ b/src/Network/BitTorrent/Exchange/Message.hs
@@ -28,19 +28,25 @@
 {-# LANGUAGE TemplateHaskell   #-}
 {-# OPTIONS  -fno-warn-orphans #-}
 module Network.BitTorrent.Exchange.Message
-       ( -- * Initial handshake
-         Handshake(..)
-       , handshake
-       , handshakeCaps
-       , recvHandshake
-       , sendHandshake
-
-         -- ** Defaults
+       ( -- * Extensions
+         Extension (..)
+       , Caps
+       , requires
+       , allowed
+       , toCaps
+       , fromCaps
+
+         -- * Handshake
+       , Handshake(..)
        , defaultHandshake
        , defaultBTProtocol
-       , defaultReserved
        , handshakeMaxSize
 
+         -- * TODO remove this section from this module
+       , handshake
+       , recvHandshake
+       , sendHandshake
+
          -- * Messages
        , Message        (..)
        , StatusUpdate   (..)
@@ -51,10 +57,13 @@ module Network.BitTorrent.Exchange.Message
 import Control.Applicative
 import Control.Exception
 import Control.Monad
+import Data.Bits
 import Data.ByteString as BS
 import Data.ByteString.Char8 as BC
 import Data.ByteString.Lazy  as BL
 import Data.Default
+import Data.List as L
+import Data.Monoid
 import Data.Serialize as S
 import Data.Word
 import Network
@@ -67,7 +76,66 @@ import Data.Torrent.Block
 import Data.Torrent.InfoHash
 import Network.BitTorrent.Core.PeerId
 import Network.BitTorrent.Core.PeerAddr ()
-import Network.BitTorrent.Exchange.Extension
+
+{-----------------------------------------------------------------------
+--  Extensions
+-----------------------------------------------------------------------}
+
+--  | See <http://www.bittorrent.org/beps/bep_0004.html> for more
+--  information.
+--
+data Extension
+  = ExtDHT  -- ^ BEP 5
+  | ExtFast -- ^ BEP 6
+    deriving (Show, Eq, Ord, Enum, Bounded)
+
+instance Pretty Extension where
+  pretty ExtDHT  = "DHT"
+  pretty ExtFast = "Fast Extension"
+
+capMask :: Extension -> Caps
+capMask ExtDHT  = Caps 0x01
+capMask ExtFast = Caps 0x04
+
+{-----------------------------------------------------------------------
+--  Capabilities
+-----------------------------------------------------------------------}
+
+-- | A set of 'Extension's.
+newtype Caps = Caps { unCaps :: Word64 }
+  deriving (Show, Eq)
+
+instance Pretty Caps where
+  pretty = hcat . punctuate ", " . L.map pretty . fromCaps
+
+instance Default Caps where
+  def = Caps 0
+  {-# INLINE def #-}
+
+instance Monoid Caps where
+  mempty  = Caps (-1)
+  {-# INLINE mempty #-}
+
+  mappend (Caps a) (Caps b) = Caps (a .&. b)
+  {-# INLINE mappend #-}
+
+instance Serialize Caps where
+  put (Caps caps) = S.putWord64be caps
+  {-# INLINE put #-}
+
+  get = Caps <$> S.getWord64be
+  {-# INLINE get #-}
+
+allowed :: Caps -> Extension -> Bool
+allowed (Caps caps) = testMask . capMask
+  where
+    testMask (Caps bits) = (bits .&. caps) == bits
+
+toCaps :: [Extension] -> Caps
+toCaps = Caps . L.foldr (.|.) 0 . L.map (unCaps . capMask)
+
+fromCaps :: Caps -> [Extension]
+fromCaps caps = L.filter (allowed caps) [minBound..maxBound]
 
 {-----------------------------------------------------------------------
     Handshake
@@ -77,11 +145,11 @@ import Network.BitTorrent.Exchange.Extension
 -- to establish connection between peers.
 --
 data Handshake = Handshake {
-    -- | Identifier of the protocol.
+    -- | Identifier of the protocol. This is usually equal to defaultProtocol
     hsProtocol    :: BS.ByteString
 
     -- | Reserved bytes used to specify supported BEP's.
-  , hsReserved    :: Capabilities
+  , hsReserved    :: Caps
 
     -- | Info hash of the info part of the metainfo file. that is
     -- transmitted in tracker requests. Info hash of the initiator
@@ -98,17 +166,17 @@ data Handshake = Handshake {
   } deriving (Show, Eq)
 
 instance Serialize Handshake where
-  put hs = do
-    S.putWord8 (fromIntegral (BS.length (hsProtocol hs)))
-    S.putByteString (hsProtocol hs)
-    S.putWord64be   (hsReserved hs)
-    S.put (hsInfoHash hs)
-    S.put (hsPeerId hs)
+  put Handshake {..} = do
+    S.putWord8 (fromIntegral (BS.length hsProtocol))
+    S.putByteString hsProtocol
+    S.put hsReserved
+    S.put hsInfoHash
+    S.put hsPeerId
 
   get = do
     len  <- S.getWord8
     Handshake <$> S.getBytes (fromIntegral len)
-              <*> S.getWord64be
+              <*> S.get
               <*> S.get
               <*> S.get
 
@@ -116,11 +184,6 @@ instance Pretty Handshake where
   pretty Handshake {..}
     = text (BC.unpack hsProtocol) <+> pretty (clientInfo hsPeerId)
 
--- | Extract capabilities from a peer handshake message.
-handshakeCaps :: Handshake -> Capabilities
-handshakeCaps = hsReserved
-
-
 -- | Get handshake message size in bytes from the length of protocol
 -- string.
 handshakeSize :: Word8 -> Int
@@ -128,21 +191,18 @@ handshakeSize n = 1 + fromIntegral n + 8 + 20 + 20
 
 -- | Maximum size of handshake message in bytes.
 handshakeMaxSize :: Int
-handshakeMaxSize = handshakeSize 255
+handshakeMaxSize = handshakeSize maxBound
 
 -- | Default protocol string "BitTorrent protocol" as is.
 defaultBTProtocol :: BS.ByteString
 defaultBTProtocol = "BitTorrent protocol"
 
--- | Default reserved word is 0.
-defaultReserved :: Word64
-defaultReserved = 0
-
 -- | Length of info hash and peer id is unchecked, so it /should/ be
 -- equal 20.
 defaultHandshake :: InfoHash -> PeerId -> Handshake
-defaultHandshake = Handshake defaultBTProtocol defaultReserved
+defaultHandshake = Handshake defaultBTProtocol def
 
+-- | TODO remove socket stuff to corresponding module
 sendHandshake :: Socket -> Handshake -> IO ()
 sendHandshake sock hs = sendAll sock (S.encode hs)
 
@@ -206,6 +266,14 @@ data RegularMessage =
   | Cancel  !BlockIx
     deriving (Show, Eq)
 
+-- TODO
+-- data Availability = Have | Bitfield
+-- data Transfer
+--   = Request !BlockIx
+--   | Piece   !(Block BL.ByteString)
+--   | Cancel  !BlockIx
+
+
 instance Pretty RegularMessage where
   pretty (Have     ix ) = "Have"     <+> int ix
   pretty (Bitfield _  ) = "Bitfield"
@@ -349,3 +417,11 @@ putFast  HaveNone          = putInt 1  >> S.putWord8 0x0F
 putFast (SuggestPiece pix) = putInt 5  >> S.putWord8 0x0D >> putInt pix
 putFast (RejectRequest i ) = putInt 13 >> S.putWord8 0x10 >> S.put i
 putFast (AllowedFast   i ) = putInt 5  >> S.putWord8 0x11 >> putInt i
+
+
+requires :: Message -> Maybe Extension
+requires  KeepAlive  = Nothing
+requires (Status  _) = Nothing
+requires (Regular _) = Nothing
+requires (Port    _) = Just ExtDHT
+requires (Fast    _) = Just ExtFast
\ No newline at end of file
diff --git a/src/Network/BitTorrent/Exchange/Status.hs b/src/Network/BitTorrent/Exchange/Status.hs
index 683ac594..7920f2a1 100644
--- a/src/Network/BitTorrent/Exchange/Status.hs
+++ b/src/Network/BitTorrent/Exchange/Status.hs
@@ -4,6 +4,7 @@ module Network.BitTorrent.Exchange.Status
          PeerStatus(..)
        , choking
        , interested
+       , updateStatus
 
          -- * Session status
        , SessionStatus(..)
@@ -24,6 +25,9 @@ import Data.Aeson.TH
 import Data.List as L
 import Data.Default
 
+import Network.BitTorrent.Exchange.Message
+
+
 -- |
 data PeerStatus = PeerStatus {
     _choking    :: !Bool
@@ -36,6 +40,15 @@ $(deriveJSON L.tail ''PeerStatus)
 instance Default PeerStatus where
   def = PeerStatus True False
 
+updateStatus :: StatusUpdate -> PeerStatus -> PeerStatus
+updateStatus Choke         = choking    .~ True
+updateStatus Unchoke       = choking    .~ False
+updateStatus Interested    = interested .~ True
+updateStatus NotInterested = interested .~ False
+
+statusUpdates :: PeerStatus -> PeerStatus -> [StatusUpdate]
+statusUpdates a b = undefined
+
 -- |
 data SessionStatus = SessionStatus {
     _clientStatus :: !PeerStatus