convert bitfield from a peer into piece list

[functorrent.git] / src / FuncTorrent / Peer.hs

diff --git a/src/FuncTorrent/Peer.hs b/src/FuncTorrent/Peer.hs
index c2268e894c94b3275a025f1742971745849f9dda..93e4a1510bd841be2ae43f96fadb3d7b64fc98c4 100644 (file)
--- a/src/FuncTorrent/Peer.hs
+++ b/src/FuncTorrent/Peer.hs
@@ -1,20 +1,24 @@
 {-# LANGUAGE OverloadedStrings #-}
 module FuncTorrent.Peer
     (Peer(..),
-     handShake
+     handShake,
+     msgLoop
     ) where
 
-import Prelude hiding (lookup, concat, replicate, splitAt)
+import Prelude hiding (lookup, concat, replicate, splitAt, empty)
 
 import System.IO
 import Data.ByteString (ByteString, pack, unpack, concat, hGet, hPut, singleton)
-import qualified Data.ByteString.Char8 as BC (replicate, pack)
+import Data.ByteString.Lazy (fromStrict, fromChunks)
+import qualified Data.ByteString.Char8 as BC (replicate, pack, length)
 import Network (connectTo, PortID(..))
-import Data.Binary (Binary(..))
+import Data.Binary (Binary(..), decode)
 import Data.Binary.Put (putWord32be, putWord16be, putWord8)
-import Data.Binary.Get (getWord32be, getWord16be, getWord8)
-import Control.Monad (replicateM, liftM)
+import Data.Binary.Get (getWord32be, getWord16be, getWord8, runGet)
+import Control.Monad (replicateM, liftM, forever)
 import Control.Applicative ((<$>), liftA3)
+import Data.Bits
+import Data.Word (Word8)
 
 type ID = String
 type IP = String
@@ -24,7 +28,7 @@ data PeerState = PeerState { handle :: Handle
                            , am_choking :: Bool
                            , am_interested :: Bool
                            , peer_choking :: Bool
-                           , peer_interested :: Bool }
+                           , peer_interested :: Bool}
 
 -- Maintain info on every piece and the current state of it.
 -- should probably be a TVar.
@@ -69,9 +73,8 @@ handShake (Peer _ ip port) infoHash peerid = do
   h <- connectTo ip (PortNumber (fromIntegral port))
   hSetBuffering h LineBuffering
   hPut h hs
-  rlenBS <- hGet h 1
-  let rlen = fromIntegral $ (unpack rlenBS) !! 0
-  hGet h rlen
+  rlenBS <- hGet h (length (unpack hs))
+  putStrLn $ "got handshake from peer: " ++ show rlenBS
   return h
 
 instance Binary PeerMsg where
@@ -130,7 +133,41 @@ instance Binary PeerMsg where
      8 -> liftA3 CancelMsg getInteger getInteger getInteger
        where getInteger = fromIntegral <$> getWord32be
      9 -> liftM (PortMsg . fromIntegral) getWord16be
-     _ -> error "unknown message ID"
+     _ -> error ("unknown message ID: " ++ show msgid)
+
+getMsg :: Handle -> IO PeerMsg
+getMsg h = do
+  lBS <- hGet h 4
+  let l = bsToInt lBS
+  if l == 0
+    then return KeepAliveMsg
+    else do
+    msg <- hGet h l
+    return $ decode $ fromStrict $ concat [lBS, msg]
+
+bsToInt :: ByteString -> Int
+bsToInt x = fromIntegral (runGet getWord32be (fromChunks (return x)))
+
+bitfieldToList :: [Word8] -> Integer -> [Integer]
+bitfieldToList [] pos = []
+bitfieldToList (b:bs) pos =
+  let setBits = [pos*8 + (toInteger i) | i <- [0..8], testBit b i]
+  in
+   setBits ++ (bitfieldToList bs (pos + 1))
 
 -- loop1 :: shake hands with all peers, find out the pieces they have, form PieceData.
 -- recvMsg :: Peer -> Handle -> Msg
+
+msgLoop :: Handle -> ByteString -> IO ()
+msgLoop h pieceHash =
+  let numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
+  in
+   forever $ do
+     msg <- getMsg h
+     putStrLn $ "got a " ++ show msg
+     case msg of
+      BitFieldMsg bss -> do
+                         let pieceList = bitfieldToList (unpack bss) 0
+                         putStrLn (show pieceList)
+                         -- doenload each of the piece
+      _ -> putStrLn (show msg)