]> git.rkrishnan.org Git - functorrent.git/blobdiff - src/FuncTorrent/Peer.hs
peer is needed only for debug print
[functorrent.git] / src / FuncTorrent / Peer.hs
index ae768a88fd765d2709268b55cfc8ee3af711cece..88e5a7105e73c8e74719519659cedbcb30e280c1 100644 (file)
@@ -4,10 +4,10 @@ module FuncTorrent.Peer
      handlePeerMsgs
     ) where
 
-import Prelude hiding (lookup, concat, replicate, splitAt, writeFile)
+import Prelude hiding (lookup, concat, replicate, splitAt, take)
 
-import System.IO (Handle, BufferMode(..), IOMode(..), SeekMode(..), withFile, hSeek, hSetBuffering)
-import Data.ByteString (ByteString, pack, unpack, concat, hGet, hPut, singleton, writeFile)
+import System.IO (Handle, BufferMode(..), hSetBuffering)
+import Data.ByteString (ByteString, pack, unpack, concat, hGet, hPut, singleton, take, empty)
 import Data.ByteString.Lazy (fromStrict, fromChunks, toStrict)
 import qualified Data.ByteString.Char8 as BC (replicate, pack, length)
 import Network (connectTo, PortID(..))
@@ -19,9 +19,11 @@ 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)
+import FuncTorrent.Fileops (createDummyFile, writeFileAtOffset)
 
 type ID = String
 type IP = String
@@ -75,7 +77,7 @@ mkPieceMap numPieces pieceHash pLengths = fromList kvs
                              , hash = h
                              , len = pLen })
               | (i, h, pLen) <- zip3 [0..numPieces] hashes pLengths]
-        hashes = splitN (fromIntegral numPieces) pieceHash
+        hashes = splitN 20 pieceHash
 
 havePiece :: PieceMap -> Integer -> Bool
 havePiece pm index =
@@ -88,16 +90,20 @@ genHandShakeMsg infoHash peer_id = concat [pstrlen, pstr, reserved, infoHash, pe
         reserved = BC.replicate 8 '\0'
         peerID = BC.pack peer_id
 
-handShake :: Peer -> ByteString -> String -> IO Handle
-handShake peer@(Peer _ ip port) infoHash peerid = do
-  let hs = genHandShakeMsg infoHash peerid
+connectToPeer :: Peer -> IO Handle
+connectToPeer peer@(Peer _ ip port) = do
   h <- connectTo ip (PortNumber (fromIntegral port))
   hSetBuffering h LineBuffering
+  return h
+
+doHandShake :: Handle -> Peer -> ByteString -> String -> IO ()
+doHandShake h peer infoHash peerid = do
+  let hs = genHandShakeMsg infoHash peerid
   hPut h hs
   putStrLn $ "--> handhake to peer: " ++ show peer
   _ <- hGet h (length (unpack hs))
   putStrLn $ "<-- handshake from peer: " ++ show peer
-  return h
+  return ()
 
 instance Binary PeerMsg where
   put msg = case msg of
@@ -168,10 +174,8 @@ getMsg h = do
     return $ decode $ fromStrict $ concat [lBS, msg]
 
 sendMsg :: Handle -> PeerMsg -> IO ()
-sendMsg h msg =
-  let bsMsg = toStrict $ encode msg
-  in
-   hPut h bsMsg
+sendMsg h msg = hPut h bsMsg
+  where bsMsg = toStrict $ encode msg
 
 bsToInt :: ByteString -> Int
 bsToInt x = fromIntegral (runGet getWord32be (fromChunks (return x)))
@@ -184,38 +188,32 @@ bitfieldToList bs = go bs 0
           in
            setBits ++ go bs' (pos + 1)
 
-createDummyFile :: FilePath -> Int -> IO ()
-createDummyFile path size =
-  writeFile path (BC.replicate size '\0')
-
--- 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 &&
-                            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 pState)
-                              msgLoop (pState { meInterested = True }) pieceStatus
-                          | meInterested pState == True &&
-                            heChoking pState == False =
+msgLoop pState pieceStatus | not (meInterested pState) && heChoking pState = 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 pState)
+                               msgLoop (pState { meInterested = True }) pieceStatus
+                          | meInterested pState && not (heChoking pState) =
                               -- 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
                                  let pLen = len (pieceStatus ! workPiece)
