keep hlint happy

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

diff --git a/src/FuncTorrent/Peer.hs b/src/FuncTorrent/Peer.hs
index 8007ab057432be2b793a49c440b6777619322dfc..54761211d38eaa98fe1c7ee58bd97465df7a8ac7 100644 (file)
--- a/src/FuncTorrent/Peer.hs
+++ b/src/FuncTorrent/Peer.hs
@@ -1,45 +1,54 @@
 {-# LANGUAGE OverloadedStrings #-}
 module FuncTorrent.Peer
     (Peer(..),
-     handShake,
-     msgLoop
+     handlePeerMsgs
     ) where
 
-import Prelude hiding (lookup, concat, replicate, splitAt)
+import Prelude hiding (lookup, concat, replicate, splitAt, empty, writeFile)
 
-import System.IO
-import Data.ByteString (ByteString, pack, unpack, concat, hGet, hPut, singleton)
+import System.IO (Handle, BufferMode(..), hSetBuffering)
+import Data.ByteString (ByteString, pack, unpack, concat, hGet, hPut, singleton, writeFile)
 import Data.ByteString.Lazy (fromStrict, fromChunks)
-import qualified Data.ByteString.Char8 as BC (replicate, pack)
+import qualified Data.ByteString.Char8 as BC (replicate, pack, length)
 import Network (connectTo, PortID(..))
 import Data.Binary (Binary(..), decode)
 import Data.Binary.Put (putWord32be, putWord16be, putWord8)
 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)
+import Data.Map (Map(..), fromList, (!))
+
+import FuncTorrent.Metainfo (Info(..), Metainfo(..))
+import FuncTorrent.Utils (splitN)
 
 type ID = String
 type IP = String
 type Port = Integer
 
 data PeerState = PeerState { handle :: Handle
-                           , am_choking :: Bool
-                           , am_interested :: Bool
-                           , peer_choking :: Bool
-                           , peer_interested :: Bool}
+                           , amChoking :: Bool
+                           , amInterested :: Bool
+                           , peerChoking :: Bool
+                           , peerInterested :: Bool}
 
 -- Maintain info on every piece and the current state of it.
 -- should probably be a TVar.
 type Pieces = [PieceData]
 
-data PieceState = Pending
-                | InProgress
-                | Have
-                deriving (Show)
+data PieceDlState = Pending
+                  | InProgress
+                  | Have
+                  deriving (Show, Eq)
+
+-- todo - map with index to a new data structure (peers who have that piece amd state)
+data PieceData = PieceData { peers :: [Peer]        -- ^ list of peers who have this piece
+                           , state :: PieceDlState  -- ^ state of the piece from download perspective.
+                           , hash  :: ByteString }      -- ^ piece hash
 
-data PieceData = PieceData { index :: Int           -- ^ Piece number
-                           , peers :: [Peer]        -- ^ list of peers who have this piece
-                           , state :: PieceState }  -- ^ state of the piece from download perspective.
+-- which piece is with which peers
+type PieceMap = Map Integer PieceData
 
 -- | Peer is a PeerID, IP address, port tuple
 data Peer = Peer ID IP Port
@@ -58,6 +67,19 @@ data PeerMsg = KeepAliveMsg
              | PortMsg Port
              deriving (Show)
 
+-- Make the initial Piece map, with the assumption that no peer has the
+-- piece and that every piece is pending download.
+mkPieceMap :: Integer -> ByteString -> PieceMap
+mkPieceMap numPieces pieceHash = fromList kvs
+  where kvs = [(i, PieceData { peers = []
+                             , state = Pending
+                             , hash = h }) | (i, h) <- zip [0..numPieces] hashes]
+        hashes = splitN (fromIntegral numPieces) pieceHash
+
+havePiece :: PieceMap -> Integer -> Bool
+havePiece pm index =
+  state (pm ! index) == Have
+
 genHandShakeMsg :: ByteString -> String -> ByteString
 genHandShakeMsg infoHash peer_id = concat [pstrlen, pstr, reserved, infoHash, peerID]
   where pstrlen = singleton 19
@@ -143,13 +165,44 @@ getMsg h = 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]
+bitfieldToList bs = go bs 0
+  where go [] _ = []
+        go (b:bs') pos =
+          let setBits = [pos*8 + toInteger i | i <- [0..8], testBit b i]
+          in
+           setBits ++ go bs' (pos + 1)
+
+-- downloadPiece :: Integer -> Handle -> IO ()
+
+createDummyFile :: FilePath -> Int -> IO ()
+createDummyFile path size =
+  writeFile path (BC.replicate size '\0')
+
 -- loop1 :: shake hands with all peers, find out the pieces they have, form PieceData.
 -- recvMsg :: Peer -> Handle -> Msg
-
-msgLoop :: Handle -> IO ()
-msgLoop h = forever $ do
-  msg <- getMsg h
-  putStrLn $ "got a " ++ show msg
+msgLoop :: Handle -> ByteString -> IO ()
+msgLoop h pieceHash =
+  let numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
+      pieceStatus = mkPieceMap numPieces pieceHash
+  in
+   forever $ do
+     msg <- getMsg h
+     putStrLn $ "got a " ++ show msg
+     case msg of
+      BitFieldMsg bss -> do
+        let pieceList = bitfieldToList (unpack bss)
+        print pieceList
+        -- download each of the piece in order
+      _ -> print msg
+
+handlePeerMsgs :: Peer -> Metainfo -> String -> (String -> IO ()) -> IO ()
+handlePeerMsgs p m peerId logFn = do
+  h <- handShake p (infoHash m) peerId
+  logFn "handShake"
+  msgLoop h (pieces (info m))
+