handshake for server and client

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

diff --git a/src/FuncTorrent/Peer.hs b/src/FuncTorrent/Peer.hs
index e4d7a7222e2849d9c3d9572a3532fa0f25951c96..2caf377ac6849a79d3e38afe4424d50ee77c32d0 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
@@ -83,14 +83,31 @@ 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 isClient h peer infoHash peerid =
   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 ()
+  in
+   if isClient
+   then do
+     hPut h hs
+     putStrLn $ "--> handhake to peer: " ++ show peer
+     _ <- hGet h (length (unpack hs))
+     putStrLn $ "<-- handshake from peer: " ++ show peer
+     return ()
+   else do
+     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
@@ -132,10 +149,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