]> git.rkrishnan.org Git - functorrent.git/blobdiff - src/FuncTorrent/Peer.hs
verify SHA1 hash of each piece after assembling the blocks
[functorrent.git] / src / FuncTorrent / Peer.hs
index 33d72a384047bc9c1b0a37a0875a416a8e1fc5f4..7802b08cb055a23d2f23ab14e8f2373e8a2e989a 100644 (file)
@@ -4,24 +4,25 @@ module FuncTorrent.Peer
      handlePeerMsgs
     ) where
 
-import Prelude hiding (lookup, concat, replicate, splitAt, writeFile)
+import Prelude hiding (lookup, concat, replicate, splitAt, writeFile, take)
 
-import System.IO (Handle, BufferMode(..), hSetBuffering)
-import Data.ByteString (ByteString, pack, unpack, concat, hGet, hPut, singleton, writeFile)
+import System.IO (Handle, BufferMode(..), IOMode(..), SeekMode(..), withFile, hSeek, hSetBuffering)
+import Data.ByteString (ByteString, pack, unpack, concat, hGet, hPut, singleton, writeFile, take)
 import Data.ByteString.Lazy (fromStrict, fromChunks, toStrict)
 import qualified Data.ByteString.Char8 as BC (replicate, pack, length)
 import Network (connectTo, PortID(..))
 import Data.Binary (Binary(..), decode, encode)
 import Data.Binary.Put (putWord32be, putWord16be, putWord8)
 import Data.Binary.Get (getWord32be, getWord16be, getWord8, runGet)
-import Control.Monad (replicateM, liftM, forever)
+import Control.Monad (replicateM, liftM, forM)
 import Control.Applicative ((<$>), liftA3)
 import Data.Bits
 import Data.Word (Word8)
-import Data.Map (Map(..), fromList, toList, (!), mapWithKey)
+import Data.Map (Map, fromList, toList, (!), mapWithKey, adjust)
+import qualified Crypto.Hash.SHA1 as SHA1 (hash)
 
 import FuncTorrent.Metainfo (Info(..), Metainfo(..))
-import FuncTorrent.Utils (splitN)
+import FuncTorrent.Utils (splitN, splitNum)
 
 type ID = String
 type IP = String
@@ -35,10 +36,6 @@ data PeerState = PeerState { handle :: Handle
                            , heChoking :: Bool
                            , heInterested :: Bool}
 
--- Maintain info on every piece and the current state of it.
--- should probably be a TVar.
-type Pieces = [PieceData]
-
 data PieceDlState = Pending
                   | InProgress
                   | Have
@@ -47,7 +44,8 @@ data PieceDlState = Pending
 -- 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
+                           , hash  :: ByteString    -- ^ piece hash
+                           , len :: Integer }       -- ^ piece length
 
 -- which piece is with which peers
 type PieceMap = Map Integer PieceData
@@ -71,11 +69,13 @@ data PeerMsg = KeepAliveMsg
 
 -- 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
+mkPieceMap :: Integer -> ByteString -> [Integer] -> PieceMap
+mkPieceMap numPieces pieceHash pLengths = fromList kvs
   where kvs = [(i, PieceData { peers = []
                              , state = Pending
-                             , hash = h }) | (i, h) <- zip [0..numPieces] hashes]
+                             , hash = h
+                             , len = pLen })
+              | (i, h, pLen) <- zip3 [0..numPieces] hashes pLengths]
         hashes = splitN (fromIntegral numPieces) pieceHash
 
 havePiece :: PieceMap -> Integer -> Bool
@@ -185,60 +185,68 @@ bitfieldToList bs = go bs 0
           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.
+-- write into a file at a specific offet
+writeFileAtOffset :: FilePath -> Integer -> ByteString -> IO ()
+writeFileAtOffset path offset block =
+  withFile path WriteMode $ (\h -> do
+                                _ <- hSeek h AbsoluteSeek offset
+                                hPut h block)
+
 -- recvMsg :: Peer -> Handle -> Msg
 msgLoop :: PeerState -> PieceMap -> IO ()
-msgLoop state pieceStatus | meInterested state == False &&
-                            heChoking state == True = do
-                              -- if meInterested and he NOT Choking, pick a piece to download
-                              -- and send a requestmsg.
-                              let h = handle state
+msgLoop pState pieceStatus | meInterested pState == False &&
+                            heChoking pState == True = do
+                              -- if me NOT Interested and she is Choking, tell her that
+                              -- I am interested.
+                              let h = handle pState
                               sendMsg h InterestedMsg
-                              putStrLn $ "--> InterestedMsg to peer: " ++ show (peer state)
-                              msgLoop (state { meInterested = True }) pieceStatus
-                          | meInterested state == True &&
-                            heChoking state == False =
+                              putStrLn $ "--> InterestedMsg to peer: " ++ show (peer pState)
+                              msgLoop (pState { meInterested = True }) pieceStatus
+                          | meInterested pState == True &&
+                            heChoking pState == False =
+                              -- if me Interested and she not Choking, send her a request
+                              -- for a piece.
                               case pickPiece pieceStatus of
                                Nothing -> putStrLn "Nothing to download"
                                Just workPiece -> do
