Peer: rename identifiers that are shadowing already defined identifiers

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

diff --git a/src/FuncTorrent/Peer.hs b/src/FuncTorrent/Peer.hs
index e4d7a7222e2849d9c3d9572a3532fa0f25951c96..c6629a4c980146a734f3e003d2ef7751506356ad 100644 (file)
--- a/src/FuncTorrent/Peer.hs
+++ b/src/FuncTorrent/Peer.hs
@@ -8,10 +8,10 @@ module FuncTorrent.Peer
      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 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
@@ -31,8 +31,6 @@ data PState = PState { handle :: Handle
                      , heChoking :: Bool
                      , heInterested :: Bool}
 
-type PeerState = State PState
-
 data PieceDlState = Pending
                   | Downloading
                   | Have
@@ -60,18 +58,18 @@ initPieceMap pieceHash fileLen pieceLen = fromList kvs
                          , len = pLen })
           | (i, h, pLen) <- zip3 [0..numPieces] hashes pLengths]
     hashes = splitN 20 pieceHash
-    pLengths = (splitNum fileLen pieceLen)
+    pLengths = splitNum fileLen pieceLen
 
 pieceMapFromFile :: FilePath -> PieceMap -> IO PieceMap
-pieceMapFromFile filePath pieceMap = do
+pieceMapFromFile filePath pieceMap =
   traverseWithKey f pieceMap
-    where
-      f k v = do
-        let offset = if k == 0 then 0 else k * len (pieceMap ! (k - 1))
-        isHashValid <- (flip verifyHash) (hash v) <$> (readFileAtOffset filePath offset (len v))
-        if isHashValid
-          then return $ v { dlstate = Have }
-          else return $ v
+  where
+    f k v = do
+      let offset = if k == 0 then 0 else k * len (pieceMap ! (k - 1))
+      isHashValid <- flip verifyHash (hash v) <$> readFileAtOffset filePath offset (len v)
+      if isHashValid
+        then return $ v { dlstate = Have }
+        else return v
 
 havePiece :: PieceMap -> Integer -> Bool
 havePiece pm index =
@@ -83,14 +81,29 @@ connectToPeer (Peer _ ip port) = do
   hSetBuffering h LineBuffering
   return h
 
-doHandshake :: Handle -> Peer -> ByteString -> String -> IO ()
-doHandshake h peer infoHash peerid = do
-  let hs = genHandshakeMsg infoHash peerid
+doHandshake :: Bool -> Handle -> Peer -> ByteString -> String -> IO ()
+doHandshake True h p infohash peerid = do
+  let hs = genHandshakeMsg infohash peerid
   hPut h hs
-  putStrLn $ "--> handhake to peer: " ++ show peer
+  putStrLn $ "--> handhake to peer: " ++ show p
   _ <- hGet h (length (unpack hs))
-  putStrLn $ "<-- handshake from peer: " ++ show peer
+  putStrLn $ "<-- handshake from peer: " ++ show p
   return ()
+doHandshake False h p infohash peerid = do
+  let hs = genHandshakeMsg infohash peerid
+  putStrLn "waiting for a handshake"
+  hsMsg <- hGet h (length (unpack hs))
+  putStrLn $ "<-- handshake from peer: " ++ show p
+  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 p
+    return ()
 
 bitfieldToList :: [Word8] -> [Integer]
 bitfieldToList bs = go bs 0
@@ -132,10 +145,10 @@ updatePieceAvailability pieceStatus p pieceList =
                        then (pd { peers = p : peers pd })
                        else pd) pieceStatus
 
-handlePeerMsgs :: Peer -> String -> Metainfo -> PieceMap -> IO ()
-handlePeerMsgs p peerId m pieceMap = 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
       filePath = name (info m)
   _ <- runStateT (msgLoop pieceMap filePath) pstate
@@ -149,7 +162,7 @@ msgLoop pieceStatus file = do
     PState { meInterested = False, heChoking = True } -> do
       liftIO $ sendMsg h InterestedMsg
       gets peer >>= (\p -> liftIO $ putStrLn $ "--> InterestedMsg to peer: " ++ show p)
-      modify (\st -> st { meInterested = True })
+      modify (\st' -> st' { meInterested = True })
       msgLoop pieceStatus file
     PState { meInterested = True, heChoking = False } ->
       case pickPiece pieceStatus of
@@ -160,7 +173,7 @@ msgLoop pieceStatus file = do
           pBS <- liftIO $ downloadPiece h workPiece pLen
           if not $ verifyHash pBS (hash (pieceStatus ! workPiece))
             then
-            liftIO $ putStrLn $ "Hash mismatch"
+            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
@@ -184,8 +197,23 @@ msgLoop pieceStatus file = do
           -- download each of the piece in order
           msgLoop pieceStatus' file
         UnChokeMsg -> do
-          modify (\st -> st {heChoking = False })
+          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 _ _ _ -> -- check if valid index, begin, length
+          msgLoop pieceStatus file
+        PortMsg _ ->
           msgLoop pieceStatus file
+        -- handle RequestMsg, HaveMsg. No need to handle PieceMsg here.
+        -- also BitFieldMsg
 
 
 downloadPiece :: Handle -> Integer -> Integer -> IO ByteString