refactor handshake msg

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

diff --git a/src/FuncTorrent/Peer.hs b/src/FuncTorrent/Peer.hs
index 845c14aa0d51f675a9ee00184c761ebeae0db486..22dc48f485ef6fd51be47a6aabb185510ab0b0f2 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, take)
+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, take, empty)
+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(..))
@@ -23,6 +23,7 @@ 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
@@ -89,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@(Peer _ ip port) 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
@@ -169,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)))
@@ -185,17 +188,6 @@ 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 ReadWriteMode (\h -> do
-                                  _ <- hSeek h AbsoluteSeek offset
-                                  hPut h block)
-
 -- recvMsg :: Peer -> Handle -> Msg
 msgLoop :: PeerState -> PieceMap -> IO ()
 msgLoop pState pieceStatus | not (meInterested pState) && heChoking pState = do
@@ -218,7 +210,8 @@ msgLoop pState pieceStatus | not (meInterested pState) && heChoking pState = do
                                    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 (len (pieceStatus ! (workPiece - 1)))
+                                   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
@@ -261,7 +254,8 @@ updatePieceAvailability pieceStatus p pieceList =
 
 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