handle Choke, interested, notinterested, cancel and port msgs

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

diff --git a/src/FuncTorrent/Peer.hs b/src/FuncTorrent/Peer.hs
index 2b9ff5ffd2264d26c311407d697ed4a0f7e6418e..e4a53495d0949ddda016beaa175914abd79c1a41 100644 (file)
--- a/src/FuncTorrent/Peer.hs
+++ b/src/FuncTorrent/Peer.hs
@@ -1,27 +1,27 @@
 {-# LANGUAGE OverloadedStrings #-}
 module FuncTorrent.Peer
     (Peer(..),
+     PieceMap,
      handlePeerMsgs,
-     bytesDownloaded
+     bytesDownloaded,
+     initPieceMap,
+     pieceMapFromFile
     ) where
 
-import Prelude hiding (lookup, concat, replicate, splitAt, take, filter)
+import Prelude hiding (lookup, concat, replicate, splitAt, take, drop, filter)
 
-import System.IO (Handle, BufferMode(..), hSetBuffering)
-import System.Directory (doesFileExist)
-import Data.ByteString (ByteString, unpack, concat, hGet, hPut, take, empty)
+import System.IO (Handle, BufferMode(..), hSetBuffering, hClose)
+import Data.ByteString (ByteString, unpack, concat, hGet, hPut, take, drop, empty)
 import qualified Data.ByteString.Char8 as BC (length)
 import Network (connectTo, PortID(..))
 import Control.Monad.State
 import Data.Bits
 import Data.Word (Word8)
 import Data.Map (Map, fromList, toList, (!), mapWithKey, traverseWithKey, adjust, filter)
-import qualified Crypto.Hash.SHA1 as SHA1 (hash)
 import Safe (headMay)
 
 import FuncTorrent.Metainfo (Info(..), Metainfo(..))
-import FuncTorrent.Utils (splitN, splitNum)
-import FuncTorrent.Fileops (createDummyFile, writeFileAtOffset, readFileAtOffset)
+import FuncTorrent.Utils (splitN, splitNum, writeFileAtOffset, readFileAtOffset, verifyHash)
 import FuncTorrent.PeerMsgs (Peer(..), PeerMsg(..), sendMsg, getMsg, genHandshakeMsg)
 
 data PState = PState { handle :: Handle
@@ -62,14 +62,6 @@ initPieceMap pieceHash fileLen pieceLen = fromList kvs
     hashes = splitN 20 pieceHash
     pLengths = (splitNum fileLen pieceLen)
 
-updatePieceMap :: FilePath -> PieceMap -> IO PieceMap
-updatePieceMap filePath pieceMap = do
-  dfe <- doesFileExist filePath
-  -- TODO: this is not enough, file should have the same size as well
-  if dfe
-    then pieceMapFromFile filePath pieceMap
-    else return pieceMap
-
 pieceMapFromFile :: FilePath -> PieceMap -> IO PieceMap
 pieceMapFromFile filePath pieceMap = do
   traverseWithKey f pieceMap
@@ -91,14 +83,29 @@ connectToPeer (Peer _ ip port) = do
   hSetBuffering h LineBuffering
   return h
 
-doHandshake :: Handle -> Peer -> ByteString -> String -> IO ()
-doHandshake h peer infoHash peerid = do
+doHandshake :: Bool -> Handle -> Peer -> ByteString -> String -> IO ()
+doHandshake True 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 ()
+doHandshake False h peer infoHash peerid = do
+  let hs = genHandshakeMsg infoHash peerid
+  putStrLn $ "waiting for a handshake"
+  hsMsg <- hGet h (length (unpack hs))
+  putStrLn $ "<-- handshake from peer: " ++ show peer
+  let rxInfoHash = take 20 $ drop 28 hsMsg
+  if rxInfoHash /= infoHash
+    then do
+    putStrLn $ "infoHashes does not match"
+    hClose h
+    return ()
+    else do
+    _ <- hPut h hs
+    putStrLn $ "--> handhake to peer: " ++ show peer
+    return ()
 
 bitfieldToList :: [Word8] -> [Integer]
 bitfieldToList bs = go bs 0
@@ -140,19 +147,13 @@ updatePieceAvailability pieceStatus p pieceList =
                        then (pd { peers = p : peers pd })
                        else pd) pieceStatus
 
-handlePeerMsgs :: Peer -> String -> Metainfo -> IO ()
-handlePeerMsgs p peerId m = do
+handlePeerMsgs :: Peer -> String -> Metainfo -> PieceMap -> Bool -> IO ()
+handlePeerMsgs p peerId m pieceMap isClient = do
   h <- connectToPeer p
-  doHandshake h p (infoHash m) peerId
+  doHandshake isClient h p (infoHash m) peerId
   let pstate = toPeerState h p False False True True
-      pieceHash = pieces (info m)
-      pLen = pieceLength (info m)
-      fileLen = lengthInBytes (info m)
-      fileName = name (info m)
-      pieceStatus = initPieceMap pieceHash fileLen pLen
-  pieceStatus' <- updatePieceMap fileName pieceStatus
-  createDummyFile fileName (fromIntegral fileLen)
-  _ <- runStateT (msgLoop pieceStatus' fileName) pstate
+      filePath = name (info m)
+  _ <- runStateT (msgLoop pieceMap filePath) pstate
   return ()
 
 msgLoop :: PieceMap -> FilePath -> StateT PState IO ()
@@ -174,7 +175,7 @@ msgLoop pieceStatus file = do
           pBS <- liftIO $ downloadPiece h workPiece pLen
           if not $ verifyHash pBS (hash (pieceStatus ! workPiece))
             then
-            liftIO $ putStrLn $ "Hash mismatch: " ++ show (hash (pieceStatus ! workPiece)) ++ " vs " ++ show (take 20 (SHA1.hash pBS))
+            liftIO $ putStrLn $ "Hash mismatch"
             else do
             let fileOffset = if workPiece == 0 then 0 else workPiece * len (pieceStatus ! (workPiece - 1))
             liftIO $ putStrLn $ "Write into file at offset: " ++ show fileOffset
@@ -200,6 +201,21 @@ msgLoop pieceStatus file = do
         UnChokeMsg -> do
           modify (\st -> st {heChoking = False })
           msgLoop pieceStatus file
+        ChokeMsg -> do
+          modify (\st -> st {heChoking = True })
+          msgLoop pieceStatus file
+        InterestedMsg -> do
+          modify (\st -> st {heInterested = True})
+          msgLoop pieceStatus file
+        NotInterestedMsg -> do
+          modify (\st -> st {heInterested = False})
+          msgLoop pieceStatus file
+        CancelMsg _ _ _ -> do -- check if valid index, begin, length
+          msgLoop pieceStatus file
+        PortMsg _ -> do
+          msgLoop pieceStatus file
+        -- handle RequestMsg, HaveMsg. No need to handle PieceMsg here.
+        -- also BitFieldMsg
 
 
 downloadPiece :: Handle -> Integer -> Integer -> IO ByteString
@@ -222,6 +238,3 @@ downloadPiece h index pieceLength = do
                                                   putStrLn "ignoring irrelevant msg"
                                                   return empty)
 
-verifyHash :: ByteString -> ByteString -> Bool
-verifyHash bs pieceHash =
-  take 20 (SHA1.hash bs) == pieceHash