~ Fix bitfield encoding.

5 files changed, 50 insertions, 24 deletions

diff --git a/bittorrent.cabal b/bittorrent.cabal
index da69e48d..1a28c880 100644
--- a/bittorrent.cabal
+++ b/bittorrent.cabal
@@ -59,7 +59,8 @@ library
 
                        -- Data packages
                      , array      >= 0.4
-                     , bytestring >= 0.10.2
+                     , bytestring
+ -- >= 0.10.2
                      , containers >= 0.4
                      , intset     >= 0.1
                      , text       >= 0.11.0
@@ -94,11 +95,12 @@ library
 
 test-suite info-hash
   type:                exitcode-stdio-1.0
-  main-is:             info-hash.hs
+  main-is:             InfoHash.hs
   hs-source-dirs:      tests
 
   build-depends:       base       == 4.*
-                     , bytestring >= 0.10.2.0
+                     , bytestring
+ -- >= 0.10.2.0
                      , containers >= 0.4.2.1
                      , bencoding  >= 0.1.0.0
                      , bittorrent
@@ -116,7 +118,8 @@ test-suite properties
   main-is:             Main.hs
   hs-source-dirs:      tests
   build-depends:       base       == 4.*
-                     , bytestring >= 0.10.2
+                     , bytestring
+-- >= 0.10.2
                      , cereal     >= 0.3.5.2
                      , network    >= 2.4.0.13
                      , text
@@ -141,7 +144,7 @@ benchmark benchmarks
   hs-source-dirs:      bench
 
   build-depends:       base == 4.*
-                     , bytestring >= 0.10.2.0
+                     , bytestring
                      , cereal
                      , network
 
diff --git a/src/Data/Bitfield.hs b/src/Data/Bitfield.hs
index ee0570e7..56365bf7 100644
--- a/src/Data/Bitfield.hs
+++ b/src/Data/Bitfield.hs
@@ -17,6 +17,7 @@ module Data.Bitfield
 
          -- * Construction
        , haveAll, haveNone, have
+       , adjustSize
 
          -- * Query
        , Data.Bitfield.null
@@ -31,8 +32,7 @@ module Data.Bitfield
        , difference
 
          -- * Serialization
-       , getBitfield, putBitfield
-       , bitfieldByteCount
+       , fromBitmap, toBitmap
 
 #if  defined (TESTING)
          -- * Debug
@@ -42,15 +42,18 @@ module Data.Bitfield
 
 import Control.Monad
 import Control.Monad.ST
+import           Data.ByteString (ByteString)
+import qualified Data.ByteString as B
+import qualified Data.ByteString.Lazy as Lazy
 import           Data.Vector.Unboxed (Vector)
 import qualified Data.Vector.Unboxed as V
 import qualified Data.Vector.Unboxed.Mutable as VM
 import           Data.IntervalSet (IntSet)
 import qualified Data.IntervalSet as S
+import qualified  Data.IntervalSet.ByteString as S
 import           Data.List (foldl')
 import           Data.Monoid
 import           Data.Ratio
-import           Data.Serialize
 import Network.BitTorrent.PeerWire.Block
 
 
@@ -97,6 +100,11 @@ have ix Bitfield {..}
   | 0 <= ix && ix < bfSize = Bitfield bfSize (S.insert ix bfSet)
   |      otherwise         = Bitfield bfSize bfSet
 
+-- | Assign new size to bitfield. FIXME Normally, size should be only
+-- decreased, otherwise exception raised.
+adjustSize :: PieceCount -> Bitfield -> Bitfield
+adjustSize s Bitfield {..} = Bitfield s bfSet
+
 {-----------------------------------------------------------------------
     Query
 -----------------------------------------------------------------------}
@@ -200,17 +208,22 @@ unions = foldl' union (haveNone 0)
     Serialization
 -----------------------------------------------------------------------}
 