-                                 sendMsg (handle state) (RequestMsg workPiece 0 16384)
-                                 putStrLn $ "--> RequestMsg for Piece " ++ (show workPiece) ++ "to peer: " ++ show (peer state)
-                                 msg <- getMsg (handle state)
-                                 putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer state)
-                                 -- msgLoop state pieceStatus
+                                 let pLen = len (pieceStatus ! workPiece)
+                                 pBS <- downloadPiece (handle pState) workPiece pLen
+                                 if not $ verifyHash pBS (hash (pieceStatus ! workPiece))
+                                   then
+                                   putStrLn $ "Hash mismatch: " ++ show (hash (pieceStatus ! workPiece)) ++ " vs " ++ show (take 20 (SHA1.hash pBS))
+                                   else
+                                   msgLoop pState (adjust (\pieceData -> pieceData { state = Have }) workPiece pieceStatus)
                           | otherwise = do
-                              msg <- getMsg (handle state)
-                              putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer state)
+                              msg <- getMsg (handle pState)
+                              putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer pState)
                               case msg of
                                KeepAliveMsg -> do
-                                 sendMsg (handle state) KeepAliveMsg
-                                 putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show (peer state)
-                                 msgLoop state pieceStatus
+                                 sendMsg (handle pState) KeepAliveMsg
+                                 putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show (peer pState)
+                                 msgLoop pState pieceStatus
                                BitFieldMsg bss -> do
                                  let pieceList = bitfieldToList (unpack bss)
-                                     pieceStatus' = updatePieceAvailability pieceStatus (peer state) pieceList
+                                     pieceStatus' = updatePieceAvailability pieceStatus (peer pState) pieceList
                                  print pieceList
                                  -- for each pieceIndex in pieceList, make an entry in the pieceStatus
                                  -- map with pieceIndex as the key and modify the value to add the peer.
                                  -- download each of the piece in order
-                                 msgLoop state pieceStatus'
+                                 msgLoop pState pieceStatus'
                                UnChokeMsg -> do
-                                 msgLoop (state { heChoking = False }) pieceStatus
+                                 msgLoop (pState { heChoking = False }) pieceStatus
                                _ -> do
-                                 msgLoop state pieceStatus
+                                 msgLoop pState pieceStatus
 
 -- simple algorithm to pick piece.
 -- pick the first piece from 0 that is not downloaded yet.
 pickPiece :: PieceMap -> Maybe Integer
 pickPiece m =
   let pieceList = toList m
-      allPending = filter (\(k, v) -> state v == Pending) pieceList
+      allPending = filter (\(_, v) -> state v == Pending) pieceList
   in
    case allPending of
     [] -> Nothing
@@ -250,18 +258,38 @@ updatePieceAvailability pieceStatus p pieceList =
                        then (pd { peers = p : (peers pd) })
                        else pd) pieceStatus
 
-handlePeerMsgs :: Peer -> Metainfo -> String -> (String -> IO ()) -> IO ()
-handlePeerMsgs p m peerId logFn = do
+handlePeerMsgs :: Peer -> Metainfo -> String -> IO ()
+handlePeerMsgs p m peerId = do
   h <- handShake p (infoHash m) peerId
-  -- logFn "handShake"
   let state = PeerState { handle = h
                         , peer = p
                         , heInterested = False
                         , heChoking = True
                         , meInterested = False
                         , meChoking = True }
-      pieceHash = (pieces (info m))
+      pieceHash = pieces (info m)
       numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
-      pieceStatus = mkPieceMap numPieces pieceHash
+      pLen = pieceLength (info m)
+      fileLen = lengthInBytes (info m)
+      pieceStatus = mkPieceMap numPieces pieceHash (splitNum fileLen pLen)
   msgLoop state pieceStatus
   
+downloadPiece :: Handle -> Integer -> Integer -> IO ByteString
+downloadPiece h index pieceLength = do
+  let chunks = splitNum pieceLength 16384
+  liftM concat $ forM (zip [0..] chunks) (\(i, pLen) -> do
+                                             sendMsg h (RequestMsg index (i*pLen) pLen)
+                                             putStrLn $ "--> " ++ "RequestMsg for Piece "
+                                               ++ (show index) ++ ", part: " ++ show i ++ " of length: "
+                                               ++ show pLen
+                                             -- TODO: non exhaustive pattern matching. :(
+                                             PieceMsg index begin block <- getMsg h
+                                             putStrLn $ " <-- PieceMsg for Piece: "
+                                               ++ show index
+                                               ++ ", offset: "
+                                               ++ show begin
+                                             return block)
+
+verifyHash :: ByteString -> ByteString -> Bool
+verifyHash bs pieceHash =
+  take 20 (SHA1.hash bs) == pieceHash