verify SHA1 hash of each piece after assembling the blocks

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

diff --git a/src/FuncTorrent/Peer.hs b/src/FuncTorrent/Peer.hs
index dba33bc288e876bf720dc174317d9c8d9af22575..7802b08cb055a23d2f23ab14e8f2373e8a2e989a 100644 (file)
--- a/src/FuncTorrent/Peer.hs
+++ b/src/FuncTorrent/Peer.hs
@@ -4,10 +4,10 @@ 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(..))
@@ -19,6 +19,7 @@ import Control.Applicative ((<$>), liftA3)
 import Data.Bits
 import Data.Word (Word8)
 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, splitNum)
@@ -188,7 +189,13 @@ 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 pState pieceStatus | meInterested pState == False &&
@@ -207,12 +214,12 @@ msgLoop pState pieceStatus | meInterested pState == False &&
                                Nothing -> putStrLn "Nothing to download"
                                Just workPiece -> do
                                  let pLen = len (pieceStatus ! workPiece)
-                                 _ <- downloadPiece (handle pState) workPiece pLen
-                                 -- sendMsg (handle state) (RequestMsg workPiece 0 pLen)
-                                 -- putStrLn $ "--> RequestMsg for Piece " ++ (show workPiece) ++ "to peer: " ++ show (peer state) ++ " of length: " ++ show pLen
-                                 -- msg <- getMsg (handle state)
-                                 -- putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer state)
-                                 msgLoop pState (adjust (\pieceData -> pieceData { state = Have }) workPiece pieceStatus)
+                                 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 pState)
                               putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer pState)
@@ -267,11 +274,22 @@ handlePeerMsgs p m peerId = do
       pieceStatus = mkPieceMap numPieces pieceHash (splitNum fileLen pLen)
   msgLoop state pieceStatus
   
-downloadPiece :: Handle -> Integer -> Integer -> IO [PeerMsg]
+downloadPiece :: Handle -> Integer -> Integer -> IO ByteString
 downloadPiece h index pieceLength = do
   let chunks = splitNum pieceLength 16384
-  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
-                              -- putStrLn $ "--> RequestMsg for Piece " ++ (show workPiece) ++ "to peer: " ++ show (peer state) ++ " of length: " ++ show pLen
-                              getMsg h)
+  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