peer: more refactoring. Supports all messages as before

diff --git a/src/FuncTorrent/Peer.hs b/src/FuncTorrent/Peer.hs
index 99906b074293bd260ec748cb5b027fdbf15c2aa9..cf5b66d35f35e6069b26cf1f56ecb1d893984fd8 100644 (file)
--- a/src/FuncTorrent/Peer.hs
+++ b/src/FuncTorrent/Peer.hs
@@ -99,53 +99,6 @@ toPeerState h p meCh meIn heCh heIn =
          , meChoking = meCh
          , meInterested = meIn }
 
--- -- recvMsg :: Peer -> Handle -> Msg
--- msgLoop :: PeerState -> PieceMap -> FilePath -> IO ()
--- msgLoop pState@(PeerState { meInterested = False, heChoking = True }) pieceStatus file = 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 file
--- msgLoop pState@(PeerState { meInterested = True, heChoking = False }) pieceStatus file =
---   -- 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)
---      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 file fileOffset pBS
---        msgLoop pState (adjust (\pieceData -> pieceData { state = Have }) workPiece pieceStatus) file
--- msgLoop pState pieceStatus file = do
---   msg <- getMsg (handle pState)
---   putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer pState)
---   case msg of
---    KeepAliveMsg -> do
---      sendMsg (handle pState) KeepAliveMsg
---      putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show (peer pState)
---      msgLoop pState pieceStatus file
---    BitFieldMsg bss -> do
---      let pieceList = bitfieldToList (unpack bss)
---          pieceStatus' = updatePieceAvailability pieceStatus (peer pState) 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' file
---    UnChokeMsg ->
---      msgLoop (pState { heChoking = False }) pieceStatus file
---    _ ->
---      msgLoop pState pieceStatus file
-
 -- simple algorithm to pick piece.
 -- pick the first piece from 0 that is not downloaded yet.
 pickPiece :: PieceMap -> Maybe Integer
@@ -181,13 +134,49 @@ handlePeerMsgs p m peerId = do
 msgLoop :: PieceMap -> FilePath -> StateT PState IO ()
 msgLoop pieceStatus file = do
   h <- gets handle
-  msg <- liftIO $ getMsg h
-  gets peer >>= (\p -> liftIO $ putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show p)
-  case msg of
-    KeepAliveMsg -> do
-      liftIO $ sendMsg h KeepAliveMsg
-      gets peer >>= (\p -> liftIO $ putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show p)
+  st <- get
+  case st of
+    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 })
       msgLoop pieceStatus file
+    PState { meInterested = True, heChoking = False } -> do
+      case pickPiece pieceStatus of
+        Nothing -> liftIO $ putStrLn "Nothing to download"
+        Just workPiece -> do
+          let pLen = len (pieceStatus ! workPiece)
+          liftIO $ putStrLn $ "piece length = " ++ show pLen
+          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))
+            else do
+            let fileOffset = if workPiece == 0 then 0 else workPiece * len (pieceStatus ! (workPiece - 1))
+            liftIO $ putStrLn $ "Write into file at offset: " ++ show fileOffset
+            liftIO $ writeFileAtOffset file fileOffset pBS
+            msgLoop (adjust (\pieceData -> pieceData { dlstate = Have }) workPiece pieceStatus) file
+    _ -> do
+      msg <- liftIO $ getMsg h
+      gets peer >>= (\p -> liftIO $ putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show p)
+      case msg of
+        KeepAliveMsg -> do
+          liftIO $ sendMsg h KeepAliveMsg
+          gets peer >>= (\p -> liftIO $ putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show p)
+          msgLoop pieceStatus file
+        BitFieldMsg bss -> do
+          p <- gets peer
+          let pieceList = bitfieldToList (unpack bss)
+              pieceStatus' = updatePieceAvailability pieceStatus p pieceList
+          liftIO $ 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 pieceStatus' file
+        UnChokeMsg -> do
+          modify (\st -> st {heChoking = False })
+          msgLoop pieceStatus file
+
 
 downloadPiece :: Handle -> Integer -> Integer -> IO ByteString
 downloadPiece h index pieceLength = do