-                                 _ <- downloadPiece (handle pState) workPiece pLen
-                                 -- TODO: verify the hash
-                                 msgLoop pState (adjust (\pieceData -> pieceData { state = Have }) workPiece pieceStatus)
+                                 putStrLn $ "piece length = " ++ show pLen
+                                 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 do
+                                   let fileOffset = if workPiece == 0 then 0 else workPiece * len (pieceStatus ! (workPiece - 1))
+                                   putStrLn $ "Write into file at offset: " ++ show fileOffset
+                                   writeFileAtOffset "/tmp/download.file" fileOffset pBS
+                                   msgLoop pState (adjust (\pieceData -> pieceData { state = Have }) workPiece pieceStatus)
                           | otherwise = do
                               msg <- getMsg (handle pState)
                               putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer pState)
@@ -227,14 +225,14 @@ msgLoop pState pieceStatus | meInterested pState == False &&
                                BitFieldMsg bss -> do
                                  let pieceList = bitfieldToList (unpack bss)
                                      pieceStatus' = updatePieceAvailability pieceStatus (peer pState) pieceList
-                                 print pieceList
+                                 putStrLn $ show (length pieceList) ++ " Pieces"
                                  -- 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 pState pieceStatus'
-                               UnChokeMsg -> do
+                               UnChokeMsg ->
                                  msgLoop (pState { heChoking = False }) pieceStatus
-                               _ -> do
+                               _ ->
                                  msgLoop pState pieceStatus
 
 -- simple algorithm to pick piece.
@@ -251,12 +249,13 @@ pickPiece m =
 updatePieceAvailability :: PieceMap -> Peer -> [Integer] -> PieceMap
 updatePieceAvailability pieceStatus p pieceList =
   mapWithKey (\k pd -> if k `elem` pieceList
-                       then (pd { peers = p : (peers pd) })
+                       then (pd { peers = p : peers pd })
                        else pd) pieceStatus
 
 handlePeerMsgs :: Peer -> Metainfo -> String -> IO ()
 handlePeerMsgs p m peerId = do
-  h <- handShake p (infoHash m) peerId
+  h <- connectToPeer p
+  doHandShake h p (infoHash m) peerId
   let state = PeerState { handle = h
                         , peer = p
                         , heInterested = False
@@ -268,19 +267,29 @@ handlePeerMsgs p m peerId = do
       pLen = pieceLength (info m)
       fileLen = lengthInBytes (info m)
       pieceStatus = mkPieceMap numPieces pieceHash (splitNum fileLen pLen)
+  createDummyFile "/tmp/download.file" (fromIntegral fileLen)
   msgLoop state pieceStatus
   
-downloadPiece :: Handle -> Integer -> Integer -> IO [ByteString]
+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
-                              PieceMsg index begin block <- getMsg h
-                              putStrLn $ " <-- PieceMsg for Piece: "
-                                ++ show index
-                                ++ ", offset: "
-                                ++ show begin
-                              return block)
+  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
+                                             msg <- getMsg h
+                                             case msg of
+                                              PieceMsg index begin block -> do
+                                                putStrLn $ " <-- PieceMsg for Piece: "
+                                                  ++ show index
+                                                  ++ ", offset: "
+                                                  ++ show begin
+                                                return block
+                                              _ -> do
+                                                putStrLn "ignoring irrelevant msg"
+                                                return empty)
+
+verifyHash :: ByteString -> ByteString -> Bool
+verifyHash bs pieceHash =
+  take 20 (SHA1.hash bs) == pieceHash