--- |
-getBitfield :: Int -> Get Bitfield
-getBitfield = error "getBitfield"
-
--- |
-putBitfield :: Bitfield -> Put
-putBitfield = error "putBitfield"
+-- | Unpack 'Bitfield' from tightly packed bit array. Note resulting
+-- size might be more than real bitfield size, use 'adjustSize'.
+fromBitmap :: ByteString -> Bitfield
+fromBitmap bs = Bitfield {
+    bfSize = B.length bs * 8
+  , bfSet  = S.fromByteString bs
+  }
+{-# INLINE fromBitmap #-}
 
--- |
-bitfieldByteCount :: Bitfield -> Int
-bitfieldByteCount = error "bitfieldByteCount"
+-- | Pack a 'Bitfield' to tightly packed bit array.
+toBitmap :: Bitfield -> Lazy.ByteString
+toBitmap Bitfield {..} = Lazy.fromChunks [intsetBM, alignment]
+  where
+    byteSize  = bfSize `div` 8 + if bfSize `mod` 8 == 0 then 0 else 1
+    alignment = B.replicate (byteSize - B.length intsetBM) 0
+    intsetBM  = S.toByteString bfSet
 
 {-----------------------------------------------------------------------
     Debug
diff --git a/src/Network/BitTorrent/PeerWire/Message.hs b/src/Network/BitTorrent/PeerWire/Message.hs
index 6515fdf2..3895ed5f 100644
--- a/src/Network/BitTorrent/PeerWire/Message.hs
+++ b/src/Network/BitTorrent/PeerWire/Message.hs
@@ -7,6 +7,7 @@ module Network.BitTorrent.PeerWire.Message
 
 import Control.Applicative
 import qualified Data.ByteString as B
+import qualified Data.ByteString.Lazy as Lazy
 import Data.Serialize
 import Text.PrettyPrint
 import Network
@@ -91,7 +92,7 @@ instance Serialize Message where
           0x02 -> return Interested
           0x03 -> return NotInterested
           0x04 -> Have     <$> getInt
-          0x05 -> Bitfield <$> getBitfield (pred len)
+          0x05 -> (Bitfield . fromBitmap) <$> getByteString (pred len)
           0x06 -> Request  <$> get
           0x07 -> Piece    <$> getBlock (len - 9)
           0x08 -> Cancel   <$> get
@@ -118,8 +119,9 @@ instance Serialize Message where
   put Interested    = putInt 1  >> putWord8 0x02
   put NotInterested = putInt 1  >> putWord8 0x03
   put (Have i)      = putInt 5  >> putWord8 0x04 >> putInt i
-  put (Bitfield b)  = putInt l  >> putWord8 0x05 >> putBitfield b
-    where l = succ (bitfieldByteCount b)
+  put (Bitfield bf) = putInt l  >> putWord8 0x05 >> putLazyByteString b
+    where b = toBitmap bf
+          l = succ (fromIntegral (Lazy.length b))
           {-# INLINE l #-}
   put (Request blk) = putInt 13 >> putWord8 0x06 >> put blk
   put (Piece   blk) = putInt l  >> putWord8 0x07 >> putBlock
diff --git a/src/Network/BitTorrent/Tracker.hs b/src/Network/BitTorrent/Tracker.hs
index 2e599002..aaa08f3c 100644
--- a/src/Network/BitTorrent/Tracker.hs
+++ b/src/Network/BitTorrent/Tracker.hs
@@ -145,7 +145,7 @@ data TSession = TSession {
 newSession :: Progress -> Int -> [PeerAddr] -> IO TSession
 newSession pr i ps = TSession <$> newTVarIO pr
                               <*> newIORef i
-                              <*> newTVarIO psx
+                              <*> newTVarIO ps
 
 getPeerList :: TSession -> IO [PeerAddr]
 getPeerList = readTVarIO . sePeers
diff --git a/tests/Main.hs b/tests/Main.hs
index 0e18a06b..84870d66 100644
--- a/tests/Main.hs
+++ b/tests/Main.hs
@@ -185,6 +185,15 @@ instance Arbitrary Message where
     , Cancel   <$> arbitrary
     , Port     <$> arbitrary
     ]
+-- todo add all messages
+
+prop_messageEncoding :: Message -> Bool
+prop_messageEncoding msg @ (Bitfield bf)
+  = case S.decode (S.encode msg) of
+      Right (Bitfield bf') -> bf == adjustSize (totalCount bf) bf'
+      Left  _   -> False
+prop_messageEncoding msg
+  = S.decode (S.encode msg) == Right msg
 
 {-----------------------------------------------------------------------
     Main
@@ -209,7 +218,6 @@ main = defaultMain $
     -- handshake module
   , testProperty "handshake encoding" $
       prop_cerealEncoding  (T :: T Handshake)
-  , testProperty "message encoding" $
-      prop_cerealEncoding (T :: T Message)
+  , testProperty "message encoding" prop_messageEncoding
 
   ] ++ test_scrape_url